#!/usr/bin/env python3
"""Unit tests for measured-only aggregation and rendering helpers."""

from __future__ import annotations

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

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

from aggregate_metrics import evalplus_metrics, training_metrics


class AggregationTests(unittest.TestCase):
    def test_evalplus_first_sample_is_pass_at_one(self):
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "eval.json"
            path.write_text(json.dumps({"eval": {
                "HumanEval/0": [{"base_status": "pass", "plus_status": "fail"}],
                "HumanEval/1": [{"base_status": "pass", "plus_status": "pass"}],
            }}))
            result = evalplus_metrics(path, "test")
        self.assertEqual(result["tasks"], 2)
        self.assertEqual(result["pass_at_1"], {"base": 1.0, "plus": 0.5})

    def test_evalplus_rejects_summary_without_task_statuses(self):
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "eval.json"
            path.write_text(json.dumps({"pass_at_k": {"base": {"1": 0.5}}}))
            with self.assertRaises(ValueError):
                evalplus_metrics(path, "HumanEval")

    def test_evalplus_rejects_partial_official_run(self):
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "eval.json"
            path.write_text(json.dumps({"eval": {
                "HumanEval/0": [{"base_status": "pass", "plus_status": "pass"}],
            }}))
            with self.assertRaisesRegex(ValueError, "incomplete HumanEval"):
                evalplus_metrics(path, "HumanEval")

    def test_training_excludes_warmup_and_nonfinite_values(self):
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "metrics.jsonl"
            path.write_text("\n".join((
                json.dumps({"step": 10, "tokens_per_second": 100}),
                json.dumps({"step": 20, "tok_s": 200}),
                json.dumps({"step": 30, "tokens_per_second": "unknown"}),
            )))
            result = training_metrics([path], warmup_steps=10)
        self.assertEqual(result["samples"], 1)
        self.assertEqual(result["mean_tokens_per_second"], 200)


class ModelCardTests(unittest.TestCase):
    def test_missing_metrics_are_not_measured(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            report = root / "report.json"
            output = root / "README.md"
            report.write_text(json.dumps({"schema_version": 1, "benchmarks": {}, "training": None, "model": None}))
            subprocess.run([
                sys.executable, str(HERE / "generate_model_card.py"), "--report", str(report),
                "--out", str(output), "--model-name", "Test Model",
            ], check=True, capture_output=True, text=True)
            text = output.read_text()
        self.assertIn("license: apache-2.0", text)
        self.assertIn("not measured", text)
        self.assertNotIn("0.00%", text)

    def test_charts_render_with_no_measurements(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            report = root / "report.json"
            report.write_text(json.dumps({"benchmarks": {}, "training": None}))
            outputs = []
            for script in ("chart_training.py", "chart_benchmarks.py"):
                output = root / f"{script}.png"
                subprocess.run([
                    sys.executable, str(HERE / script), "--report", str(report), "--out", str(output),
                ], check=True, capture_output=True, text=True)
                outputs.append(output)
            self.assertTrue(all(path.stat().st_size > 0 for path in outputs))


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