"""Tests for eval export utilities."""

from __future__ import annotations

import json

from openjarvis.evals.core.export import (
    _compute_efficiency,
    _compute_normalized,
    export_artifacts_manifest,
    export_jsonl,
    export_summary_json,
)
from openjarvis.evals.core.trace import QueryTrace, TurnTrace


def _make_traces(n=3):
    traces = []
    for i in range(n):
        traces.append(
            QueryTrace(
                query_id=f"q{i:04d}",
                workload_type="test",
                query_text=f"Question {i}",
                response_text=f"Answer {i}",
                turns=[
                    TurnTrace(
                        turn_index=0,
                        input_tokens=100 + i * 10,
                        output_tokens=50 + i * 5,
                        wall_clock_s=1.0 + i * 0.5,
                        gpu_energy_joules=5.0 + i,
                        cost_usd=0.01,
                    ),
                ],
                total_wall_clock_s=1.0 + i * 0.5,
                completed=True,
                is_resolved=i % 2 == 0,
            )
        )
    return traces


class TestExportJsonl:
    def test_basic_export(self, tmp_path):
        traces = _make_traces()
        path = tmp_path / "traces.jsonl"
        result = export_jsonl(traces, path)
        assert result == path
        assert path.exists()

        lines = path.read_text().strip().split("\n")
        assert len(lines) == 3
        for line in lines:
            d = json.loads(line)
            assert "query_id" in d
            assert "turns" in d

    def test_empty_traces(self, tmp_path):
        path = tmp_path / "empty.jsonl"
        export_jsonl([], path)
        assert path.read_text() == ""

    def test_creates_parent_dirs(self, tmp_path):
        path = tmp_path / "a" / "b" / "c" / "traces.jsonl"
        export_jsonl(_make_traces(1), path)
        assert path.exists()


class TestExportSummaryJson:
    def test_basic_summary(self, tmp_path):
        traces = _make_traces()
        path = tmp_path / "summary.json"
        result = export_summary_json(traces, {"model": "test"}, path)
        assert result == path
        assert path.exists()

        summary = json.loads(path.read_text())
        assert summary["totals"]["queries"] == 3
        assert summary["totals"]["completed"] == 3
        assert summary["totals"]["resolved"] == 2
        assert summary["totals"]["input_tokens"] > 0
        assert summary["totals"]["output_tokens"] > 0
        assert summary["config"]["model"] == "test"
        assert "statistics" in summary

    def test_statistics_keys(self, tmp_path):
        traces = _make_traces()
        path = tmp_path / "summary.json"
        export_summary_json(traces, {}, path)
        summary = json.loads(path.read_text())
        stats = summary["statistics"]
        expected_stat_keys = {
            "wall_clock_s",
            "gpu_energy_joules",
            "cpu_energy_joules",
            "gpu_power_watts",
            "cpu_power_watts",
            "input_tokens",
            "output_tokens",
            "total_tokens",
            "throughput_tokens_per_sec",
            "energy_per_token_joules",
            "cost_usd",
            "turns",
            "tool_calls",
            "mbu_avg_pct",
        }
        assert set(stats.keys()) == expected_stat_keys

    def test_agg_stats_fields(self, tmp_path):
        traces = _make_traces()
        path = tmp_path / "summary.json"
        export_summary_json(traces, {}, path)
        summary = json.loads(path.read_text())
        wc_stats = summary["statistics"]["wall_clock_s"]
        assert "avg" in wc_stats
        assert "median" in wc_stats
        assert "min" in wc_stats
        assert "max" in wc_stats
        assert "std" in wc_stats

    def test_empty_traces(self, tmp_path):
        path = tmp_path / "summary.json"
        export_summary_json([], {}, path)
        summary = json.loads(path.read_text())
        assert summary["totals"]["queries"] == 0


