"""Generic worker loop shared by every model adapter.

The loop is model-agnostic: it reads the requests file, resumes from the results file,
calls ``Adapter.run_batch`` in fixed-size batches, and appends one fsynced result line per
request. A failing batch is retried item by item so a single bad clip cannot take its
neighbours down; every failure is recorded with its exception type and message (never
dropped) and counted in the summary. Model loading failures are fatal: they propagate and
the process exits non-zero.
"""

from __future__ import annotations

import abc
import argparse
import hashlib
import importlib
import importlib.metadata
import json
import logging
import sys
import time
import traceback
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, replace
from datetime import UTC, datetime
from pathlib import Path
from typing import ClassVar

from nida.bench.protocol import (
    ProtocolError,
    WorkerSummary,
    WorkRequest,
    WorkResult,
    append_result,
    drop_results,
    load_requests,
    load_results,
    pending_requests,
    read_summary,
    write_summary,
)

__all__ = [
    "MAX_CONSECUTIVE_FAILURES",
    "Adapter",
    "ModelRef",
    "TooManyFailuresError",
    "WorkerOptions",
    "main_for",
    "run_worker",
]

log = logging.getLogger(__name__)

#: Abort when this many items in a row fail: the model/device is broken, not the clips.
MAX_CONSECUTIVE_FAILURES = 10
#: Progress is logged when this many more items are done (and at the end).
PROGRESS_EVERY = 200
#: When the *same* failure run (same clip ids) aborts more than this many times in a row,
#: it is evidently repeatable rather than transient and is kept as recorded failures so a
#: rerun can move past it. Unrelated aborts never accumulate (see ``_abort_count``).
MAX_ABORTS_BEFORE_KEEPING = 2


class TooManyFailuresError(RuntimeError):
    """Consecutive item failures suggest a broken model or device (e.g. CUDA context)."""


@dataclass(frozen=True)
class ModelRef:
    """Which model to load.

    Attributes:
        model_id: Hub id, model card name or local path, interpreted by the adapter.
        revision: Pinned revision (commit SHA) or ``""`` when the source has none.
    """

    model_id: str
    revision: str


class Adapter(abc.ABC):
    """One model family. Subclasses keep heavy imports inside their methods."""

    name: ClassVar[str]
    #: Distribution names whose installed versions are recorded in the summary.
    distributions: ClassVar[tuple[str, ...]] = ()

    @abc.abstractmethod
    def load(self, model: ModelRef, device: str, batch_size: int) -> None:
        """Load the model onto ``device``. Exceptions here abort the worker."""

    @abc.abstractmethod
    def run_batch(self, batch: Sequence[WorkRequest]) -> list[Mapping[str, object]]:
        """Return one JSON-compatible output mapping per request, in the same order."""

    def peak_vram_gb(self) -> float | None:
        """Peak CUDA memory allocated by this process, if torch with CUDA is present."""
        try:
            torch = importlib.import_module("torch")
        except ImportError:
            return None
        if not torch.cuda.is_available():
            return None
        return round(float(torch.cuda.max_memory_allocated()) / 1024**3, 2)


@dataclass(frozen=True)
class WorkerOptions:
    requests: Path
    results: Path
    summary: Path
    model: ModelRef
    device: str
    batch_size: int
    retry_failed: bool


def _now() -> str:
    return datetime.now(UTC).isoformat(timespec="seconds")


def _describe(err: BaseException) -> str:
    return f"{type(err).__name__}: {err}"


def _package_versions(names: Sequence[str]) -> dict[str, str]:
    versions: dict[str, str] = {}
    for name in names:
        try:
            versions[name] = importlib.metadata.version(name)
        except importlib.metadata.PackageNotFoundError:
            versions[name] = "not-installed"
    return versions


def _abort_counter_path(results: Path) -> Path:
    return results.with_name(results.name + ".aborts.json")


