"""Pipeline building blocks: agent races (with grace), single-slot, schemas, prompts, guide status.

Holds the mutable pipeline state (generation semaphore, cancel set).
Access the cancel set ONLY through the functions here — copied references
in other modules would diverge on a re-assign.
"""

import asyncio
import logging
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Callable

from agents import run_agent, kill_process, cancel_scope, clear_scope
from config import MAX_CONCURRENT_GENERATIONS, TEMPLATES_DIR, TIMEOUTS
from database import update_guide

log = logging.getLogger("creator.pipeline")

_semaphore = asyncio.Semaphore(MAX_CONCURRENT_GENERATIONS)
_cancelled: set[str] = set()


async def cancel_guide(guide_id: str) -> bool:
    _cancelled.add(guide_id)
    cancel_scope(f"{guide_id}-")  # waiting agents bail before spawn
    kill_process(guide_id)        # kill running subprocesses
    now = datetime.now(timezone.utc).isoformat()
    await update_guide(guide_id, status="error", progress=None, error_msg="Cancelled — progress is preserved", updated_at=now)
    return True


def is_guide_cancelled(guide_id: str) -> bool:
    return guide_id in _cancelled


def clear_guide_cancelled(guide_id: str) -> None:
    _cancelled.discard(guide_id)
    clear_scope(f"{guide_id}-")  # clear scope → restart not blocked


async def _set_progress(guide_id: str, progress: str) -> None:
    now = datetime.now(timezone.utc).isoformat()
    await update_guide(guide_id, progress=progress, updated_at=now)


async def _set_step(guide_id: str, step: int, progress: str) -> None:
    now = datetime.now(timezone.utc).isoformat()
    await update_guide(guide_id, step=step, progress=progress, updated_at=now)


async def _fail(guide_id: str, msg: str) -> None:
    now = datetime.now(timezone.utc).isoformat()
    await update_guide(guide_id, status="error", progress=None, error_msg=msg, updated_at=now)


def _prompt(name: str, **kwargs) -> str:
    template = (TEMPLATES_DIR / "Prompt" / f"{name}.md").read_text(encoding="utf-8")
    return template.format(**kwargs)


def _extra(instructions: str) -> str:
    return f"\n\nADDITIONAL INSTRUCTIONS FROM THE USER:\n{instructions}\n" if instructions else ""


def _log(topic: str, msg: str) -> None:
    log.info("[%s] %s", topic, msg)


def _claude_error(label: str, returncode: int, stdout: str, stderr: str) -> str:
    stderr = (stderr or "").strip()
    if stderr:
        return f"{label}: {stderr[:1000]}"
    tail = (stdout or "").strip()[-500:]
    if tail:
        return f"{label} (exit {returncode}, stderr empty): …{tail}"
    return f"{label} (exit {returncode}, no output)"


_INFRA_MARKERS = ("HTTP 429", "HTTP 5", "rate_limit", "Timeout after",
                  "ConnectError", "ConnectTimeout", "ReadError", "RemoteProtocolError")


def _is_infra(err: str) -> bool:
    """Transport-/Infra-Fehler (retry + pause) statt inhaltlichem Fehlschlag."""
    return any(m in (err or "") for m in _INFRA_MARKERS)


def _gather_error(label: str, results: list) -> str:
    for r in results:
        if isinstance(r, BaseException):
            return f"{label}: {type(r).__name__}: {r}"
        returncode, stdout, stderr = r
        if returncode != 0:
            return _claude_error(label, returncode, stdout, stderr)
    return f"{label}: no usable result"


def _timeout(step: str, n: int = 0) -> int:
    base, per = TIMEOUTS[step]
    return base + per * n


def _problems_schema(data):
    """{"ok": true} → [] · {"problems": [str]} → list · else None."""
    if not isinstance(data, dict):
        return None
    if data.get("ok") is True:
        return []
    p = data.get("problems")
    if not isinstance(p, list) or not p:
        return None
    out = [str(x).strip() for x in p if str(x).strip()]
    return out or None


def _str_list(val) -> list[str] | None:
    """List of non-empty strings → stripped list (empty allowed) · else None."""
    if not isinstance(val, list) or not all(isinstance(x, str) for x in val):
        return None
    out = [x.strip() for x in val]
    return None if any(not x for x in out) else out




def _runde_schema(data, final: bool = False):
    """{"keep": [str], "rest": [str]} → (include, rest) · else None.

    final=True: last clarification round — a non-empty rest is invalid.
    """
    if not isinstance(data, dict):
        return None
    include = _str_list(data.get("keep"))
    rest = _str_list(data.get("rest"))
    if include is None or rest is None or (final and rest):
        return None
    return include, rest


_YESNO = ("ja", "nein")


