import json

from coremate import cli
from coremate.ai import TrainingPlan, TrainingPreset
from coremate.cli import main


def test_cli_ai_report(monkeypatch, capsys):
    monkeypatch.setattr(cli, "detect_ai_stack", lambda: {"cuda_available": True})
    exit_code = main(["ai", "report"])
    payload = json.loads(capsys.readouterr().out)

    assert exit_code == 0
    assert payload["cuda_available"] is True


def test_cli_ai_recommend(monkeypatch, capsys):
    monkeypatch.setattr(
        cli,
        "recommend_training_preset",
        lambda task, aggressive=False: TrainingPreset(
            task=task,
            accelerator="cuda",
            precision="float16",
            batch_size=2,
            gradient_accumulation_steps=2,
            gradient_checkpointing=True,
            compile_model=True,
            dataloader_num_workers=4,
            pin_memory=True,
            optimizer="adamw_fused",
            flash_attention=True,
        ),
    )
    exit_code = main(["ai", "recommend", "--task", "llm"])
    payload = json.loads(capsys.readouterr().out)

    assert exit_code == 0
    assert payload["task"] == "llm"
    assert payload["precision"] == "float16"


def test_cli_ai_plan(monkeypatch, capsys):
    monkeypatch.setattr(
        cli,
        "build_training_plan",
        lambda **kwargs: TrainingPlan(
            task="llm",
            model_scale="base",
            accelerator="cuda",
            device="cuda",
            precision="float16",
            amp_dtype="torch.float16",
            per_device_batch_size=2,
            gradient_accumulation_steps=8,
            global_batch_size=16,
            gradient_checkpointing=True,
            compile_model=True,
            flash_attention=True,
            dataloader_num_workers=4,
            pin_memory=True,
            optimizer="adamw_fused",
            learning_rate=2e-4,
            weight_decay=0.1,
            warmup_ratio=0.03,
            max_grad_norm=1.0,
            notes=[],
        ),
    )
    monkeypatch.setattr(
        cli,
        "recommend_torch_training_kwargs",
        lambda plan: {"device": "cuda", "trainer": {"global_batch_size": 16}},
    )
    exit_code = main(["ai", "plan", "--task", "llm", "--target-global-batch", "16"])
    payload = json.loads(capsys.readouterr().out)

    assert exit_code == 0
    assert payload["plan"]["task"] == "llm"
    assert payload["training_kwargs"]["trainer"]["global_batch_size"] == 16


def test_cli_ai_seed(monkeypatch, capsys):
    monkeypatch.setattr(
        cli,
        "set_global_seed",
        lambda seed, deterministic=False: {"seed": seed, "deterministic": deterministic},
    )
    exit_code = main(["ai", "seed", "--seed", "42", "--deterministic"])
    payload = json.loads(capsys.readouterr().out)

    assert exit_code == 0
    assert payload["seed"] == 42
    assert payload["deterministic"] is True


def test_cli_ai_tune_without_torch_returns_error(monkeypatch, capsys):
    monkeypatch.setattr(
        cli,
        "tune_torch_runtime",
        lambda **kwargs: (_ for _ in ()).throw(RuntimeError("PyTorch is not installed")),
    )
    exit_code = main(["ai", "tune"])
    payload = json.loads(capsys.readouterr().out)

    assert exit_code == 2
    assert "PyTorch is not installed" in payload["error"]


def test_cli_ai_optimize(monkeypatch, capsys):
    monkeypatch.setattr(
        cli,
        "optimize_torch_cuda",
        lambda **kwargs: {"ok": True, "task": kwargs["task"]},
    )
    exit_code = main(["ai", "optimize", "--task", "cv"])
    payload = json.loads(capsys.readouterr().out)

    assert exit_code == 0
    assert payload == {"ok": True, "task": "cv"}
