import csv
import json
from io import StringIO
from pathlib import Path

import pytest

from dlb.matrix import MatrixTask
from dlb.timing_summary import main, summarize_timing, write_timing_csv


PROJECT_ROOT = Path(__file__).parents[1]


def prepare_registry(root: Path) -> None:
    (root / "configs").mkdir(parents=True, exist_ok=True)
    (root / "configs/experiments.yaml").write_text(
        (PROJECT_ROOT / "configs/experiments.yaml").read_text(encoding="utf-8"),
        encoding="utf-8",
    )


def fake_task(root: Path, steps: int) -> MatrixTask:
    timing = root / f"results/timing/lm1b/fmlm/steps_{steps}/timing.json"
    return MatrixTask(
        task_id=f"fmlm-lm1b-steps-{steps}",
        category="few",
        model="fmlm",
        dataset="lm1b",
        steps=steps,
        sample_count=1024,
        seed=42,
        environment="dlb-flm",
        adapter="flm",
        source="flm",
        provenance="official",
        sample_dir=str(root / f"results/samples/lm1b/fmlm/steps_{steps}"),
        metrics_path=str(root / f"results/metrics/lm1b/fmlm/steps_{steps}/metrics.json"),
        timing_path=str(timing),
    )


def write_timing(path: Path, seconds: float) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        json.dumps(
            {
                "schema": "dlb-generation-timing-v1",
                "timing": {
                    "mode": "primary_latency",
                    "warmups": 5,
                    "repeats": 32,
                    "batch_size": 1,
                    "num_timed_samples": 32,
                    "seconds_per_sample": seconds,
                },
            }
        ),
        encoding="utf-8",
    )


def test_summarize_timing_reports_present_rows_and_missing_steps(tmp_path: Path) -> None:
    tasks = [fake_task(tmp_path, step) for step in (1, 2, 4)]
    write_timing(Path(tasks[0].timing_path), 0.125)
    write_timing(Path(tasks[2].timing_path), 0.5)

    report = summarize_timing(tasks)

    assert [row.as_csv_row() for row in report.rows] == [
        {
            "dataset": "lm1b",
            "model": "fmlm",
            "steps": "1",
            "seconds_per_sample": "0.125000",
        },
        {
            "dataset": "lm1b",
            "model": "fmlm",
            "steps": "4",
            "seconds_per_sample": "0.500000",
        },
    ]
    assert [(item.dataset, item.model, item.steps) for item in report.missing] == [
        ("lm1b", "fmlm", 2),
    ]


def test_write_timing_csv_uses_requested_column_order(tmp_path: Path) -> None:
    tasks = [fake_task(tmp_path, 2)]
    write_timing(Path(tasks[0].timing_path), 0.25)
    report = summarize_timing(tasks)
    output = StringIO()

    write_timing_csv(report.rows, output)

    parsed = list(csv.DictReader(StringIO(output.getvalue())))
    assert output.getvalue().splitlines()[0] == "dataset,model,steps,seconds_per_sample"
    assert parsed == [
        {
            "dataset": "lm1b",
            "model": "fmlm",
            "steps": "2",
            "seconds_per_sample": "0.250000",
        }
    ]


@pytest.mark.parametrize("seconds", [None, -0.1, "0.1"])
def test_summarize_timing_rejects_invalid_seconds_per_sample(
    tmp_path: Path, seconds: object
) -> None:
    task = fake_task(tmp_path, 1)
    path = Path(task.timing_path)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        json.dumps(
            {
                "schema": "dlb-generation-timing-v1",
                "timing": {
                    "mode": "primary_latency",
                    "warmups": 5,
                    "repeats": 32,
                    "batch_size": 1,
                    "num_timed_samples": 32,
                    "seconds_per_sample": seconds,
                },
            }
        ),
        encoding="utf-8",
    )

    report = summarize_timing([task])

    assert report.rows == ()
    assert len(report.invalid) == 1
    assert report.invalid[0].steps == 1


def test_cli_filters_matrix_and_reports_missing_counts(
    tmp_path: Path, capsys: pytest.CaptureFixture[str]
) -> None:
    prepare_registry(tmp_path)
    write_timing(
        tmp_path / "results/timing/lm1b/fmlm/steps_2/timing.json",
        0.25,
    )

    assert (
        main(
            [
                "--root",
                str(tmp_path),
                "--model",
                "fmlm",
                "--dataset",
                "lm1b",
            ]
        )
        == 0
    )

    captured = capsys.readouterr()
    assert captured.out.splitlines() == [
        "dataset,model,steps,seconds_per_sample",
        "lm1b,fmlm,2,0.250000",
    ]
    assert "expected=6 present=1 missing=5 invalid=0" in captured.err
    assert "missing,lm1b,fmlm,1," in captured.err