"""Prime Verifiers v1 episode JSONL enters Bench without a second harness."""

from __future__ import annotations

import json
import sys
import tempfile
import unittest
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

from kbench.verifiers_import import (
    VerifiersImportError,
    import_quality,
    read_episodes,
)


def trace(
    trace_id: str,
    reward: float,
    *,
    agent: str = "agent",
    trainable: bool = True,
    answer: str = "done",
) -> dict:
    return {
        "id": trace_id,
        "verifiers": {"version": "0.3.1"},
        "agent": {"name": agent, "trainable": trainable},
        "ok": True,
        "rewards": {"task": {"score": reward, "weight": 1.0}},
        "metrics": {"strict": reward},
        "nodes": [{"message": {"role": "assistant", "content": answer}}],
        "timing": {"agent": {"start": 10.0, "end": 12.0}},
    }


def episode(episode_id: str, traces: list[dict], task_key: str = "task-1") -> dict:
    return {
        "id": episode_id,
        "env": {"id": "example-tool-loop-v1"},
        "task": {
            "type": "OfficeJobTask",
            "key": task_key,
            "data": {"private_prompt": "must never enter the score card"},
        },
        "ok": True,
        "errors": [],
        "traces": traces,
    }


class VerifiersImportTests(unittest.TestCase):
    def write(self, root: Path, rows: list[dict]) -> Path:
        path = root / "traces.jsonl"
        path.write_text("".join(json.dumps(row) + "\n" for row in rows))
        return path

    def test_one_episode_becomes_one_existing_sample_outcome(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            path = self.write(Path(tmp), [episode("ep-1", [trace("tr-1", 1.0)])])
            result = import_quality(path, "example-tool-loop")
        self.assertEqual(result.n_samples, 1)
        self.assertEqual(result.metrics["accuracy"], 1.0)
        self.assertEqual(result.samples[0].sample_id, "ep-1")
        self.assertEqual(result.samples[0].excerpt, "done")
        self.assertNotIn("private_prompt", json.dumps(result.samples[0].metadata))
        self.assertEqual(result.dataset_version, "verifiers:0.3.1")

    def test_multi_agent_episode_ignores_non_trainable_user_simulator(self) -> None:
        rows = [
            episode(
                "ep-1",
                [
                    trace("assistant", 0.75, agent="assistant"),
                    trace("user", 0.0, agent="user", trainable=False),
                ],
            )
        ]
        with tempfile.TemporaryDirectory() as tmp:
            result = import_quality(self.write(Path(tmp), rows), "user-sim", pass_threshold=0.5)
        self.assertEqual(result.samples[0].score, 0.75)
        self.assertTrue(result.samples[0].passed)
        self.assertEqual(result.samples[0].metadata["agents"], ["assistant"])

    def test_named_metric_can_be_the_score_source(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            path = self.write(Path(tmp), [episode("ep-1", [trace("tr-1", 0.25)])])
            result = import_quality(path, "strict", pass_threshold=0.2, primary_metric="strict")
        self.assertEqual(result.samples[0].score, 0.25)
        self.assertEqual(result.samples[0].metadata["score_source"], "metric:strict")

    def test_fingerprint_uses_task_identity_not_line_order(self) -> None:
        rows = [
            episode("ep-1", [trace("tr-1", 1.0)], "a"),
            episode("ep-2", [trace("tr-2", 0.0)], "b"),
        ]
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            a = import_quality(self.write(root, rows), "t").dataset_fingerprint
            b = import_quality(self.write(root, list(reversed(rows))), "t").dataset_fingerprint
        self.assertEqual(a, b)

    def test_duplicate_episode_ids_are_refused(self) -> None:
        row = episode("same", [trace("tr-1", 1.0)])
        with tempfile.TemporaryDirectory() as tmp:
            path = self.write(Path(tmp), [row, row])
            with self.assertRaises(VerifiersImportError):
                read_episodes(path)

    def test_directory_resolves_the_upstream_traces_filename(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            self.write(root, [episode("ep-1", [trace("tr-1", 1.0)])])
            self.assertEqual(len(read_episodes(root)), 1)


if __name__ == "__main__":
    unittest.main()