class TestHarnessErrorExclusion:
    """Harness errors must never count as model misses in resolve-rate."""

    def _traces(self):
        return [
            QueryTrace(
                query_id="q0000",
                workload_type="agentic",
                completed=True,
                is_resolved=True,
                turns=[TurnTrace(turn_index=0, input_tokens=10, output_tokens=5)],
            ),
            QueryTrace(
                query_id="q0001",
                workload_type="agentic",
                completed=True,
                is_resolved=False,  # genuine model miss: stays in denominator
                turns=[TurnTrace(turn_index=0, input_tokens=10, output_tokens=5)],
            ),
            QueryTrace(
                query_id="q0002",
                workload_type="agentic",
                completed=False,
                is_resolved=False,  # stamped by run_tests despite zero contact
                error="zero_model_requests: the agent never contacted the model",
                error_kind="harness_error",
            ),
        ]

    def test_summary_excludes_harness_errors_from_accuracy(self, tmp_path):
        path = tmp_path / "summary.json"
        export_summary_json(self._traces(), {"model": "m"}, path)
        summary = json.loads(path.read_text())

        totals = summary["totals"]
        assert totals["harness_errors"] == 1
        assert totals["resolved"] == 1
        assert totals["unresolved"] == 1  # NOT 2: harness error excluded
        assert totals["accuracy"] == 0.5  # NOT 1/3
        # Flat table_gen metrics exclude the harness error too.
        assert summary["metrics"]["accuracy"]["n"] == 2
        assert summary["metrics"]["accuracy"]["mean"] == 0.5
        # Details surfaced for diagnosis.
        details = summary["harness_error_details"]
        assert details[0]["query_id"] == "q0002"
        assert "zero_model_requests" in details[0]["error"]

    def test_efficiency_excludes_harness_errors(self):
        result = _compute_efficiency(self._traces(), None, None)
        assert result["accuracy"] == 0.5

    def test_no_harness_errors_key_absent(self, tmp_path):
        path = tmp_path / "summary.json"
        export_summary_json(_make_traces(), {}, path)
        summary = json.loads(path.read_text())
        assert summary["totals"]["harness_errors"] == 0
        assert "harness_error_details" not in summary


class TestTraceErrorFieldRoundTrip:
    def test_round_trip_preserves_error_fields(self):
        trace = QueryTrace(
            query_id="q0",
            workload_type="agentic",
            error="zero_model_requests: ...",
            error_kind="harness_error",
        )
        restored = QueryTrace.from_dict(trace.to_dict())
        assert restored.error == trace.error
        assert restored.error_kind == "harness_error"

    def test_old_trace_dicts_still_load(self):
        """Backward compat: traces.jsonl written before the schema change."""
        restored = QueryTrace.from_dict({"query_id": "q0", "workload_type": "agentic"})
        assert restored.error is None
        assert restored.error_kind is None


class TestExportArtifactsManifest:
    def test_no_artifacts_dir(self, tmp_path):
        result = export_artifacts_manifest(tmp_path)
        assert result is None

    def test_with_artifacts(self, tmp_path):
        art_dir = tmp_path / "artifacts"
        q_dir = art_dir / "q0001_test"
        q_dir.mkdir(parents=True)
        (q_dir / "response.txt").write_text("hello")
        (q_dir / "metadata.json").write_text("{}")

        result = export_artifacts_manifest(tmp_path)
        assert result is not None
        assert result.exists()

        manifest = json.loads(result.read_text())
        assert len(manifest) == 1
        assert manifest[0]["query_dir"] == "q0001_test"
        assert len(manifest[0]["files"]) == 2


class TestComputeEfficiency:
    def test_with_resolved_traces(self):
        traces = _make_traces()
        result = _compute_efficiency(traces, 15.0, 3.0)
        # 2 resolved out of 3 scored (is_resolved=True for i=0,2; False for i=1)
        assert result["accuracy"] == 2 / 3
        assert result["total_gpu_energy_joules"] == 15.0
        assert result["total_cpu_energy_joules"] == 3.0
        assert result["ipj"] is not None
        assert result["ipw"] is None  # no gpu power data on traces

    def test_no_scored_traces(self):
        traces = [
            QueryTrace(
                query_id="q0",
                workload_type="test",
                completed=True,
                is_resolved=None,
            ),
        ]
        result = _compute_efficiency(traces, 5.0, 1.0)
        assert result["accuracy"] is None
        assert result["ipj"] is None

    def test_no_energy(self):
        traces = _make_traces(1)
        result = _compute_efficiency(traces, None, None)
        assert result["total_gpu_energy_joules"] is None
        assert result["ipj"] is None

    def test_with_gpu_power(self):
        traces = [
            QueryTrace(
                query_id="q0",
                workload_type="test",
                completed=True,
                is_resolved=True,
                query_gpu_power_avg_watts=100.0,
            ),
        ]
        result = _compute_efficiency(traces, 50.0, None)
        assert result["accuracy"] == 1.0
        assert result["avg_gpu_power_watts"] == 100.0
        assert result["ipw"] == 1.0 / 100.0