def _enum_map_schema(key: str, allowed):
    """Factory for `{"<key>": {"1": value, …}}` → `{id: value}` parsers; value ∈ `allowed`
    (casefolded). If `ids` are given, at least these must be covered (extras allowed). None
    on any invalid id/value or wrong shape. The caller filters the result to `ids`."""
    def parse(data, ids: set[int] | None = None):
        if not isinstance(data, dict) or not isinstance(data.get(key), dict) or not data[key]:
            return None
        out: dict[int, str] = {}
        for k, v in data[key].items():
            try:
                num = int(k)
            except (ValueError, TypeError):
                return None
            value = str(v).strip().casefold()
            if value not in allowed:
                return None
            out[num] = value
        if ids is not None and not ids <= set(out):
            return None
        return out
    return parse


_yesno_schema = _enum_map_schema("relevant", _YESNO)            # triage gate ∈ ja/nein


from config import (MAX_RESTARTS as _MAX_RESTARTS, HEDGE_NACH_S as _HEDGE_NACH_S,  # noqa: E402 — zentral tunebar
                    INFRA_MAX_RETRIES as _INFRA_MAX_RETRIES, INFRA_BACKOFF_BASE as _INFRA_BACKOFF)


async def _race(topic: str, label: str, slots: list[dict], quorum: int, timeout: int, provider: str, cancelled=None, *, grace: int | None = None) -> list | None:
    """Starts all slots in parallel and collects `quorum` valid results.

    Slot spec: {key, prompt, role, capabilities, payload}. `payload(result)`
    checks validity and returns the slot result or None.
    Error/timeout/invalid → slot restart (max. _MAX_RESTARTS). As soon as the
    quorum stands, the remaining agents are killed. None = quorum missed.
    `cancelled()` → True aborts (no restarts, returns None).

    With `grace`, `quorum` becomes the minimum: the first valid result starts
    a timer of `grace` seconds. After it expires, running agents are only
    killed if the minimum stands — otherwise the race, including restarts,
    keeps running until it stands. Returns: `quorum` to `len(slots)` results.
    """
    attempts = {i: 0 for i in range(len(slots))}
    infra_attempts = {i: 0 for i in range(len(slots))}
    infra_erschoepft: set[int] = set()
    tasks: dict[asyncio.Task, int] = {}
    keys: dict[asyncio.Task, str] = {}
    born: dict[asyncio.Task, float] = {}
    hedged: set[int] = set()   # slot got its one twin — no hedge cascades
    fertig: set[int] = set()   # slot delivered a valid result (late twins are ignored)
    # Hedge-Schwelle relativ zum Call-Timeout (HEDGE_NACH_S = Untergrenze): pauschale 90 s
    # hedgten jeden gesunden langen Call — z. B. Guide-Fixes, die normal 110–135 s laufen.
    hedge_s = max(_HEDGE_NACH_S, timeout / 2) if _HEDGE_NACH_S else 0
    loop = asyncio.get_running_loop()
    deadline: float | None = None

    def spawn(i: int, suffix: str = "") -> None:
        slot = slots[i]
        lbl = slot.get("label") or (label if len(slots) == 1 else f"{label} {i + 1}")
        key = slot["key"] + suffix
        task = asyncio.create_task(run_agent(
            key, slot["prompt"], timeout,
            provider=provider, role=slot["role"], capabilities=slot["capabilities"],
            scope=topic, on_line=slot.get("on_line"), label=lbl,
        ))
        tasks[task] = i
        keys[task] = key
        born[task] = loop.time()

    for i in range(len(slots)):
        spawn(i)

    results: list = []
    try:
        while tasks:
            if cancelled and cancelled():
                return None
            if deadline is not None and len(results) >= quorum and loop.time() >= deadline:
                return results
            # Hedge: a slot running HEDGE_NACH_S without result gets ONE parallel twin
            # (key -h) — first valid result wins. Stalled provider calls burned the full
            # timeout cap before the restart even began (measured: 160–230 s per stall).
            if hedge_s:
                now = loop.time()
                for t in [t for t in list(tasks) if tasks[t] not in hedged | fertig
                          and now - born[t] >= hedge_s]:
                    i = tasks[t]
                    hedged.add(i)
                    spawn(i, suffix="-h")
                    _log(topic, f"{label} {i + 1}: {round(hedge_s)}s ohne Ergebnis — Hedge-Zwilling gestartet")
            # Wake up for the earliest relevant deadline (grace or next hedge).
            waits = []
            if deadline is not None and len(results) >= quorum:
                waits.append(deadline - loop.time())
            if hedge_s:
                naechste = [born[t] + hedge_s - loop.time() for t in tasks
                            if tasks[t] not in hedged | fertig]
                if naechste:
                    waits.append(max(0.0, min(naechste)))
            wait_timeout = max(0.0, min(waits)) if waits else None
            done, _ = await asyncio.wait(tasks.keys(), return_when=asyncio.FIRST_COMPLETED, timeout=wait_timeout)
            if not done:
                continue
            for task in done:
                i = tasks.pop(task)
                keys.pop(task, None)
                born.pop(task, None)
                if i in fertig:
                    continue  # späte Zwillinge eines bereits gewerteten Slots
                payload, err = None, None
                try:
                    result = task.result()
                    if result[0] != 0:
                        err = _claude_error("Error", *result)
                    else:
                        payload = slots[i]["payload"](result)
                        if payload is None:
                            err = "result invalid/not parseable"
                except asyncio.TimeoutError:
                    err = f"Timeout after {timeout}s"
                except Exception as e:
                    err = f"{type(e).__name__}: {e}"

                if payload is not None:
                    results.append(payload)
                    fertig.add(i)
                    for t2 in [t2 for t2, i2 in tasks.items() if i2 == i]:  # Zwilling killen
                        kill_process(keys.get(t2, slots[i]["key"]))
                        t2.cancel()
                    if grace is not None and deadline is None:
                        deadline = loop.time() + grace
                        _log(topic, f"{label}: first result — grace {grace}s running")
                    if len(results) >= quorum and (grace is None or loop.time() >= deadline):
                        return results
                    continue

                infra = _is_infra(err)
                _log(topic, f"{label} {i + 1} (attempt {attempts[i] + 1}{' infra' if infra else ''}): {err}")
                attempts[i] += 1
                # If the minimum already stands, restarts are pointless — the restart
                # would be killed at the grace end anyway. A still-running twin IS the retry.
                enough = grace is not None and len(results) >= quorum
                zwilling = any(i2 == i for i2 in tasks.values())
                if enough or zwilling or (cancelled and cancelled()):
                    continue
                if infra:
                    # 429/Timeout/Netz: eigener Zähler + wachsende Pause, dann Eskalation
                    infra_attempts[i] += 1
                    if infra_attempts[i] > _INFRA_MAX_RETRIES:
                        infra_erschoepft.add(i)  # kein Respawn → raise am Quorum-Miss
                        continue
                    await asyncio.sleep(_INFRA_BACKOFF * 2 ** (infra_attempts[i] - 1))
                    spawn(i)
                elif attempts[i] <= _MAX_RESTARTS:
                    spawn(i)
        if len(results) >= quorum:  # all slots done, minimum stands (only reachable with grace)
            return results
        _log(topic, f"{label}: quorum {quorum} not reached ({len(results)} valid)")
        if infra_erschoepft:
            raise AgentInfraError(
                f"{label}: {len(infra_erschoepft)} Slot(s) nach {_INFRA_MAX_RETRIES} "
                "Infra-Retries erschöpft (429/Timeout)")
        return None
    finally:
        for task, i in tasks.items():
            kill_process(keys.get(task, slots[i]["key"]))
            task.cancel()
        if tasks:
            await asyncio.gather(*tasks.keys(), return_exceptions=True)


