from datetime import datetime

from giant import config as gconfig


def _write_config(path, git_hash):
    path.write_text(
        f"""
[train]
epochs = 5

[model]
hidden_dim = 64

[meta]
git_hash = "{git_hash}"
"""
    )


def test_merge_cli_overrides_applies_file_then_cli(tmp_path, monkeypatch):
    monkeypatch.setattr(gconfig, "git_hash", lambda: "abc123")
    path = tmp_path / "config.toml"
    _write_config(path, "abc123")

    cfg = gconfig.merge_cli_overrides(
        gconfig.DEFAULT_CONFIG,
        path,
        train_overrides={},
        model_overrides={"hidden_dim": 128},
    )
    assert cfg["train"]["epochs"] == 5  # from file
    assert cfg["model"]["hidden_dim"] == 128  # CLI override wins over file


def test_merge_cli_overrides_warns_on_git_hash_mismatch(tmp_path, monkeypatch, capsys):
    monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
    path = tmp_path / "config.toml"
    _write_config(path, "old111")

    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}, {})

    assert cfg["train"]["epochs"] == 5  # does not fail, config still applied
    captured = capsys.readouterr()
    assert "warning" in captured.err
    assert "old111" in captured.err
    assert "current999" in captured.err


def test_merge_cli_overrides_no_warning_on_matching_git_hash(
    tmp_path, monkeypatch, capsys
):
    monkeypatch.setattr(gconfig, "git_hash", lambda: "same123")
    path = tmp_path / "config.toml"
    _write_config(path, "same123")

    gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}, {})
    assert capsys.readouterr().err == ""


def test_merge_cli_overrides_no_warning_when_git_hash_unknown(
    tmp_path, monkeypatch, capsys
):
    monkeypatch.setattr(gconfig, "git_hash", lambda: "unknown")
    path = tmp_path / "config.toml"
    _write_config(path, "abc123")

    gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}, {})
    assert capsys.readouterr().err == ""


def test_merge_cli_overrides_no_warning_when_meta_section_absent(
    tmp_path, monkeypatch, capsys
):
    monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
    path = tmp_path / "config.toml"
    path.write_text("[train]\nepochs = 5\n")

    gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}, {})
    assert capsys.readouterr().err == ""


def test_warn_if_checkpoint_config_mismatch_finds_sibling_toml(
    tmp_path, monkeypatch, capsys
):
    monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
    ckpt_path = tmp_path / "best.pt"
    ckpt_path.write_bytes(b"")  # contents irrelevant, only its directory is used
    _write_config(tmp_path / "config.toml", "old111")

    gconfig.warn_if_checkpoint_config_mismatch(ckpt_path)

    captured = capsys.readouterr()
    assert "warning" in captured.err
    assert "old111" in captured.err
    assert "current999" in captured.err


def test_warn_if_checkpoint_config_mismatch_no_warning_when_toml_absent(
    tmp_path, monkeypatch, capsys
):
    monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
    ckpt_path = tmp_path / "best.pt"
    ckpt_path.write_bytes(b"")

    gconfig.warn_if_checkpoint_config_mismatch(ckpt_path)
    assert capsys.readouterr().err == ""


def test_warn_if_checkpoint_config_mismatch_no_warning_when_hashes_match(
    tmp_path, monkeypatch, capsys
):
    monkeypatch.setattr(gconfig, "git_hash", lambda: "same123")
    ckpt_path = tmp_path / "best.pt"
    ckpt_path.write_bytes(b"")
    _write_config(tmp_path / "config.toml", "same123")

    gconfig.warn_if_checkpoint_config_mismatch(ckpt_path)
    assert capsys.readouterr().err == ""


def test_resolve_expert_dims_default_config_inherits_hidden_dim_and_n_blocks():
    # The unset sentinel (expert_hidden_dim/n_blocks == 0 in DEFAULT_CONFIG)
    # is exactly the bug fixed by resolve_expert_dims: it must not silently
    # fall back to some other hardcoded default, only to the monolith's own
    # hidden_dim/n_blocks, so --hidden-dim/--n-blocks reach the experts too.
    router_cfg = dict(gconfig.DEFAULT_CONFIG["model"]["router"])
    assert router_cfg["expert_hidden_dim"] == 0
    assert router_cfg["expert_n_blocks"] == 0

    hidden_dim, n_blocks = gconfig.resolve_expert_dims(router_cfg, 512, 6)
    assert (hidden_dim, n_blocks) == (512, 6)


def test_resolve_expert_dims_missing_keys_also_inherit():
    hidden_dim, n_blocks = gconfig.resolve_expert_dims({}, 512, 6)
    assert (hidden_dim, n_blocks) == (512, 6)


def test_default_config_gumbel_router_defaults_off():
    # Straight-through Gumbel-softmax combine weights (giant.model.network.
    # Router.combine_weights) must be opt-in — existing routed configs and
    # checkpoints should be unaffected unless gumbel is explicitly enabled.
    router_cfg = gconfig.DEFAULT_CONFIG["model"]["router"]
    assert router_cfg["gumbel"] is False
    assert router_cfg["gumbel_tau_start"] == 1.0
    assert router_cfg["gumbel_tau_end"] == 0.1


def test_resolve_expert_dims_explicit_override_wins():
    router_cfg = {"expert_hidden_dim": 128, "expert_n_blocks": 3}
    hidden_dim, n_blocks = gconfig.resolve_expert_dims(router_cfg, 512, 6)
    assert (hidden_dim, n_blocks) == (128, 3)


