from __future__ import annotations

import json
from pathlib import Path
from typing import Any

from ..models import RolloutTrace
from ..scoring.state import acp_events_to_state, read_acp_events


VARIANT_TO_COMPILE_TYPE = {
    "model_skill": "model_compile",
    "ori_skill": "ori",
}


def event_to_state(trajectory_path: Path) -> tuple[list[dict[str, str]], bool]:
    events = read_acp_events(trajectory_path)
    skill_invoked = any(
        event.get("type") == "tool_call"
        and any(
            str(event.get(field, "")).strip().lower() == "skill"
            for field in ("title", "kind")
        )
        for event in events
    )
    return acp_events_to_state(events), skill_invoked


def trajectory_for_test(test_dir: Path) -> Path | None:
    candidates = sorted(test_dir.rglob("acp_trajectory.jsonl"))
    if not candidates:
        return None
    canonical = [path for path in candidates if "trajectory" in path.parts]
    return canonical[0] if canonical else candidates[0]


def _result_for_trajectory(trajectory_path: Path) -> dict[str, Any]:
    result_path = trajectory_path.parent.parent / "result.json"
    if not result_path.is_file():
        raise ValueError(f"missing structured BenchFlow result: {result_path}")
    value = json.loads(result_path.read_text(encoding="utf-8"))
    if not isinstance(value, dict):
        raise ValueError(f"BenchFlow result must be a JSON object: {result_path}")
    return value


def load_benchflow_trace(
    test_dir: Path,
    task_name: str,
    compile_type: str,
) -> RolloutTrace:
    trajectory = trajectory_for_test(test_dir)
    if trajectory is None:
        raise ValueError(f"missing acp_trajectory.jsonl under {test_dir}")
    result = _result_for_trajectory(trajectory)
    agent_timeout = result.get("agent_timeout_info")
    idle_timeout = result.get("idle_timeout_info")
    timeout_info = (
        agent_timeout
        if isinstance(agent_timeout, dict)
        else idle_timeout if isinstance(idle_timeout, dict) else None
    )
    timed_out = timeout_info is not None
    metadata: dict[str, Any] = {
        "source": str(test_dir.resolve()),
        "termination": "timeout" if timed_out else "completed",
    }
    timing = result.get("timing") if isinstance(result.get("timing"), dict) else {}
    execution_seconds = timing.get("agent_execution")
    if execution_seconds is None and isinstance(timeout_info, dict):
        execution_seconds = timeout_info.get(
            "wall_clock_elapsed_sec", timeout_info.get("timeout_sec")
        )
    if isinstance(execution_seconds, (int, float)):
        metadata["agent_execution_seconds"] = float(execution_seconds)
    if isinstance(result.get("n_tool_calls"), int):
        metadata["tool_calls"] = result["n_tool_calls"]
    if timed_out:
        metadata["timeout_reason"] = timeout_info.get("reason")
        metadata["timeout_seconds"] = timeout_info.get(
            "timeout_sec", timeout_info.get("idle_timeout_sec")
        )
        metadata["partial_trajectory"] = bool(result.get("partial_trajectory", False))
        metadata["error_category"] = result.get("error_category")
    state, skill_invoked = event_to_state(trajectory)
    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=metadata,
    )


def _variant_dirs(input_path: Path) -> list[tuple[Path, str, str]]:
    if input_path.name in VARIANT_TO_COMPILE_TYPE:
        return [(
            input_path,
            input_path.parent.name,
            VARIANT_TO_COMPILE_TYPE[input_path.name],
        )]
    direct_variants = [
        (input_path / variant_name, input_path.name, compile_type)
        for variant_name, compile_type in VARIANT_TO_COMPILE_TYPE.items()
        if (input_path / variant_name).is_dir()
    ]
    if direct_variants:
        return direct_variants
    variants = []
    for task_dir in sorted(path for path in input_path.iterdir() if path.is_dir()):
        for variant_name, compile_type in VARIANT_TO_COMPILE_TYPE.items():
            variant_dir = task_dir / variant_name
            if variant_dir.is_dir():
                variants.append((variant_dir, task_dir.name, compile_type))
    return variants


def load_benchflow_traces(input_path: Path) -> list[RolloutTrace]:
    input_path = input_path.resolve()
    if not input_path.is_dir():
        raise ValueError(f"BenchFlow input does not exist: {input_path}")
    variants = _variant_dirs(input_path)
    if not variants:
        raise ValueError(f"no model_skill or ori_skill directories under {input_path}")
    traces = []
    for variant_dir, task_name, compile_type in variants:
        for test_dir in sorted(path for path in variant_dir.glob("test-*") if path.is_dir()):
            traces.append(load_benchflow_trace(test_dir, task_name, compile_type))
    return traces