"""Batch answering: many questions asked of one document as one server-side job.

Each question is an ordinary answer — retrieval runs for it, the answer step
fills the same `Answer` a single question produces — so a batch introduces no
second answer format. What this module adds is the run around them: a job that
continues if the page is closed, a bounded number of questions in flight, a
retry with backoff for a provider that returns 529 under load, and one event
per question completing so progress is a count filling in rather than a log
being tailed.

The two steps are two stages with their own bounds. Retrieval for one question
does not wait for another question's answer: up to `concurrency` retrievals
run while up to `concurrency` answers run, and at most `concurrency` retrieved
contexts wait between them. The run then moves at the pace of the slower stage
rather than the sum of both. Retrieval alone took about as long as the answer
on the Jev ranker, so serial slots spent half their life on work that held no
answer-model call.

The runner holds no reference to the app. The caller hands `start` two
coroutines, one that retrieves for a question and one that answers it from
what was retrieved; everything here is scheduling, accounting and delivery,
which is what makes it testable without a model.

Runs live in the server's memory, capped at the twenty most recent so a
long-lived process does not accumulate them. Given a store, the runner also
records each run and every finished row in the database, and answers for a
run it does not hold in memory from there: a run started by another task,
or by a process since replaced. A run known to neither reaches the page as
an id it kept and a run nobody can answer for. That is never silence — the
route 404s and the page reports the run as gone, because a batch of forty
figures quietly becoming thirty-nine is the failure this feature exists to
prevent.
"""

from __future__ import annotations

import asyncio
import uuid
from collections import OrderedDict
from dataclasses import dataclass, field
from typing import Any, AsyncIterator, Awaitable, Callable, Dict, List, Optional, Protocol

from pydantic import BaseModel

from quber.settings import get_settings

#: Runs kept in memory, most recent first. Chosen to put eviction out of reach
#: of a working session rather than to be relied on.
MAX_RUNS = 20

#: HTTP statuses worth retrying: overload and transient server failures.
RETRY_STATUS = {429, 500, 502, 503, 529}
MAX_ATTEMPTS = 4
FIRST_RETRY_DELAY = 2.0

#: The two coroutines the caller supplies. Retrieval: one question in, an
#: opaque context out, whatever the answer stage needs. Answer: the question
#: and that context in, the answer contract out.
RetrieveOne = Callable[[str], Awaitable[Any]]
AnswerOne = Callable[[str, Any], Awaitable[Dict[str, Any]]]

#: How often a run served from the store is re-read for new rows.
STORE_POLL_SECONDS = 2.0


class RunStore(Protocol):
    """The durable side of a run: what the runner writes as it goes and what
    it reads back for a run it does not hold in memory."""

    def save_run(self, job_id: str, doc_id: int, want: str, questions: List[str]) -> None: ...
    def save_row(
        self, job_id: str, index: int, question: str, error: Optional[str], response: Optional[Dict[str, Any]]
    ) -> None: ...
    def finish_run(self, job_id: str) -> None: ...
    def load_run(self, job_id: str) -> Optional[Dict[str, Any]]: ...
    def running_for(self, doc_id: int) -> bool: ...


def parse_questions(text: str) -> List[str]:
    """One question per line, as the batch screen's prompt window states.

    Lines are trimmed and blank lines dropped. A repeated line is kept — a
    person who pasted a question twice gets two rows, because a question that
    silently vanishes from a run is exactly what a batch must never do.
    """
    return [line.strip() for line in text.splitlines() if line.strip()]


class BatchRow(BaseModel):
    """One finished question. Either `response` (the /api/chat contract,
    verbatim) or `error` is set, never both."""

    index: int
    question: str
    error: Optional[str] = None
    response: Optional[Dict[str, Any]] = None


class RunState(BaseModel):
    """What `GET /api/batch/{id}` returns: the run's counts and the rows
    finished so far, in question order."""

    job_id: str
    doc_id: int
    want: str
    total: int
    done: int
    failed: int
    running: bool
    rows: List[BatchRow]


