"""JSONL protocol between the benchmark orchestrator and model workers.

Each model family runs in its own virtual environment (dependency pins conflict, see
docs/benchmarks/B1-B2-tasarim.md §1.3), so the orchestrator and the workers only share
files. This module is the single definition of those files and must stay stdlib-only:
it is imported inside every worker environment.

Layout for one worker invocation::

    requests.jsonl        one ``WorkRequest`` per line (written by the orchestrator)
    results.jsonl         one ``WorkResult`` per line (appended by the worker)
    results.summary.json  one ``WorkerSummary`` (written by the worker at the end)

Results are appended one line at a time and flushed + fsynced, so an interrupted worker
leaves a valid prefix plus at most one partial trailing line. On restart the worker
calls :func:`load_results` with ``repair=True``, which drops the partial line, and
skips every request whose ``clip_id`` already has a result. Failed items are recorded
with an ``error`` and are never retried silently: callers decide via ``retry_failed``.
"""

from __future__ import annotations

import json
import os
from collections.abc import Iterable, Iterator, Mapping
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import IO

__all__ = [
    "ProtocolError",
    "WorkRequest",
    "WorkResult",
    "WorkerSummary",
    "append_result",
    "drop_results",
    "load_requests",
    "load_results",
    "pending_requests",
    "read_summary",
    "write_requests",
    "write_summary",
]


class ProtocolError(ValueError):
    """A protocol file is malformed or inconsistent."""


@dataclass(frozen=True)
class WorkRequest:
    """One unit of work for a worker.

    Attributes:
        clip_id: Content-addressed clip id; unique within a requests file.
        audio_path: Absolute path of the input audio file.
        params: Worker-specific string parameters (e.g. ``language``, ``context``,
            ``out_path``). Values are strings so the format stays trivially portable.
    """

    clip_id: str
    audio_path: str
    params: Mapping[str, str] = field(default_factory=dict)

    def to_json(self) -> str:
        return json.dumps(
            {"clip_id": self.clip_id, "audio_path": self.audio_path, "params": dict(self.params)},
            ensure_ascii=False,
            sort_keys=True,
        )

    @staticmethod
    def from_obj(obj: object, where: str) -> WorkRequest:
        if not isinstance(obj, dict):
            raise ProtocolError(f"{where}: request must be a JSON object")
        clip_id = obj.get("clip_id")
        audio_path = obj.get("audio_path")
        params = obj.get("params", {})
        if not isinstance(clip_id, str) or not clip_id:
            raise ProtocolError(f"{where}: clip_id must be a non-empty string")
        if not isinstance(audio_path, str) or not audio_path:
            raise ProtocolError(f"{where}: audio_path must be a non-empty string")
        if not isinstance(params, dict) or not all(
            isinstance(k, str) and isinstance(v, str) for k, v in params.items()
        ):
            raise ProtocolError(f"{where}: params must map strings to strings")
        return WorkRequest(clip_id=clip_id, audio_path=audio_path, params=dict(params))


@dataclass(frozen=True)
class WorkResult:
    """Outcome for one request: exactly one of ``output`` or ``error`` is set.

    Attributes:
        clip_id: Id of the request this answers.
        output: Worker-specific JSON-compatible payload (e.g. ``{"text": ...}``).
        error: ``"<ExceptionType>: <message>"`` when the item failed.
        infer_seconds: Wall-clock inference time attributed to this item (batch time
            divided by batch size), excluding model loading.
    """

    clip_id: str
    output: Mapping[str, object] | None
    error: str | None
    infer_seconds: float

    def __post_init__(self) -> None:
        if (self.output is None) == (self.error is None):
            raise ProtocolError(f"{self.clip_id}: exactly one of output/error must be set")
        if self.infer_seconds < 0:
            raise ProtocolError(f"{self.clip_id}: infer_seconds must be >= 0")

    @property
    def ok(self) -> bool:
        return self.error is None

    def to_json(self) -> str:
        return json.dumps(
            {
                "clip_id": self.clip_id,
                "output": None if self.output is None else dict(self.output),
                "error": self.error,
                "infer_seconds": self.infer_seconds,
            },
            ensure_ascii=False,
            sort_keys=True,
        )

    @staticmethod
    def from_obj(obj: object, where: str) -> WorkResult:
        if not isinstance(obj, dict):
            raise ProtocolError(f"{where}: result must be a JSON object")
        clip_id = obj.get("clip_id")
        output = obj.get("output")
        error = obj.get("error")
        seconds = obj.get("infer_seconds")
        if not isinstance(clip_id, str) or not clip_id:
            raise ProtocolError(f"{where}: clip_id must be a non-empty string")
        if output is not None and not isinstance(output, dict):
            raise ProtocolError(f"{where}: output must be an object or null")
        if error is not None and not isinstance(error, str):
            raise ProtocolError(f"{where}: error must be a string or null")
        if isinstance(seconds, bool) or not isinstance(seconds, int | float):
            raise ProtocolError(f"{where}: infer_seconds must be a number")
        return WorkResult(clip_id=clip_id, output=output, error=error, infer_seconds=float(seconds))


@dataclass(frozen=True)
class WorkerSummary:
    """Facts about one worker invocation, written once at the end."""

    worker: str
    model_id: str
    revision: str
    load_seconds: float
    peak_vram_gb: float | None
    n_ok: int
    n_error: int
    started_at: str
    ended_at: str
    packages: Mapping[str, str] = field(default_factory=dict)
    extra: Mapping[str, str] = field(default_factory=dict)


def _fsync_line(handle: IO[str], line: str) -> None:
    handle.write(line + "\n")
    handle.flush()
    os.fsync(handle.fileno())