def _abort_count(results: Path, failure_ids: Sequence[str]) -> int:
    """Record an abort and return how often *this* failure run has aborted in a row.

    The counter stores a signature of the failing clip ids: a different failure run
    restarts at 1, and :func:`_clear_abort_count` removes it after any run that finishes
    without aborting, so transient incidents in different sessions never add up.
    """
    signature = hashlib.sha1("\n".join(sorted(failure_ids)).encode("utf-8")).hexdigest()
    path = _abort_counter_path(results)
    count = 1
    if path.exists():
        try:
            previous = json.loads(path.read_text(encoding="utf-8"))
            if previous.get("signature") == signature:
                count = int(previous.get("count", 0)) + 1
        except (ValueError, AttributeError, TypeError) as err:
            # A damaged counter must not wedge the worker; recount from this abort.
            log.warning("ignoring unreadable abort counter %s (%s); restarting at 1", path, err)
    tmp = path.with_name(f".{path.name}.tmp")
    tmp.write_text(json.dumps({"signature": signature, "count": count}), encoding="utf-8")
    tmp.replace(path)
    return count


def _clear_abort_count(results: Path) -> None:
    _abort_counter_path(results).unlink(missing_ok=True)


def _drop_failed(path: Path) -> None:
    """Rewrite the results file without failed entries (for ``--retry-failed``)."""
    results = load_results(path, repair=True)
    kept = [r.to_json() + "\n" for r in results.values() if r.ok]
    tmp = path.with_name(f".{path.name}.tmp")
    tmp.write_text("".join(kept), encoding="utf-8")
    tmp.replace(path)


def _validated_outputs(batch: Sequence[WorkRequest], outputs: object) -> list[Mapping[str, object]]:
    if not isinstance(outputs, list) or len(outputs) != len(batch):
        got = len(outputs) if isinstance(outputs, list) else type(outputs).__name__
        raise ProtocolError(f"adapter returned {got} outputs for a batch of {len(batch)}")
    for output in outputs:
        if not isinstance(output, Mapping):
            raise ProtocolError(f"adapter output must be a mapping, got {type(output).__name__}")
    return outputs


def _run_one_batch(adapter: Adapter, batch: Sequence[WorkRequest]) -> list[WorkResult]:
    started = time.perf_counter()
    try:
        outputs = _validated_outputs(batch, adapter.run_batch(batch))
    except Exception as err:  # noqa: BLE001 - recorded per item below, never swallowed
        if len(batch) == 1:
            log.warning("item %s failed: %s", batch[0].clip_id, _describe(err))
            return [
                WorkResult(
                    clip_id=batch[0].clip_id,
                    output=None,
                    error=_describe(err),
                    infer_seconds=time.perf_counter() - started,
                )
            ]
        log.warning("batch of %d failed (%s); retrying items one by one", len(batch), err)
        results: list[WorkResult] = []
        for request in batch:
            results.extend(_run_one_batch(adapter, [request]))
        return results
    per_item = (time.perf_counter() - started) / len(batch)
    return [
        WorkResult(clip_id=r.clip_id, output=dict(o), error=None, infer_seconds=per_item)
        for r, o in zip(batch, outputs, strict=True)
    ]