@dataclass
class _Run:
    job_id: str
    doc_id: int
    want: str
    questions: List[str]
    rows: Dict[int, BatchRow] = field(default_factory=dict)
    failed: int = 0
    task: Optional[asyncio.Task[None]] = None
    subscribers: List[asyncio.Queue[Optional[Dict[str, Any]]]] = field(default_factory=list)

    @property
    def running(self) -> bool:
        return self.task is not None and not self.task.done()

    def state(self) -> RunState:
        return RunState(
            job_id=self.job_id,
            doc_id=self.doc_id,
            want=self.want,
            total=len(self.questions),
            done=len(self.rows),
            failed=self.failed,
            running=self.running,
            rows=[self.rows[i] for i in sorted(self.rows)],
        )

    def publish(self, event: Dict[str, Any]) -> None:
        for q in self.subscribers:
            q.put_nowait(event)

    def close_subscribers(self) -> None:
        for q in self.subscribers:
            q.put_nowait(None)


class BatchRunner:
    """The run registry. One per server process."""

    def __init__(
        self,
        concurrency: Optional[int] = None,
        max_runs: int = MAX_RUNS,
        store: Optional[RunStore] = None,
    ) -> None:
        # Questions in flight at once, from settings unless the caller says.
        # Each holds an answer-model call, and the provider returns 529 under
        # load (observed 2026-07-29 during a sweep), so the bound is deliberate.
        self._concurrency = (
            concurrency if concurrency is not None else get_settings().playground.question_concurrency
        )
        self._max_runs = max_runs
        self._runs: OrderedDict[str, _Run] = OrderedDict()
        self._store = store

    def running_ids(self) -> List[str]:
        """The runs this process is still working on."""
        return [run.job_id for run in self._runs.values() if run.running]

    def start(
        self, doc_id: int, questions: List[str], want: str, retrieve_one: RetrieveOne, answer_one: AnswerOne
    ) -> RunState:
        """Register a run and schedule it on the running event loop."""
        job_id = uuid.uuid4().hex[:12]
        run = _Run(job_id=job_id, doc_id=doc_id, want=want, questions=questions)
        if self._store is not None:
            self._store.save_run(job_id, doc_id, want, questions)
        self._runs[job_id] = run
        while len(self._runs) > self._max_runs:
            _, evicted = self._runs.popitem(last=False)
            if evicted.task is not None and not evicted.task.done():
                evicted.task.cancel()
            evicted.close_subscribers()
        run.task = asyncio.get_running_loop().create_task(self._run(run, retrieve_one, answer_one))
        return run.state()

    def get(self, job_id: str) -> Optional[RunState]:
        run = self._runs.get(job_id)
        if run is not None:
            return run.state()
        stored = self._stored(job_id)
        return self._state_from_stored(stored) if stored else None

    def _stored(self, job_id: str) -> Optional[Dict[str, Any]]:
        return self._store.load_run(job_id) if self._store is not None else None

    @staticmethod
    def _state_from_stored(stored: Dict[str, Any]) -> RunState:
        rows = [BatchRow(**r) for r in stored["rows"]]
        return RunState(
            job_id=stored["job_id"],
            doc_id=stored["doc_id"],
            want=stored["want"],
            total=len(stored["questions"]),
            done=len(rows),
            failed=stored["failed"],
            running=stored["running"],
            rows=rows,
        )

    def running_for(self, doc_id: int) -> bool:
        """Whether any run against this document is still going — what lets
        a document removal or replacement refuse instead of pulling the
        corpus out from under a run mid-flight."""
        if any(run.doc_id == doc_id and run.running for run in self._runs.values()):
            return True
        return self._store is not None and self._store.running_for(doc_id)

    async def events(self, job_id: str) -> AsyncIterator[Dict[str, Any]]:
        """Yield one event per question completing, then a final `done` event.

        Subscribing to a finished run yields the `done` event immediately, so
        a page that attaches late still learns the run is over.
        """
        run = self._runs.get(job_id)
        if run is None:
            async for event in self._stored_events(job_id):
                yield event
            return
        if not run.running:
            yield self._done_event(run)
            return
        queue: asyncio.Queue[Optional[Dict[str, Any]]] = asyncio.Queue()
        run.subscribers.append(queue)
        try:
            while True:
                event = await queue.get()
                if event is None:
                    break
                yield event
                if event.get("type") == "done":
                    break
        finally:
            run.subscribers.remove(queue)

    async def _stored_events(self, job_id: str) -> AsyncIterator[Dict[str, Any]]:
        """Events for a run another process is working on, read from the
        store: each row as it lands there, then done once the run closes."""
        seen: set[int] = set()
        while True:
            stored = self._stored(job_id)
            if stored is None:
                return
            total = len(stored["questions"])
            for row in stored["rows"]:
                if row["index"] in seen:
                    continue
                seen.add(row["index"])
                yield {
                    "type": "row",
                    "index": row["index"],
                    "total": total,
                    "done": len(seen),
                    "failed": stored["failed"],
                    "ok": row["error"] is None,
                    "row": BatchRow(**row).model_dump(mode="json"),
                }
            if not stored["running"]:
                yield {
                    "type": "done",
                    "total": total,
                    "done": len(stored["rows"]),
                    "failed": stored["failed"],
                }
                return
            await asyncio.sleep(STORE_POLL_SECONDS)

    def _done_event(self, run: _Run) -> Dict[str, Any]:
        return {
            "type": "done",
            "total": len(run.questions),
            "done": len(run.rows),
            "failed": run.failed,
        }

    async def _run(self, run: _Run, retrieve_one: RetrieveOne, answer_one: AnswerOne) -> None:
        n = self._concurrency
        # At most 2n questions are in progress at once: n retrieving or waiting
        # with a context in hand, n answering. Retrieval runs ahead exactly as
        # far as the answer stage has fallen behind, and no further.
        in_flight = asyncio.Semaphore(2 * n)
        retrieving = asyncio.Semaphore(n)
        answering = asyncio.Semaphore(n)

        async def one(index: int, question: str) -> None:
            async with in_flight:
                row = BatchRow(index=index, question=question)
                try:
                    async with retrieving:
                        context = await self._with_retry(retrieve_one, question)
                    async with answering:
                        row.response = await self._with_retry(answer_one, question, context)
                except Exception as exc:
                    # The row carries its error rather than vanishing; the run
                    # never reports success after losing a question.
                    row.error = str(exc)
                    run.failed += 1
                run.rows[index] = row
                if self._store is not None:
                    self._store.save_row(run.job_id, index, question, row.error, row.response)
                # The event carries the finished row itself: the page appends
                # it in place rather than refetching the whole run per event.
                run.publish(
                    {
                        "type": "row",
                        "index": index,
                        "total": len(run.questions),
                        "done": len(run.rows),
                        "failed": run.failed,
                        "ok": row.error is None,
                        "row": row.model_dump(mode="json"),
                    }
                )

        try:
            await asyncio.gather(*(one(i, q) for i, q in enumerate(run.questions)))
        finally:
            if self._store is not None:
                self._store.finish_run(run.job_id)
            run.publish(self._done_event(run))
            run.close_subscribers()

    async def _with_retry(self, stage: Callable[..., Awaitable[Any]], *args: Any) -> Any:
        delay = FIRST_RETRY_DELAY
        for attempt in range(MAX_ATTEMPTS):
            try:
                return await stage(*args)
            except Exception as exc:
                status = getattr(exc, "status_code", None)
                if status not in RETRY_STATUS or attempt == MAX_ATTEMPTS - 1:
                    raise
                await asyncio.sleep(delay)
                delay *= 2
        raise RuntimeError("unreachable")  # pragma: no cover