"""Running one task once with one driver model, and recording it."""

import io
import shlex
import subprocess
import tempfile
import time
import uuid
from dataclasses import asdict, dataclass, field
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from nisse_core.execution.ssh import SSHHost
from nisse_core.log.canonical import (
    CanonicalEvent,
    Fact,
    StoredEvent,
    ToolCall,
    ToolResult,
    UserMessage,
)
from nisse_core.loop.driver import Driver, RunOutcome
from nisse_core.models.catalog import ModelFacts
from nisse_core.models.compact import Compactor
from nisse_core.oplog import OpLog
from nisse_dev.bench.tasks import Task
from nisse_dev.config import Tuning
from nisse_dev.loop.agent import SYSTEM
from nisse_dev.loop.tools import (
    ASK_HUMAN,
    CLOSE_ITEM,
    FINISH_MILESTONE,
    START_MILESTONE,
    WAITING_TOOLS,
)
from nisse_dev.state import StateStore
from nisse_dev.tools import EXECUTION_TOOLS, Workspace

# What the driver is offered once an item is agreed, less what the benchmark
# turns off: delegation and ask_expert (an escalation is the model's failure,
# not a frontier model's success) and review, which has no PR here.
BENCH_TOOLS = [
    ASK_HUMAN,
    START_MILESTONE,
    FINISH_MILESTONE,
    CLOSE_ITEM,
    *EXECUTION_TOOLS,
]


def outcome_of(run: RunOutcome) -> str:
    """A driver run as the benchmark records it: ``completed`` when the
    driver says the work is done, ``escalated`` when it asks the operator,
    otherwise how the run ended (capped, stuck, failed)."""
    if run.kind == "wait":
        return "escalated" if run.tool == "ask_human" else "completed"
    return run.kind


class MeteringClient:
    """Counts what a run spends, and notices a fallback: an answer from a
    model other than the one asked for is not a measurement of it."""

    def __init__(self, inner: Any, facts: ModelFacts) -> None:
        self._inner = inner
        self._facts = facts
        self.calls = 0
        self.prose_turns = 0
        self.tokens_in = 0
        self.tokens_out = 0
        self.cost = 0.0
        self.fallbacks: list[str] = []
        self.rate_limit_wait_s = 0.0

    def complete(
        self,
        events: list[StoredEvent],
        tools: list[dict[str, Any]] | None,
        model: str,
        system: str | None = None,
        **options: Any,
    ) -> list[CanonicalEvent]:
        self.calls += 1
        try:
            produced = self._inner.complete(
                events, tools, model, system=system, **options
            )
        finally:
            # Counted even when the call fails: a rate limit that never
            # cleared still cost the run its wait.
            self.rate_limit_wait_s += float(
                getattr(self._inner, "last_rate_limit_wait_s", 0.0) or 0.0
            )
        usage = getattr(self._inner, "last_usage", {}) or {}
        self.tokens_in += int(usage.get("input", 0))
        self.tokens_out += int(usage.get("output", 0))
        self.cost += float(getattr(self._inner, "last_cost", 0.0) or 0.0)
        answered = getattr(self._inner, "last_model", None)
        if not self._facts.answered(answered):
            self.fallbacks.append(str(answered))
        # Compaction's notes are prose by request, not a broken turn.
        asked_for_prose = options.get("require_tool_call") is False
        if not asked_for_prose and not any(isinstance(e, ToolCall) for e in produced):
            self.prose_turns += 1
        return produced

    def __getattr__(self, name: str) -> Any:
        # The budget (budget_for, estimate) compaction reads.
        return getattr(self._inner, name)


@dataclass
class RunRecord:
    """One run, with everything needed to compare it with runs made later
    against other models, other servers, or other tuning."""

    ts: str
    model: str
    backing: str
    provider: str
    model_info: dict[str, Any]
    note: str
    task: str
    task_hash: str
    commit: str
    tuning: dict[str, Any]
    success: bool
    outcome: str
    outcome_note: str
    turns: int
    prose_turns: int
    failed_commands: int
    invented_tools: int
    trimmed_requests: int
    tokens_in: int
    tokens_out: int
    cost: float
    wall_s: float
    fallbacks: list[str] = field(default_factory=list)
    # Part of wall_s spent waiting out the proxy's rate limit, not working.
    rate_limit_wait_s: float = 0.0
    check_output: str = ""

    def to_json(self) -> dict[str, Any]:
        return asdict(self)


