import os

import pytest

from coremate import ai


def test_detect_ai_stack_without_torch(monkeypatch):
    monkeypatch.setattr(ai, "_load_torch", lambda: None)
    monkeypatch.setattr(
        ai, "_query_nvidia_smi", lambda: {"available": False, "driver_version": None, "gpus": []}
    )

    report = ai.detect_ai_stack()
    assert report["torch_installed"] is False
    assert report["cuda_available"] is False
    assert report["cuda_device_count"] == 0


def test_recommend_training_preset_for_llm_small_vram():
    preset = ai.recommend_training_preset(
        task="llm",
        hardware={"cuda_available": True, "max_vram_gb": 8.0, "cpu_logical": 12},
    )
    assert preset.accelerator == "cuda"
    assert preset.batch_size == 1
    assert preset.gradient_checkpointing is True


def test_apply_cuda_env_defaults_sets_expected_vars(monkeypatch):
    monkeypatch.delenv("PYTORCH_CUDA_ALLOC_CONF", raising=False)
    monkeypatch.delenv("CUDA_MODULE_LOADING", raising=False)
    monkeypatch.delenv("TOKENIZERS_PARALLELISM", raising=False)

    applied = ai.apply_cuda_env_defaults(max_split_size_mb=128, expandable_segments=True)
    assert "max_split_size_mb:128" in applied["PYTORCH_CUDA_ALLOC_CONF"]
    assert applied["CUDA_MODULE_LOADING"] == "LAZY"
    assert os.environ["TOKENIZERS_PARALLELISM"] == "false"


class _DummyMatmul:
    def __init__(self):
        self.allow_tf32 = False


class _DummyCudaBackend:
    def __init__(self):
        self.matmul = _DummyMatmul()


class _DummyCudnn:
    def __init__(self):
        self.benchmark = False
        self.deterministic = False
        self.allow_tf32 = False

    @staticmethod
    def version():
        return 9000


class _DummyMps:
    @staticmethod
    def is_available():
        return False


class _DummyBackends:
    def __init__(self):
        self.cudnn = _DummyCudnn()
        self.cuda = _DummyCudaBackend()
        self.mps = _DummyMps()


class _DummyProps:
    total_memory = 12 * 1024**3
    major = 8
    minor = 6


class _DummyCuda:
    def __init__(self):
        self.seed = None

    @staticmethod
    def is_available():
        return True

    @staticmethod
    def device_count():
        return 1

    @staticmethod
    def get_device_name(_):
        return "Dummy GPU"

    @staticmethod
    def get_device_properties(_):
        return _DummyProps()

    def manual_seed_all(self, seed):
        self.seed = seed


class _DummyVersion:
    cuda = "12.4"


class _DummyTorch:
    __version__ = "2.3.0"

    def __init__(self):
        self.cuda = _DummyCuda()
        self.backends = _DummyBackends()
        self.version = _DummyVersion()
        self.manual_seed_value = None
        self.deterministic_algorithms = False
        self.matmul_precision = None
        self.num_threads = None

    def manual_seed(self, seed):
        self.manual_seed_value = seed

    def use_deterministic_algorithms(self, value):
        self.deterministic_algorithms = bool(value)

    def set_float32_matmul_precision(self, value):
        self.matmul_precision = value

    def set_num_threads(self, value):
        self.num_threads = value


def test_tune_torch_runtime_with_dummy_torch(monkeypatch):
    dummy_torch = _DummyTorch()
    monkeypatch.setattr(ai, "_load_torch", lambda: dummy_torch)
    monkeypatch.setattr(ai, "_load_numpy", lambda: None)

    tuned = ai.tune_torch_runtime(
        seed=123,
        deterministic=True,
        benchmark=True,
        allow_tf32=True,
        matmul_precision="high",
        num_threads=6,
    )
    assert tuned["seed_applied"] is True
    assert dummy_torch.manual_seed_value == 123
    assert dummy_torch.backends.cudnn.deterministic is True
    assert dummy_torch.backends.cudnn.benchmark is False
    assert dummy_torch.matmul_precision == "high"
    assert dummy_torch.num_threads == 6