@dataclass
class GenContext:
    """Pipeline parameters passed through — saves long argument signatures."""
    topic: str
    provider: str
    is_cancelled: Callable[[], bool]
    guide_id: str | None = None


# Result status of run_single_slot
OK, CANCELLED, FAILED = "ok", "cancelled", "failed"


class AgentInfraError(Exception):
    """Slot nach Infra-Retries erschöpft (429/Timeout). KEIN inhaltliches Urteil —
    der Aufrufer pausiert den Lauf, nie fail-open."""


async def run_single_slot(
    ctx: GenContext, label: str, *,
    key: str, prompt: str, role: str, capabilities: str, payload, timeout: int, on_line=None,
) -> tuple[str, object]:
    """One agent, one valid result (race with quorum 1).

    → (OK, value) | (CANCELLED, None) | (FAILED, None)
    """
    slots = [{"key": key, "prompt": prompt, "role": role, "capabilities": capabilities, "payload": payload, "on_line": on_line}]
    res = await _race(ctx.topic, label, slots, 1, timeout, ctx.provider, cancelled=ctx.is_cancelled)
    if ctx.is_cancelled():
        return CANCELLED, None
    if res is None:
        return FAILED, None
    return OK, res[0]


async def _gather_progress(coros, total, report, start=0):
    """Runs `coros` concurrently and reports live progress: `await report(done, total)`
    after each completion (and once initially). Results in order, return_exceptions=True."""
    done = start

    async def wrap(c):
        nonlocal done
        try:
            return await c
        finally:
            done += 1
            await report(done, total)

    await report(done, total)
    return await asyncio.gather(*[wrap(c) for c in coros], return_exceptions=True)