from __future__ import annotations

from agent_preflight.models import RunResult
from agent_preflight.runner import run_multi, run_single


class TestRunMulti:
    def test_preflight_models_are_merged_with_preflight_kwargs(self, monkeypatch, tmp_path):
        captured = {}

        def fake_preflight(**kwargs):
            captured.update(kwargs)
            return []

        def fake_run_container(**kwargs):
            return RunResult(
                alias=kwargs["alias"],
                model_id=kwargs["model_id"],
                budget_usd=kwargs["budget_usd"],
            )

        monkeypatch.setattr("agent_preflight.runner.preflight", fake_preflight)
        monkeypatch.setattr("agent_preflight.runner.run_container", fake_run_container)

        results = run_multi(
            models=[
                ("z-ai/glm-5", "writer"),
                ("google/gemini-2.5-flash", "planner"),
            ],
            image_name="test-image:latest",
            admin_key="admin",
            or_key="or",
            budget_usd=0.05,
            runs_dir=tmp_path,
            parallel=False,
            preflight_kwargs={"models": ["google/gemini-2.5-flash", "anthropic/claude-sonnet-4"]},
        )

        assert [r.alias for r in results] == ["writer", "planner"]
        assert captured["models"] == [
            "z-ai/glm-5",
            "google/gemini-2.5-flash",
            "anthropic/claude-sonnet-4",
        ]


class TestRunSingle:
    def test_extra_models_are_forwarded_to_preflight(self, monkeypatch, tmp_path):
        captured = {}

        def fake_run_multi(**kwargs):
            captured.update(kwargs)
            return [
                RunResult(
                    alias="glm-5",
                    model_id="z-ai/glm-5",
                    budget_usd=kwargs["budget_usd"],
                )
            ]

        monkeypatch.setattr("agent_preflight.runner.run_multi", fake_run_multi)

        result = run_single(
            model_id="z-ai/glm-5",
            image_name="test-image:latest",
            admin_key="admin",
            or_key="or",
            budget_usd=0.05,
            runs_dir=tmp_path,
            extra_models=["google/gemini-2.5-flash"],
            preflight_kwargs={"models": ["google/gemini-2.5-flash", "anthropic/claude-sonnet-4"]},
        )

        assert result.model_id == "z-ai/glm-5"
        assert captured["models"] == [("z-ai/glm-5", "glm-5")]
        assert captured["parallel"] is False
        assert captured["preflight_kwargs"] == {
            "models": ["google/gemini-2.5-flash", "anthropic/claude-sonnet-4"],
        }