def test_resolve_expert_dims_partial_override_mixes_explicit_and_inherited():
    router_cfg = {"expert_hidden_dim": 128, "expert_n_blocks": 0}
    hidden_dim, n_blocks = gconfig.resolve_expert_dims(router_cfg, 512, 6)
    assert (hidden_dim, n_blocks) == (128, 6)


def _default_cfg(**overrides):
    train_overrides = {
        k: v for k, v in overrides.items() if k in gconfig.DEFAULT_CONFIG["train"]
    }
    model_overrides = {
        k: v for k, v in overrides.items() if k in gconfig.DEFAULT_CONFIG["model"]
    }
    router_overrides = overrides.get("router")
    if router_overrides:
        model_overrides["router"] = router_overrides
    return gconfig.merge_cli_overrides(
        gconfig.DEFAULT_CONFIG, None, train_overrides, model_overrides
    )


_NOW = datetime(2026, 7, 29, 14, 30)


def test_default_out_dir_name_all_defaults_is_just_the_timestamp():
    cfg = _default_cfg()
    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430"


def test_default_out_dir_name_single_non_default_field():
    cfg = _default_cfg(hidden_dim=512)
    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_h512"


def test_default_out_dir_name_conditioning_embedding_shown_abbreviated():
    cfg = _default_cfg(conditioning="embedding")
    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_cemb"


def test_default_out_dir_name_conditioning_default_omitted():
    cfg = _default_cfg(conditioning="physical")
    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430"


def test_default_out_dir_name_router_enabled_shown_as_unit():
    cfg = _default_cfg(router={"enabled": True, "type": "energy", "n_experts": 8})
    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_r-energy8"


def test_default_out_dir_name_router_disabled_omitted_even_if_subfields_nondefault():
    cfg = _default_cfg(router={"enabled": False, "type": "pdg", "n_experts": 8})
    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430"


def test_default_out_dir_name_router_gumbel_shown_when_enabled():
    cfg = _default_cfg(
        router={"enabled": True, "type": "energy", "n_experts": 8, "gumbel": True}
    )
    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_r-energy8_gum"


def test_default_out_dir_name_router_gumbel_omitted_when_router_disabled():
    cfg = _default_cfg(
        router={"enabled": False, "type": "energy", "n_experts": 8, "gumbel": True}
    )
    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430"


def test_default_out_dir_name_router_learn_centers_shown_only_when_disabled():
    cfg_default = _default_cfg(
        router={"enabled": True, "type": "energy", "n_experts": 8}
    )
    assert (
        gconfig.default_out_dir_name(cfg_default, now=_NOW) == "20260729_1430_r-energy8"
    )

    cfg_off = _default_cfg(
        router={
            "enabled": True,
            "type": "energy",
            "n_experts": 8,
            "learn_centers": False,
        }
    )
    assert (
        gconfig.default_out_dir_name(cfg_off, now=_NOW)
        == "20260729_1430_r-energy8_nolc"
    )


def test_default_out_dir_name_router_learn_width_and_temperature_shown():
    cfg = _default_cfg(
        router={
            "enabled": True,
            "type": "energy",
            "n_experts": 8,
            "learn_width": True,
        }
    )
    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_r-energy8_lw"

    cfg2 = _default_cfg(
        router={
            "enabled": True,
            "type": "energy",
            "n_experts": 8,
            "learn_temperature": True,
        }
    )
    assert gconfig.default_out_dir_name(cfg2, now=_NOW) == "20260729_1430_r-energy8_lt"


def test_default_out_dir_name_mode_shown_bare_no_prefix():
    cfg = _default_cfg(mode="wgan")
    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_wgan"


def test_default_out_dir_name_overflow_caps_and_hashes_remainder():
    cfg = _default_cfg(
        mode="wgan",
        router={"enabled": True, "type": "energy", "n_experts": 8},
        conditioning="embedding",
        hidden_dim=512,
        n_blocks=8,
        emb_dim=32,
        lr=1e-3,
        batch_size=2048,
        seed=3,
        epochs=200,
    )
    name = gconfig.default_out_dir_name(cfg, now=_NOW)
    # First 6 by priority: mode, router, conditioning, hidden_dim, n_blocks, emb_dim.
    assert name.startswith("20260729_1430_wgan_r-energy8_cemb_h512_b8_e32_+4more-")
    digest = name.split("-")[-1]
    assert len(digest) == 6


def test_default_out_dir_name_overflow_hash_is_deterministic_and_value_sensitive():
    base = dict(
        mode="wgan",
        router={"enabled": True, "type": "energy", "n_experts": 8},
        conditioning="embedding",
        hidden_dim=512,
        n_blocks=8,
        emb_dim=32,
        lr=1e-3,
        batch_size=2048,
        seed=3,
        epochs=200,
    )
    name_a = gconfig.default_out_dir_name(_default_cfg(**base), now=_NOW)
    name_b = gconfig.default_out_dir_name(_default_cfg(**base), now=_NOW)
    assert name_a == name_b  # stable across calls with the same overflow set

    changed = dict(base, epochs=999)
    name_c = gconfig.default_out_dir_name(_default_cfg(**changed), now=_NOW)
    assert (
        name_c != name_a
    )  # differs when an overflowed value changes  # n_blocks inherited, hidden_dim not