from __future__ import annotations

from pathlib import Path

from experiments.scripts.output_test_bench import (
    CaseSpec,
    build_progress_snapshot,
    build_benchmark_prompt,
    discover_cases,
    load_existing_results,
    merge_artifacts,
    scan_output_directory,
)


def test_discover_cases_extracts_case_metadata(tmp_path: Path) -> None:
    dataset_root = tmp_path / "output_test"
    case_dir = dataset_root / "simple" / "0000" / "00000797"
    case_dir.mkdir(parents=True)
    (case_dir / "description.txt").write_text("Create a part", encoding="utf-8")
    (case_dir / "00000797.png").write_bytes(b"png")
    (case_dir / "00000797_cip.py").write_text("# cip", encoding="utf-8")
    (case_dir / "00000797_cq.py").write_text("# cq", encoding="utf-8")
    (case_dir / "00000797.step").write_text("step", encoding="utf-8")

    cases = discover_cases(dataset_root)

    assert len(cases) == 1
    case = cases[0]
    assert case.case_id == "simple/0000/00000797"
    assert case.case_slug == "simple__0000__00000797"
    assert Path(case.description_path).name == "description.txt"
    assert Path(case.image_path or "").name == "00000797.png"
    assert Path(case.reference_cip_path or "").name == "00000797_cip.py"
    assert Path(case.reference_cq_path or "").name == "00000797_cq.py"
    assert Path(case.reference_step_path or "").name == "00000797.step"


def test_discover_cases_works_when_dataset_root_is_already_simple_dir(
    tmp_path: Path,
) -> None:
    dataset_root = tmp_path / "simple"
    case_dir = dataset_root / "0000" / "00000797"
    case_dir.mkdir(parents=True)
    (case_dir / "description.txt").write_text("Create a part", encoding="utf-8")
    (case_dir / "00000797.png").write_bytes(b"png")

    cases = discover_cases(dataset_root)

    assert len(cases) == 1
    case = cases[0]
    assert case.case_id == "0000/00000797"
    assert case.bucket_id == "0000"
    assert case.sample_id == "00000797"


def test_build_benchmark_prompt_requires_full_automation(tmp_path: Path) -> None:
    case = CaseSpec(
        case_id="simple/0000/00000797",
        case_slug="simple__0000__00000797",
        bucket_id="0000",
        sample_id="00000797",
        case_dir=str(tmp_path),
        description_path=str(tmp_path / "description.txt"),
        image_path=str(tmp_path / "00000797.png"),
        reference_cip_path=None,
        reference_cq_path=None,
        reference_step_path=None,
    )

    prompt = build_benchmark_prompt(
        case=case,
        description_text="Build a flange from the reference image.",
        requested_output_dir=Path(
            "/repo/workspace/experiments/output_test/r1/simple/0000/00000797"
        ),
        execution_root=Path("/repo/workspace"),
    )

    assert "full user confirmation and authorization" in prompt
    assert "Do not stop to ask for confirmation" in prompt
    assert (
        "target_output_dir: ./experiments/output_test/r1/simple/0000/00000797" in prompt
    )
    assert "<|code_file|>path/to/model.py</|code_file|>" in prompt
    assert "Build a flange from the reference image." in prompt


def test_merge_artifacts_falls_back_to_directory_scan(tmp_path: Path) -> None:
    output_dir = tmp_path / "workspace" / "case"
    output_dir.mkdir(parents=True)
    (output_dir / "model.py").write_text("print('ok')", encoding="utf-8")
    (output_dir / "part.stl").write_text("solid", encoding="utf-8")
    (output_dir / "part.step").write_text("step", encoding="utf-8")

    scanned = scan_output_directory(output_dir)
    merged = merge_artifacts(
        tagged_artifacts={"code_path": None, "output_paths": []},
        scanned_outputs=scanned,
        requested_output_dir=output_dir,
    )

    assert merged["artifact_source"] == "dir_scan"
    assert merged["code_path"] == output_dir / "model.py"
    assert merged["stl_path"] == output_dir / "part.stl"
    assert merged["step_path"] == output_dir / "part.step"


def test_load_existing_results_keeps_latest_per_case(tmp_path: Path) -> None:
    case_results = tmp_path / "case_results.jsonl"
    case_results.write_text(
        "\n".join(
            [
                '{"case_id":"simple/0000/0001","status":"failed"}',
                '{"case_id":"simple/0000/0002","status":"success"}',
                '{"case_id":"simple/0000/0001","status":"success"}',
            ]
        )
        + "\n",
        encoding="utf-8",
    )

    loaded = load_existing_results(case_results)

    assert sorted(loaded) == ["simple/0000/0001", "simple/0000/0002"]
    assert loaded["simple/0000/0001"]["status"] == "success"


def test_build_progress_snapshot_counts_skipped_running_and_pending(
    tmp_path: Path,
) -> None:
    results = [
        {"case_id": "simple/0000/0001", "status": "success"},
        {"case_id": "simple/0000/0002", "status": "failed"},
    ]

    snapshot = build_progress_snapshot(
        run_id="r1",
        dataset_root=tmp_path,
        results=results,
        total_cases=5,
        skipped_existing=2,
        running_case_ids=["simple/0000/0003"],
        pending_case_ids=["simple/0000/0004", "simple/0000/0005"],
        status="running",
    )

    assert snapshot["run_id"] == "r1"
    assert snapshot["total_cases"] == 5
    assert snapshot["completed_cases"] == 2
    assert snapshot["skipped_existing"] == 2
    assert snapshot["running_cases"] == 1
    assert snapshot["pending_cases"] == 2
    assert snapshot["success"] == 1
    assert snapshot["failed"] == 1