from __future__ import annotations

import sys
from pathlib import Path
from types import ModuleType, SimpleNamespace
from typing import Any

from yolo_tui.config import MlflowConfig, TrainingConfig
from yolo_tui.trainer import TrainingEvent, TrainingRunner


def test_runner_wires_yolo_callbacks_and_returns_output(
    monkeypatch: Any, tmp_path: Path
) -> None:
    settings_updates: list[dict[str, bool]] = []
    constructed: list[tuple[str, str]] = []
    train_arguments: list[dict[str, Any]] = []

    class FakeSettings:
        def update(self, values: dict[str, bool]) -> None:
            settings_updates.append(values)

    class FakeYOLO:
        def __init__(self, model: str, task: str) -> None:
            constructed.append((model, task))
            self.callbacks: dict[str, Any] = {}
            self.trainer = SimpleNamespace(
                args=SimpleNamespace(epochs=2),
                epoch=0,
                metrics={},
                stop=False,
                save_dir=tmp_path / "run",
            )

        def add_callback(self, name: str, callback: Any) -> None:
            self.callbacks[name] = callback

        def train(self, **kwargs: Any) -> None:
            train_arguments.append(kwargs)
            self.callbacks["on_train_start"](self.trainer)
            for epoch in range(2):
                self.trainer.epoch = epoch
                self.trainer.metrics = {"metrics/mAP50": 0.5 + epoch / 10}
                self.callbacks["on_train_epoch_end"](self.trainer)
            self.callbacks["on_train_end"](self.trainer)

    fake_ultralytics = ModuleType("ultralytics")
    fake_ultralytics.YOLO = FakeYOLO  # type: ignore[attr-defined]
    fake_ultralytics.settings = FakeSettings()  # type: ignore[attr-defined]
    monkeypatch.setitem(sys.modules, "ultralytics", fake_ultralytics)

    config = TrainingConfig(
        dataset="dataset.yaml",
        model="model.pt",
        task="pose",
        epochs=2,
        mlflow=MlflowConfig(enabled=False),
    )
    events: list[TrainingEvent] = []

    output = TrainingRunner().train(config, events.append)

    assert output == tmp_path / "run"
    assert constructed == [("model.pt", "pose")]
    assert settings_updates == [{"mlflow": False}]
    assert train_arguments[0]["data"] == "dataset.yaml"
    assert train_arguments[0]["verbose"] is False
    assert [event.kind for event in events] == ["info", "started", "epoch", "epoch", "success"]