def test_recommend_training_preset_rejects_invalid_task():
    with pytest.raises(ValueError):
        ai.recommend_training_preset(task="audio")


def test_build_training_plan_global_batch_target():
    plan = ai.build_training_plan(
        task="llm",
        model_scale="base",
        target_global_batch_size=32,
        hardware={"cuda_available": True, "max_vram_gb": 12.0, "cpu_logical": 12},
    )
    assert plan.task == "llm"
    assert plan.gradient_accumulation_steps >= 1
    assert plan.global_batch_size >= 32
    assert plan.amp_dtype in {"torch.float16", "torch.bfloat16", "torch.float32"}


def test_recommend_torch_training_kwargs_from_plan():
    plan = ai.TrainingPlan(
        task="cv",
        model_scale="small",
        accelerator="cuda",
        device="cuda",
        precision="float16",
        amp_dtype="torch.float16",
        per_device_batch_size=16,
        gradient_accumulation_steps=2,
        global_batch_size=32,
        gradient_checkpointing=False,
        compile_model=True,
        flash_attention=False,
        dataloader_num_workers=4,
        pin_memory=True,
        optimizer="adamw_fused",
        learning_rate=3e-4,
        weight_decay=0.05,
        warmup_ratio=0.05,
        max_grad_norm=1.0,
    )
    kwargs = ai.recommend_torch_training_kwargs(plan)
    assert kwargs["device"] == "cuda"
    assert kwargs["dataloader"]["pin_memory"] is True
    assert kwargs["trainer"]["global_batch_size"] == 32


def test_optimize_torch_cuda_handles_missing_torch(monkeypatch):
    monkeypatch.setattr(ai, "detect_ai_stack", lambda: {"cuda_available": False, "cpu_logical": 8})
    monkeypatch.setattr(
        ai,
        "recommend_training_preset",
        lambda task, aggressive=False, hardware=None: ai.TrainingPreset(
            task=task,
            accelerator="cpu",
            precision="float32",
            batch_size=4,
            gradient_accumulation_steps=1,
            gradient_checkpointing=False,
            compile_model=False,
            dataloader_num_workers=4,
            pin_memory=False,
            optimizer="adamw",
            flash_attention=False,
        ),
    )
    monkeypatch.setattr(ai, "apply_cuda_env_defaults", lambda **kwargs: {"CUDA_MODULE_LOADING": "LAZY"})
    monkeypatch.setattr(
        ai,
        "tune_torch_runtime",
        lambda **kwargs: (_ for _ in ()).throw(RuntimeError("PyTorch is not installed")),
    )

    result = ai.optimize_torch_cuda(task="general")
    assert result["tuned"] is None
    assert "PyTorch is not installed" in result["tune_error"]
    assert "plan" in result
    assert "training_kwargs" in result


def test_optimize_torch_cuda_returns_tuned_result(monkeypatch):
    monkeypatch.setattr(ai, "detect_ai_stack", lambda: {"cuda_available": True, "cpu_logical": 12})
    monkeypatch.setattr(
        ai,
        "recommend_training_preset",
        lambda task, aggressive=False, hardware=None: ai.TrainingPreset(
            task=task,
            accelerator="cuda",
            precision="float16",
            batch_size=8,
            gradient_accumulation_steps=1,
            gradient_checkpointing=False,
            compile_model=True,
            dataloader_num_workers=6,
            pin_memory=True,
            optimizer="adamw_fused",
            flash_attention=True,
        ),
    )
    monkeypatch.setattr(ai, "apply_cuda_env_defaults", lambda **kwargs: {"CUDA_MODULE_LOADING": "LAZY"})
    monkeypatch.setattr(ai, "tune_torch_runtime", lambda **kwargs: {"torch_available": True})

    result = ai.optimize_torch_cuda(task="cv", aggressive=True)
    assert result["preset"]["task"] == "cv"
    assert result["tuned"]["torch_available"] is True
    assert result["tune_error"] is None
    assert result["plan"]["task"] == "cv"