class TestComputeNormalized:
    def test_too_few_traces(self):
        traces = _make_traces(3)
        assert _compute_normalized(traces) is None

    def test_with_enough_traces(self):
        traces = _make_traces(10)
        result = _compute_normalized(traces)
        assert result is not None
        norm_stats = result["normalized_statistics"]
        norm_eff = result["normalized_efficiency"]
        assert "_description" in norm_stats
        assert "_outliers_removed" in norm_stats
        assert norm_stats["_outliers_removed"] == 2  # 1 from each end
        assert "wall_clock_s" in norm_stats
        assert "mbu_avg_pct" in norm_stats
        assert "accuracy" in norm_eff

    def test_large_set_trims_more(self):
        traces = _make_traces(40)
        result = _compute_normalized(traces)
        assert result is not None
        # floor(40*0.05)=2 from each end
        assert result["normalized_statistics"]["_outliers_removed"] == 4


class TestExportSummaryJsonNewSections:
    def test_totals_accuracy(self, tmp_path):
        traces = _make_traces()
        path = tmp_path / "summary.json"
        export_summary_json(traces, {}, path)
        summary = json.loads(path.read_text())
        assert "accuracy" in summary["totals"]
        assert summary["totals"]["accuracy"] == 2 / 3

    def test_efficiency_section(self, tmp_path):
        traces = _make_traces()
        path = tmp_path / "summary.json"
        export_summary_json(traces, {}, path)
        summary = json.loads(path.read_text())
        assert "efficiency" in summary
        eff = summary["efficiency"]
        assert "accuracy" in eff
        assert "total_gpu_energy_joules" in eff
        assert "ipj" in eff
        assert "ipw" in eff

    def test_mbu_avg_pct_in_statistics(self, tmp_path):
        traces = _make_traces()
        traces[0].query_mbu_avg_pct = 45.0
        traces[1].query_mbu_avg_pct = 55.0
        path = tmp_path / "summary.json"
        export_summary_json(traces, {}, path)
        summary = json.loads(path.read_text())
        mbu_stats = summary["statistics"]["mbu_avg_pct"]
        assert mbu_stats["avg"] == 50.0

    def test_normalized_present_for_enough_traces(self, tmp_path):
        traces = _make_traces(10)
        path = tmp_path / "summary.json"
        export_summary_json(traces, {}, path)
        summary = json.loads(path.read_text())
        assert "normalized_statistics" in summary
        assert "normalized_efficiency" in summary

    def test_normalized_absent_for_few_traces(self, tmp_path):
        traces = _make_traces(3)
        path = tmp_path / "summary.json"
        export_summary_json(traces, {}, path)
        summary = json.loads(path.read_text())
        assert "normalized_statistics" not in summary
        assert "normalized_efficiency" not in summary

    def test_bench_telemetry(self, tmp_path):
        traces = _make_traces()
        path = tmp_path / "summary.json"
        bench_energy = {"total_energy_joules": 100.0, "avg_power_watts": 50.0}
        export_summary_json(traces, {}, path, bench_energy=bench_energy)
        summary = json.loads(path.read_text())
        assert "bench_telemetry" in summary
        assert summary["bench_telemetry"]["total_energy_joules"] == 100.0

    def test_no_bench_telemetry_when_none(self, tmp_path):
        traces = _make_traces()
        path = tmp_path / "summary.json"
        export_summary_json(traces, {}, path)
        summary = json.loads(path.read_text())
        assert "bench_telemetry" not in summary