def write_requests(path: Path, requests: Iterable[WorkRequest]) -> int:
    """Write a requests file atomically and return the number of requests.

    Raises:
        ProtocolError: on duplicate clip ids.
    """
    from nida.utils.io import atomic_write_text  # local import: workers never call this

    seen: set[str] = set()
    lines: list[str] = []
    for request in requests:
        if request.clip_id in seen:
            raise ProtocolError(f"duplicate clip_id in requests: {request.clip_id}")
        seen.add(request.clip_id)
        lines.append(request.to_json())
    atomic_write_text(path, "".join(line + "\n" for line in lines))
    return len(lines)


def _iter_json_lines(path: Path) -> Iterator[tuple[int, str]]:
    with path.open("r", encoding="utf-8") as handle:
        yield from enumerate(handle, start=1)


def load_requests(path: Path) -> list[WorkRequest]:
    """Read and validate a requests file (no partial lines tolerated)."""
    requests: list[WorkRequest] = []
    seen: set[str] = set()
    for number, line in _iter_json_lines(path):
        where = f"{path}:{number}"
        if not line.endswith("\n"):
            raise ProtocolError(f"{where}: truncated request line")
        try:
            obj = json.loads(line)
        except json.JSONDecodeError as err:
            raise ProtocolError(f"{where}: invalid JSON ({err.msg})") from err
        request = WorkRequest.from_obj(obj, where)
        if request.clip_id in seen:
            raise ProtocolError(f"{where}: duplicate clip_id {request.clip_id}")
        seen.add(request.clip_id)
        requests.append(request)
    return requests


def load_results(path: Path, *, repair: bool = False) -> dict[str, WorkResult]:
    """Read a results file into ``{clip_id: result}``.

    A missing file yields an empty mapping. Only the *last* line may be incomplete (an
    interrupted append): with ``repair=True`` it is removed from the file, otherwise a
    ``ProtocolError`` is raised. Any malformed line before the last one, and any
    duplicate clip id, is always an error.
    """
    if not path.exists():
        return {}
    raw = path.read_bytes()
    # Split at the last newline *in bytes*: an interrupted append may end in the middle of
    # a multi-byte UTF-8 character, so the trailing fragment must not be decoded.
    keep = raw.rfind(b"\n") + 1  # 0 when there is no complete line
    body, trailing = raw[:keep], raw[keep:]
    try:
        text = body.decode("utf-8")
    except UnicodeDecodeError as err:
        raise ProtocolError(f"{path}: complete lines are not valid UTF-8 ({err})") from err
    lines = text.split("\n")[:-1] if text else []
    results: dict[str, WorkResult] = {}
    for number, line in enumerate(lines, start=1):
        where = f"{path}:{number}"
        try:
            obj = json.loads(line)
        except json.JSONDecodeError as err:
            raise ProtocolError(f"{where}: invalid JSON ({err.msg})") from err
        result = WorkResult.from_obj(obj, where)
        if result.clip_id in results:
            raise ProtocolError(f"{where}: duplicate result for {result.clip_id}")
        results[result.clip_id] = result
    if trailing:
        if not repair:
            raise ProtocolError(f"{path}: incomplete trailing line ({len(trailing)} bytes)")
        with path.open("r+b") as handle:
            handle.truncate(keep)
            handle.flush()
            os.fsync(handle.fileno())
    return results


def pending_requests(
    requests: Iterable[WorkRequest],
    results: Mapping[str, WorkResult],
    *,
    retry_failed: bool = False,
) -> list[WorkRequest]:
    """Requests that still need work, in input order.

    Raises:
        ProtocolError: if ``results`` contains clip ids that are not in ``requests``
            (the results file belongs to a different request set).
    """
    request_list = list(requests)
    known = {r.clip_id for r in request_list}
    stray = sorted(set(results) - known)
    if stray:
        raise ProtocolError(f"results contain unknown clip ids, e.g. {stray[:3]}")
    pending: list[WorkRequest] = []
    for request in request_list:
        done = results.get(request.clip_id)
        if done is None or (retry_failed and not done.ok):
            pending.append(request)
    return pending


def append_result(handle: IO[str], result: WorkResult) -> None:
    """Append one result line and fsync, so an interruption loses at most this line."""
    _fsync_line(handle, result.to_json())


def write_summary(path: Path, summary: WorkerSummary) -> None:
    """Write the summary JSON next to the results (stdlib-only atomic replace)."""
    tmp = path.with_name(f".{path.name}.tmp")
    with tmp.open("w", encoding="utf-8") as handle:
        json.dump(asdict(summary), handle, ensure_ascii=False, indent=2, sort_keys=True)
        handle.write("\n")
        handle.flush()
        os.fsync(handle.fileno())
    tmp.replace(path)


def read_summary(path: Path) -> WorkerSummary:
    """Read and validate a summary file."""
    obj = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(obj, dict):
        raise ProtocolError(f"{path}: summary must be a JSON object")
    try:
        return WorkerSummary(**obj)
    except TypeError as err:
        raise ProtocolError(f"{path}: invalid summary fields ({err})") from err


def drop_results(path: Path, clip_ids: Iterable[str]) -> int:
    """Remove the given clip ids from a results file (atomic rewrite); returns how many.

    Used when a result's side effect is gone (e.g. its output file was lost), so the
    item becomes pending again instead of being trusted.
    """
    doomed = set(clip_ids)
    results = load_results(path, repair=True)
    kept = [r.to_json() + "\n" for cid, r in results.items() if cid not in doomed]
    removed = len(results) - len(kept)
    if removed:
        tmp = path.with_name(f".{path.name}.tmp")
        with tmp.open("w", encoding="utf-8") as handle:
            handle.writelines(kept)
            handle.flush()
            os.fsync(handle.fileno())
        tmp.replace(path)
    return removed