"""Tests for `giant new-run` (config.toml + run-dir scaffolding)."""

from __future__ import annotations

import tomllib
from pathlib import Path

from typer.testing import CliRunner

from giant.cli import app

runner = CliRunner()


def test_writes_config_with_overrides_applied(tmp_path: Path):
    out_dir = tmp_path / "run1"
    result = runner.invoke(
        app,
        [
            "new-run",
            "--out",
            str(out_dir),
            "--mode",
            "ddpm",
            "--hidden-dim",
            "128",
            "--n-blocks",
            "4",
            "--lr",
            "0.0005",
        ],
    )
    assert result.exit_code == 0, result.output

    config_path = out_dir / "config.toml"
    assert config_path.exists()
    with open(config_path, "rb") as f:
        cfg = tomllib.load(f)

    assert cfg["train"]["mode"] == "ddpm"
    assert cfg["train"]["lr"] == 0.0005
    assert cfg["model"]["hidden_dim"] == 128
    assert cfg["model"]["n_blocks"] == 4
    # untouched defaults still present
    assert cfg["train"]["epochs"] == 100
    assert "router" in cfg["model"]

    assert str(out_dir) in result.output
    assert "<data.parquet>" in result.output
    assert "giant train" in result.output


def test_comment_and_provenance_recorded_in_meta(tmp_path: Path):
    out_dir = tmp_path / "run2"
    result = runner.invoke(
        app,
        ["new-run", "--out", str(out_dir), "--comment", "quick test"],
    )
    assert result.exit_code == 0, result.output

    with open(out_dir / "config.toml", "rb") as f:
        cfg = tomllib.load(f)

    assert cfg["meta"]["comment"] == "quick test"
    assert cfg["meta"]["created_by"] == "giant new-run"
    assert "created_at" in cfg["meta"]
    assert "git_hash" in cfg["meta"]


def test_data_flag_fills_printed_next_step_commands(tmp_path: Path):
    out_dir = tmp_path / "run3"
    result = runner.invoke(
        app,
        ["new-run", "--out", str(out_dir), "--data", "/ceph/lbogner/train.parquet"],
    )
    assert result.exit_code == 0, result.output
    assert "/ceph/lbogner/train.parquet" in result.output
    assert "<data.parquet>" not in result.output


def test_dry_run_writes_nothing(tmp_path: Path):
    out_dir = tmp_path / "run4"
    result = runner.invoke(
        app,
        ["new-run", "--out", str(out_dir), "--hidden-dim", "512", "--dry-run"],
    )
    assert result.exit_code == 0, result.output
    assert "dry-run" in result.output
    assert "hidden_dim = 512" in result.output
    assert not out_dir.exists()


def test_force_guard_refuses_to_clobber_existing_checkpoints(tmp_path: Path):
    out_dir = tmp_path / "run5"
    out_dir.mkdir()
    (out_dir / "last.pt").touch()

    result = runner.invoke(app, ["new-run", "--out", str(out_dir), "--mode", "ddpm"])
    assert result.exit_code != 0
    assert "already has last.pt" in result.output
    assert not (out_dir / "config.toml").exists()

    result = runner.invoke(
        app, ["new-run", "--out", str(out_dir), "--mode", "ddpm", "--force"]
    )
    assert result.exit_code == 0, result.output
    assert (out_dir / "config.toml").exists()


def test_default_out_dir_used_when_out_omitted(tmp_path: Path, monkeypatch):
    monkeypatch.chdir(tmp_path)
    result = runner.invoke(app, ["new-run", "--hidden-dim", "64"])
    assert result.exit_code == 0, result.output

    checkpoints_dir = tmp_path / "checkpoints"
    run_dirs = list(checkpoints_dir.iterdir()) if checkpoints_dir.exists() else []
    assert len(run_dirs) == 1
    assert (run_dirs[0] / "config.toml").exists()