def run_worker(adapter: Adapter, options: WorkerOptions) -> WorkerSummary:
    """Process every pending request and write the summary. Returns the summary."""
    if options.batch_size < 1:
        raise ValueError(f"batch_size must be >= 1, got {options.batch_size}")
    started_at = _now()
    requests = load_requests(options.requests)
    if options.retry_failed and options.results.exists():
        _drop_failed(options.results)
    done = load_results(options.results, repair=True)
    pending = pending_requests(requests, done)
    log.info(
        "%s: %d requests, %d already done, %d pending",
        adapter.name,
        len(requests),
        len(done),
        len(pending),
    )

    previous = read_summary(options.summary) if options.summary.exists() else None
    load_seconds = 0.0
    if pending:
        load_start = time.perf_counter()
        adapter.load(options.model, options.device, options.batch_size)
        load_seconds = time.perf_counter() - load_start
        options.results.parent.mkdir(parents=True, exist_ok=True)
        failure_run: list[WorkResult] = []
        aborted = False
        with options.results.open("a", encoding="utf-8") as handle:
            for start in range(0, len(pending), options.batch_size):
                batch = pending[start : start + options.batch_size]
                for result in _run_one_batch(adapter, batch):
                    append_result(handle, result)
                    failure_run = [] if result.ok else [*failure_run, result]
                    if len(failure_run) >= MAX_CONSECUTIVE_FAILURES:
                        aborted = True
                        break
                if aborted:
                    break
                finished = start + len(batch)
                if (
                    finished == len(pending)
                    or finished // PROGRESS_EVERY != start // PROGRESS_EVERY
                ):
                    log.info("%s: %d/%d", adapter.name, finished, len(pending))
        if aborted:
            aborts = _abort_count(options.results, [r.clip_id for r in failure_run])
            if aborts <= MAX_ABORTS_BEFORE_KEEPING:
                # Most likely infrastructure (e.g. a broken CUDA context), not bad clips: make
                # them pending again so a plain rerun retries them instead of scoring them as
                # empty hypotheses.
                drop_results(options.results, [r.clip_id for r in failure_run])
                outcome = "results removed, they will be retried"
            else:
                outcome = (
                    f"same failure run aborted {aborts} times: kept as recorded failures so "
                    "the next run moves on; use --retry-failed to try them again"
                )
            raise TooManyFailuresError(
                f"{adapter.name}: {len(failure_run)} consecutive failures ({outcome}), "
                f"last: {failure_run[-1].error}"
            )

    _clear_abort_count(options.results)  # finished without aborting
    final = load_results(options.results)
    n_ok = sum(1 for r in final.values() if r.ok)
    if not pending and previous is not None:
        # No-op resume: keep the measurements of the invocation that did the work.
        summary = replace(previous, n_ok=n_ok, n_error=len(final) - n_ok)
    else:
        peaks = [v for v in (adapter.peak_vram_gb(), previous and previous.peak_vram_gb) if v]
        invocations = int(previous.extra.get("invocations", "1")) + 1 if previous else 1
        summary = WorkerSummary(
            worker=adapter.name,
            model_id=options.model.model_id,
            revision=options.model.revision,
            load_seconds=round(load_seconds, 3),
            peak_vram_gb=max(peaks) if peaks else None,
            n_ok=n_ok,
            n_error=len(final) - n_ok,
            started_at=previous.started_at if previous else started_at,
            ended_at=_now(),
            packages=_package_versions(adapter.distributions),
            extra={
                "python": sys.version.split()[0],
                "device": options.device,
                "invocations": str(invocations),
            },
        )
    write_summary(options.summary, summary)
    return summary


def _parse(argv: Sequence[str] | None) -> WorkerOptions:
    parser = argparse.ArgumentParser(description="Run a Nida benchmark worker.")
    parser.add_argument("--requests", type=Path, required=True)
    parser.add_argument("--results", type=Path, required=True)
    parser.add_argument("--summary", type=Path, default=None)
    parser.add_argument("--model-id", required=True)
    parser.add_argument("--revision", default="")
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--retry-failed", action="store_true")
    args = parser.parse_args(argv)
    summary = args.summary or args.results.with_name(args.results.stem + ".summary.json")
    return WorkerOptions(
        requests=args.requests,
        results=args.results,
        summary=summary,
        model=ModelRef(model_id=args.model_id, revision=args.revision),
        device=args.device,
        batch_size=args.batch_size,
        retry_failed=args.retry_failed,
    )


def main_for(adapter: Adapter, argv: Sequence[str] | None = None) -> int:
    """CLI entry point used by each worker module's ``__main__`` block."""
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    options = _parse(argv)
    try:
        summary = run_worker(adapter, options)
    except (ProtocolError, OSError, ValueError, TooManyFailuresError):
        traceback.print_exc()
        return 2
    print(f"{adapter.name}: ok={summary.n_ok} error={summary.n_error}")
    return 0