"""Tests for config_loader.load_all_configs with _common/ + preset structure."""

from __future__ import annotations
from pathlib import Path
import pytest
from src.conf.config_loader import load_all_configs


def _make_minimal_preset_layout(root: Path) -> None:
    """Create _common/ + base files matching the production layout."""
    common = root / "_common"
    common.mkdir()
    (common / "training.gin").write_text("TrainingConfig.tau_init = 0.07\n")
    (common / "hardware_default.gin").write_text("HardwareConfig.batch_size = 8\n")
    (common / "hardware_no_gradckpt.gin").write_text(
        "HardwareConfig.batch_size = 8\n"
        "HardwareConfig.gradient_checkpointing = False\n",
    )
    (common / "tracking.gin").write_text("TrackingConfig.use_tb = True\n")


def _make_preset(root: Path, name: str, models_body: str, output_dir: str) -> None:
    """Create a preset directory with required pipeline.gin + models.gin."""
    preset = root / name
    preset.mkdir()
    (preset / "pipeline.gin").write_text(
        f"PipelineConfig.output_dir = '{output_dir}'\n",
    )
    (preset / "models.gin").write_text(models_body)


def test_load_returns_six_keys(tmp_path: Path) -> None:
    """load_all_configs returns the 6 expected keys."""
    _make_minimal_preset_layout(tmp_path)
    _make_preset(
        tmp_path, "test_dinov3",
        models_body="ModelsCommonConfig.backbone = 'dinov3'\n",
        output_dir="out/test_dinov3",
    )

    cfgs = load_all_configs("test_dinov3", str(tmp_path))
    assert set(cfgs.keys()) == {
        "pipeline", "hardware", "training", "tracking", "models_common", "models",
    }
    assert cfgs["pipeline"].output_dir == "out/test_dinov3"
    assert cfgs["models_common"].backbone == "dinov3"


def test_sofia_uses_no_gradckpt_hardware(tmp_path: Path) -> None:
    """Sofia backbones get hardware_no_gradckpt.gin (gradient_checkpointing=False)."""
    _make_minimal_preset_layout(tmp_path)
    _make_preset(
        tmp_path, "test_sofia",
        models_body="ModelsCommonConfig.backbone = 'sofia_v71'\n",
        output_dir="out/test_sofia",
    )

    cfgs = load_all_configs("test_sofia", str(tmp_path))
    assert cfgs["hardware"].gradient_checkpointing is False


def test_dinov3_uses_default_hardware(tmp_path: Path) -> None:
    """DINOv3/StripNet backbones get hardware_default.gin (gradient_checkpointing=True default)."""
    _make_minimal_preset_layout(tmp_path)
    # Override gradient_checkpointing=True in default profile to test routing.
    (tmp_path / "_common" / "hardware_default.gin").write_text(
        "HardwareConfig.batch_size = 8\n"
        "HardwareConfig.gradient_checkpointing = True\n",
    )
    _make_preset(
        tmp_path, "test_dinov3_hw",
        models_body="ModelsCommonConfig.backbone = 'dinov3'\n",
        output_dir="out/test_dinov3_hw",
    )
    cfgs = load_all_configs("test_dinov3_hw", str(tmp_path))
    assert cfgs["hardware"].gradient_checkpointing is True


def test_local_override_wins(tmp_path: Path) -> None:
    """Local training.gin in preset dir overrides _common/training.gin."""
    _make_minimal_preset_layout(tmp_path)
    _make_preset(
        tmp_path, "test_override",
        models_body="ModelsCommonConfig.backbone = 'dinov3'\n",
        output_dir="out/test_override",
    )
    # Drop a local training.gin that overrides tau_init.
    (tmp_path / "test_override" / "training.gin").write_text(
        "TrainingConfig.tau_init = 0.123\n",
    )

    cfgs = load_all_configs("test_override", str(tmp_path))
    assert cfgs["training"].tau_init == 0.123


def test_unknown_backbone_raises(tmp_path: Path) -> None:
    """Invalid backbone string in models.gin → ValueError."""
    _make_minimal_preset_layout(tmp_path)
    _make_preset(
        tmp_path, "test_bad",
        models_body="ModelsCommonConfig.backbone = 'mistral_42b'\n",
        output_dir="out/bad",
    )

    with pytest.raises(ValueError, match="Unknown backbone"):
        load_all_configs("test_bad", str(tmp_path))


def test_missing_preset_raises(tmp_path: Path) -> None:
    """Non-existent preset directory → FileNotFoundError."""
    _make_minimal_preset_layout(tmp_path)
    with pytest.raises(FileNotFoundError, match="Preset directory not found"):
        load_all_configs("nonexistent", str(tmp_path))


def test_missing_required_file_raises(tmp_path: Path) -> None:
    """Preset missing pipeline.gin → FileNotFoundError."""
    _make_minimal_preset_layout(tmp_path)
    bad_preset = tmp_path / "test_bad_preset"
    bad_preset.mkdir()
    (bad_preset / "models.gin").write_text("ModelsCommonConfig.backbone = 'dinov3'\n")

    with pytest.raises(FileNotFoundError, match="missing required file 'pipeline.gin'"):
        load_all_configs("test_bad_preset", str(tmp_path))