"""BenchFlow 输入与运行时轨迹适配。"""

from __future__ import annotations

import json
import re
import shutil
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from scripts.dynamic_compile.fast.models import RolloutTrace
from scripts.dynamic_compile.fast.traces.benchflow import trajectory_for_test
from scripts.dynamic_compile.fast.scoring.state import acp_events_to_state, read_acp_events


@dataclass
class BenchFlowInput:
    task_name: str
    task_dir: Path
    agent: str
    model: str
    prompt: str
    traces: list[RolloutTrace]


def _project_root() -> Path:
    return Path(__file__).resolve().parents[4]


def _run_config(test_dir: Path) -> dict[str, Any]:
    paths = sorted(test_dir.rglob("config.json"))
    if not paths:
        raise ValueError(f"missing BenchFlow run config under {test_dir}")
    value = json.loads(paths[0].read_text(encoding="utf-8"))
    if not isinstance(value, dict):
        raise ValueError(f"invalid BenchFlow run config: {paths[0]}")
    return value


def _prompt(test_dir: Path) -> str:
    paths = sorted(test_dir.rglob("prompts.json"))
    if not paths:
        raise ValueError(f"missing BenchFlow prompts.json under {test_dir}")
    value = json.loads(paths[0].read_text(encoding="utf-8"))
    if not isinstance(value, list) or not value or not isinstance(value[0], str):
        raise ValueError(f"invalid BenchFlow prompts: {paths[0]}")
    return value[0].strip()


def load_runtime_trace(test_dir: Path, task_name: str, compile_type: str) -> RolloutTrace:
    """Load agent runtime evidence without opening verifier/result artifacts."""
    trajectory = trajectory_for_test(test_dir)
    if trajectory is None:
        raise ValueError(f"missing acp_trajectory.jsonl under {test_dir}")
    events = read_acp_events(trajectory)
    state = acp_events_to_state(events)
    skill_invoked = any(
        event.get("type") == "tool_call"
        and event.get("status") == "completed"
        and any(
            str(event.get(field, "")).strip().lower() == "skill"
            for field in ("title", "kind")
        )
        for event in events
    )
    timed_out = any(event.get("type") == "agent_timeout" for event in events)
    return RolloutTrace(
        trace_id=f"{task_name}/{compile_type}/{test_dir.name}",
        task_name=task_name,
        compile_type=compile_type,
        test_name=test_dir.name,
        state=state,
        skill_invoked=skill_invoked,
        timed_out=timed_out,
        metadata={
            "termination": "timeout" if timed_out else "completed",
            "tool_calls": sum(event.get("type") == "tool_call" for event in events),
        },
    )


def load_benchflow_input(source: Path) -> BenchFlowInput:
    source = source.resolve()
    tests = sorted(path for path in source.glob("test-*") if path.is_dir())[-5:]
    if not tests:
        raise ValueError(f"no test-* BenchFlow traces under {source}")
    task_name = source.parent.name
    traces = [load_runtime_trace(test, task_name, "custom") for test in tests]
    configs = [_run_config(test) for test in tests]
    agents = {str(item.get("agent", "")).strip() for item in configs}
    models = {str(item.get("model", "")).strip() for item in configs}
    if "" in agents or len(agents) != 1:
        raise ValueError(f"BenchFlow traces do not identify one agent: {sorted(agents)}")
    if "" in models or len(models) != 1:
        raise ValueError(f"BenchFlow traces do not identify one model: {sorted(models)}")
    prompts = {_prompt(test) for test in tests}
    if len(prompts) != 1:
        raise ValueError(f"BenchFlow traces contain {len(prompts)} different task prompts")
    task_dir = _project_root() / "data" / "skills-bench" / "tasks" / task_name
    if not (task_dir / "task.md").is_file():
        raise ValueError(f"cannot resolve SkillsBench task directory: {task_dir}")
    return BenchFlowInput(
        task_name, task_dir, next(iter(agents)), next(iter(models)),
        next(iter(prompts)), traces,
    )


