"""drive one task through Pi (the agent harness) and grade the result.

Pi runs headless (`pi -p --mode json`) with built-in read/bash/edit/write tools and
native `--skill` loading; we parse its JSON event stream for the tool-call trace, then
grade the produced `solution.zig` with the compiler oracle.
"""

from __future__ import annotations

import json
import os
import subprocess
import tempfile
from dataclasses import dataclass, field
from pathlib import Path

from .config import Config, find_pi
from .tasks import Task, grade

PROMPT_SUFFIX = (
    " Write your complete solution to a file named solution.zig in the current "
    "directory, then verify it compiles by running `zig test solution.zig`."
)


@dataclass
class ArmResult:
    task: str
    arm: str  # "bare" | "zigman"
    oracle_ok: bool
    detail: str
    timed_out: bool
    tools: list[str] = field(default_factory=list)
    zigman_queries: list[str] = field(default_factory=list)
    code: str = ""

    @property
    def used_zigman(self) -> bool:
        return bool(self.zigman_queries)


def _command(cfg: Config, prompt: str, with_zigman: bool) -> list[str]:
    cmd = [find_pi(), "-p", "--mode", "json", "--no-session",
           "--provider", cfg.provider, "--model", cfg.model]
    if with_zigman and cfg.skill_dir is not None:
        cmd += ["--skill", str(cfg.skill_dir)]
    else:
        cmd += ["--no-skills"]
    cmd.append(prompt)
    return cmd


def _env(cfg: Config) -> dict[str, str]:
    env = dict(os.environ)
    # make the zigman CLI reachable from Pi's bash tool without assuming $HOME
    if cfg.zigman_bin:
        env["PATH"] = str(Path(cfg.zigman_bin).resolve().parent) + os.pathsep + env.get("PATH", "")
    return env


def parse_events(stdout: str) -> tuple[list[dict], list[str]]:
    """extract completed tool calls (deduped by id) from Pi's JSON event stream.

    the reliable source is `message_end` events for assistant messages, whose content
    carries finished `toolCall` items (partial `toolcall_*` deltas are noise).
    """
    calls: dict[str, dict] = {}
    for line in stdout.splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            e = json.loads(line)
        except json.JSONDecodeError:
            continue
        msg = e.get("message") or {}
        if e.get("type") == "message_end" and msg.get("role") == "assistant":
            for part in msg.get("content", []):
                if part.get("type") == "toolCall" and part.get("id"):
                    calls[part["id"]] = {"name": part.get("name"), "arguments": part.get("arguments", {})}
    tools = list(calls.values())
    zigman_qs = [
        str(c["arguments"].get("command", ""))
        for c in tools
        if c["name"] == "bash" and "zigman" in str(c["arguments"].get("command", ""))
    ]
    return tools, zigman_qs


def run_task(cfg: Config, task: Task, with_zigman: bool) -> ArmResult:
    arm = "zigman" if with_zigman else "bare"
    workdir = Path(tempfile.mkdtemp(prefix="pi-eval-"))
    try:
        cmd = _command(cfg, task.prompt + PROMPT_SUFFIX, with_zigman)
        timed_out = False
        try:
            proc = subprocess.run(cmd, cwd=workdir, env=_env(cfg),
                                  capture_output=True, text=True, timeout=cfg.timeout)
            stdout = proc.stdout
        except subprocess.TimeoutExpired as e:
            stdout = e.stdout.decode() if isinstance(e.stdout, bytes) else (e.stdout or "")
            timed_out = True

        tools, zigman_qs = parse_events(stdout)
        sol = workdir / "solution.zig"
        if not sol.exists():
            cands = sorted(workdir.glob("*.zig"))
            sol = cands[0] if cands else None
        if sol and sol.exists():
            code = sol.read_text()
            oracle_ok, detail = grade(code, task, zig=cfg.zig)
        else:
            code, oracle_ok, detail = "", False, "no .zig file produced"

        return ArmResult(
            task=task.id, arm=arm, oracle_ok=oracle_ok, detail=detail, timed_out=timed_out,
            tools=[t["name"] for t in tools], zigman_queries=zigman_qs, code=code,
        )
    finally:
        import shutil
        shutil.rmtree(workdir, ignore_errors=True)