class TestQueryTraceMbuRoundTrip:
    def test_mbu_serialization(self):
        trace = QueryTrace(
            query_id="q0",
            workload_type="test",
            query_mbu_avg_pct=42.5,
            query_mbu_max_pct=87.3,
        )
        d = trace.to_dict()
        assert d["query_mbu_avg_pct"] == 42.5
        assert d["query_mbu_max_pct"] == 87.3

        restored = QueryTrace.from_dict(d)
        assert restored.query_mbu_avg_pct == 42.5
        assert restored.query_mbu_max_pct == 87.3

    def test_mbu_none_by_default(self):
        trace = QueryTrace(query_id="q0", workload_type="test")
        d = trace.to_dict()
        assert d["query_mbu_avg_pct"] is None
        assert d["query_mbu_max_pct"] is None

        restored = QueryTrace.from_dict(d)
        assert restored.query_mbu_avg_pct is None
        assert restored.query_mbu_max_pct is None

    def test_mbu_in_hf_dataset_rows(self):
        trace = QueryTrace(
            query_id="q0",
            workload_type="test",
            query_mbu_avg_pct=33.3,
            query_mbu_max_pct=66.6,
        )
        # Test that to_dict includes the fields (hf_dataset uses to_dict internally)
        d = trace.to_dict()
        assert "query_mbu_avg_pct" in d
        assert "query_mbu_max_pct" in d


class TestActionEnergyBreakdown:
    def test_turn_trace_round_trip(self):
        breakdown = [
            {
                "action_type": "lm_inference",
                "duration_s": 1.5,
                "gpu_energy_joules": 10.0,
                "cpu_energy_joules": 0.5,
            },
            {
                "action_type": "tool_call:calculator",
                "duration_s": 0.2,
                "gpu_energy_joules": 0.1,
                "cpu_energy_joules": 0.01,
            },
        ]
        turn = TurnTrace(
            turn_index=0,
            action_energy_breakdown=breakdown,
        )
        d = turn.to_dict()
        assert d["action_energy_breakdown"] is not None
        assert len(d["action_energy_breakdown"]) == 2
        assert d["action_energy_breakdown"][0]["action_type"] == "lm_inference"

        restored = TurnTrace.from_dict(d)
        assert restored.action_energy_breakdown is not None
        assert len(restored.action_energy_breakdown) == 2
        action_type = restored.action_energy_breakdown[1]["action_type"]
        assert action_type == "tool_call:calculator"

    def test_turn_trace_none_by_default(self):
        turn = TurnTrace(turn_index=0)
        d = turn.to_dict()
        assert d["action_energy_breakdown"] is None

        restored = TurnTrace.from_dict(d)
        assert restored.action_energy_breakdown is None

    def test_action_energy_summary_in_export(self, tmp_path):
        traces = []
        for i in range(2):
            traces.append(
                QueryTrace(
                    query_id=f"q{i:04d}",
                    workload_type="test",
                    turns=[
                        TurnTrace(
                            turn_index=0,
                            input_tokens=100,
                            output_tokens=50,
                            wall_clock_s=2.0,
                            gpu_energy_joules=5.0,
                            action_energy_breakdown=[
                                {
                                    "action_type": "lm_inference",
                                    "duration_s": 1.5,
                                    "gpu_energy_joules": 4.0,
                                    "cpu_energy_joules": 0.3,
                                },
                                {
                                    "action_type": "tool_call:search",
                                    "duration_s": 0.5,
                                    "gpu_energy_joules": 1.0,
                                    "cpu_energy_joules": 0.1,
                                },
                            ],
                        ),
                    ],
                    total_wall_clock_s=2.0,
                    completed=True,
                )
            )
        path = tmp_path / "summary.json"
        export_summary_json(traces, {}, path)
        summary = json.loads(path.read_text())
        assert "action_energy_summary" in summary
        aes = summary["action_energy_summary"]
        assert "lm_inference" in aes
        assert aes["lm_inference"]["count"] == 2
        assert aes["lm_inference"]["total_gpu_energy_joules"] == 8.0
        assert "tool_call:search" in aes
        assert aes["tool_call:search"]["count"] == 2

    def test_no_action_energy_summary_when_empty(self, tmp_path):
        traces = _make_traces()
        path = tmp_path / "summary.json"
        export_summary_json(traces, {}, path)
        summary = json.loads(path.read_text())
        assert "action_energy_summary" not in summary


class TestHardwareInfo:
    def test_hardware_info_in_summary(self, tmp_path):
        traces = _make_traces()
        path = tmp_path / "summary.json"
        export_summary_json(traces, {}, path)
        summary = json.loads(path.read_text())
        assert "hardware_info" in summary
        hw = summary["hardware_info"]
        # Should have at least platform and cpu_count
        assert "platform" in hw
        assert "cpu_count" in hw
        assert hw["cpu_count"] > 0