def _slug(value: str) -> str:
    return re.sub(r"^-+|-+$", "", re.sub(r"[^a-z0-9]+", "-", value.lower()))


class SkillsBenchDevelopmentRolloutRunner:
    """Development-only runner; verifier output is never projected into traces."""

    def __init__(
        self,
        context: BenchFlowInput,
        work_root: Path,
        max_parallel: int = 3,
        archive_root: Path | None = None,
    ):
        self.context = context
        self.work_root = work_root
        self.max_parallel = max_parallel
        self.archive_root = archive_root

    def _variant_dir(self, jobs_root: Path) -> Path:
        return (
            jobs_root
            / _slug(f"{self.context.agent}-{self.context.model}")
            / self.context.task_name
            / "custom_skill"
        )

    def _completed_traces(self, variant: Path, batch_id: str) -> list[RolloutTrace]:
        completed: list[RolloutTrace] = []
        for test in sorted(path for path in variant.glob("test-*") if path.is_dir()):
            try:
                requirement = json.loads(
                    (test / "required-skill.json").read_text(encoding="utf-8")
                )
                if requirement.get("invoked") is not True:
                    continue
                trace = load_runtime_trace(test, self.context.task_name, batch_id)
            except (AttributeError, json.JSONDecodeError, OSError, ValueError):
                continue
            completed.append(trace)
        return completed

    def artifacts_dir(self, batch_id: str) -> Path | None:
        if self.archive_root is None:
            return None
        return self.archive_root / batch_id / "custom_skill"

    def run_batch(
        self,
        skill_package: Path,
        _prompt: str,
        batch_id: str,
        _task_name: str,
        count: int,
        progress: Any | None = None,
    ) -> list[RolloutTrace]:
        jobs_root = self.work_root / batch_id
        log_path = jobs_root / "runner.log"
        variant = self._variant_dir(jobs_root)
        archive = self.artifacts_dir(batch_id)
        if archive is not None and archive.is_dir():
            archived_traces = self._completed_traces(archive, batch_id)
            if len(archived_traces) >= count:
                traces = archived_traces
                missing = 0
            else:
                variant.parent.mkdir(parents=True, exist_ok=True)
                shutil.copytree(archive, variant, dirs_exist_ok=True)
                traces = self._completed_traces(variant, batch_id)
                missing = max(0, count - len(traces))
        else:
            traces = self._completed_traces(variant, batch_id)
            missing = max(0, count - len(traces))
        if missing:
            jobs_root.mkdir(parents=True, exist_ok=True)
            command = [
                "bash", str(_project_root() / "scripts" / "evaluate" / "run-raw-task.sh"),
                "--harness", self.context.agent,
                "--model", self.context.model,
                "--task", str(self.context.task_dir),
                "--skill-source", str(skill_package.resolve()),
                "--require-skill",
                "--repeat", str(missing),
                "--max-parallel", str(self.max_parallel),
                "--output", str(variant),
            ]
            with log_path.open("a", encoding="utf-8") as handle:
                subprocess.run(command, stdout=handle, stderr=subprocess.STDOUT, text=True)
            traces = self._completed_traces(variant, batch_id)
        if archive is not None and variant.is_dir():
            archive.parent.mkdir(parents=True, exist_ok=True)
            shutil.copytree(variant, archive, dirs_exist_ok=True)
            runner_log = jobs_root / "runner.log"
            if runner_log.is_file():
                shutil.copy2(runner_log, archive.parent / "runner.log")
            traces = self._completed_traces(archive, batch_id)
        if len(traces) < count:
            raise RuntimeError(
                f"BenchFlow produced {len(traces)}/{count} candidate traces; see {log_path}"
            )
        traces = traces[:count]
        for index, trace in enumerate(traces, 1):
            trace.trace_id = f"{batch_id}-{index:03d}"
            trace.test_name = trace.trace_id
            if progress is not None:
                progress(index, len(traces), trace.trace_id)
        return traces