def agent_commit(repo: Path | None = None) -> str:
    """The agent's commit, marked when the tree had uncommitted changes."""
    where = repo or Path(__file__).resolve().parent
    try:
        sha = subprocess.run(
            ["git", "rev-parse", "--short", "HEAD"],
            cwd=where,
            capture_output=True,
            text=True,
            check=True,
        ).stdout.strip()
        dirty = subprocess.run(
            ["git", "status", "--porcelain"],
            cwd=where,
            capture_output=True,
            text=True,
            check=True,
        ).stdout.strip()
    except (OSError, subprocess.CalledProcessError):
        return "unknown"
    return f"{sha}+dirty" if dirty else sha


def _count(events: list[StoredEvent]) -> tuple[int, int]:
    tools = {
        e.event.call_id: e.event.tool for e in events if isinstance(e.event, ToolCall)
    }
    failed = invented = 0
    for stored in events:
        event = stored.event
        if isinstance(event, ToolResult) and not event.success:
            if event.result.startswith("no such tool"):
                invented += 1
            elif tools.get(event.call_id) == "run_command":
                failed += 1
    return failed, invented


def run_task(
    task: Task,
    facts: ModelFacts,
    client: Any,
    tuning: Tuning,
    *,
    host: SSHHost,
    note: str = "",
    commit: str = "",
    workdir: Path | None = None,
) -> RunRecord:
    """Run one task once: seed a workspace, hand the driver an agreed item
    whose restatement is the task's plan, then judge the result with the
    hidden check.

    It is the loop's own driver, system prompt and tools, so what is measured
    is what the agent would do, without the harness around it: the milestone
    is started but no branch is cut, and finishing it ends the run."""
    with tempfile.TemporaryDirectory(dir=workdir) as tmp:
        root = Path(tmp)
        # The workspace is on the build VM, like the agent's own (L13).
        workspace = Workspace(
            host,
            f"/tmp/nisse-bench/{uuid.uuid4().hex}",
            timeout_s=tuning.tools.command_timeout_s,
            max_output_chars=tuning.tools.max_output_chars,
        )
        task.seed_workspace(workspace)
        store = StateStore(root / "state.db", oplog=OpLog("nisse-bench", io.StringIO()))
        # One work item per run, as the agent's own are (L1).
        store.open_work_item(
            "bench", title=task.name, position="RUNNING", cause="benchmark run"
        )
        store.transcript("bench", "orchestrator").append(
            Fact(text=f"The operator agreed to this restatement:\n\n{task.plan}")
        )
        store.transcript("bench", "operator").append(
            UserMessage(text="go", source="benchmark")
        )

        meter = MeteringClient(client, facts)
        trims_before = getattr(client, "trimmed_requests", 0)
        model = facts.alias
        driver = Driver(
            meter,
            model,
            max_turns=tuning.driver.max_turns,
            failure_limit=tuning.driver.failure_limit,
            # As the agent runs it, so a long task measures the same thing.
            compactor=Compactor(meter, tuning.context.compaction()),
        )
        system = SYSTEM.format(
            repo=f"bench/{task.name}",
            issue=1,
            title=task.title,
            labels="",
            helpers="",
            workspace=workspace.describe(),
        )
        started = time.monotonic()
        try:
            run = driver.run(
                store.transcript("bench", "driver"),
                system=system,
                tools=BENCH_TOOLS,
                waiting=WAITING_TOOLS,
                workspace=workspace,
                handlers={"start_milestone": lambda args: ("started", True)},
            )
            status, detail = outcome_of(run), run.note
        except Exception as exc:  # noqa: BLE001 -- recorded, the benchmark goes on
            status, detail = "error", repr(exc)
        wall = time.monotonic() - started

        success, check_output = task.check_workspace(workspace)
        host.execute(f"rm -rf -- {shlex.quote(str(workspace.root))}", timeout_s=60)
        failed, invented = _count(store.transcript("bench", "driver").read())
        store.close()

    return RunRecord(
        ts=datetime.now(UTC).isoformat(timespec="seconds"),
        model=model,
        backing=facts.backing,
        provider=facts.provider,
        model_info={
            "max_input_tokens": facts.max_input_tokens,
            "max_output_tokens": facts.max_output_tokens,
            "supports_tool_choice": facts.supports_tool_choice,
            "supports_function_calling": facts.supports_tools,
        },
        note=note,
        task=task.name,
        task_hash=task.content_hash,
        commit=commit,
        tuning=asdict(tuning),
        success=success and not meter.fallbacks,
        outcome=status,
        outcome_note=detail[:500],
        turns=meter.calls,
        prose_turns=meter.prose_turns,
        failed_commands=failed,
        invented_tools=invented,
        trimmed_requests=getattr(client, "trimmed_requests", 0) - trims_before,
        tokens_in=meter.tokens_in,
        tokens_out=meter.tokens_out,
        cost=round(meter.cost, 6),
        wall_s=round(wall, 1),
        fallbacks=meter.fallbacks,
        rate_limit_wait_s=meter.rate_limit_wait_s,
        check_output=check_output,
    )