# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.

"""Tests for the validate CLI script and related error-handling utilities."""

import json
import warnings
from dataclasses import asdict
from pathlib import Path
from unittest import mock

import pytest
import torch
import yaml

from litgpt import GPT
from litgpt.config import Config
from litgpt.utils import (
    CheckpointValidationResult,
    estimate_model_memory,
    validate_checkpoint,
)

# ---------------------------------------------------------------------------
# validate_checkpoint tests
# ---------------------------------------------------------------------------


class TestValidateCheckpoint:
    """Tests for the validate_checkpoint utility."""

    @staticmethod
    def _save_model_checkpoint(model: torch.nn.Module, path: Path) -> None:
        torch.save(model.state_dict(), str(path))

    def test_valid_checkpoint(self, tmp_path):
        """A checkpoint saved from the same model should pass validation."""
        config = Config.from_name("pythia-14m")
        with torch.device("meta"):
            model = GPT(config)
        # Create a real state_dict with matching shapes
        real_sd = {k: torch.randn(v.shape) for k, v in model.state_dict().items()}
        ckpt_path = tmp_path / "lit_model.pth"
        torch.save(real_sd, str(ckpt_path))

        result = validate_checkpoint(ckpt_path, model, verbose=False)
        assert result.is_valid
        assert result.missing_keys == []
        assert result.unexpected_keys == []
        assert result.shape_mismatches == []
        assert result.errors == []
        assert "passed" in result.summary().lower()

    def test_missing_keys(self, tmp_path):
        """Checkpoint missing some keys should report them."""
        config = Config.from_name("pythia-14m")
        with torch.device("meta"):
            model = GPT(config)
        real_sd = {k: torch.randn(v.shape) for k, v in model.state_dict().items()}
        # Remove a key
        removed_key = list(real_sd.keys())[0]
        del real_sd[removed_key]
        ckpt_path = tmp_path / "lit_model.pth"
        torch.save(real_sd, str(ckpt_path))

        result = validate_checkpoint(ckpt_path, model, verbose=False)
        assert not result.is_valid
        assert removed_key in result.missing_keys
        assert result.unexpected_keys == []

    def test_unexpected_keys(self, tmp_path):
        """Checkpoint with extra keys should report them."""
        config = Config.from_name("pythia-14m")
        with torch.device("meta"):
            model = GPT(config)
        real_sd = {k: torch.randn(v.shape) for k, v in model.state_dict().items()}
        real_sd["extra.unexpected.key"] = torch.randn(3)
        ckpt_path = tmp_path / "lit_model.pth"
        torch.save(real_sd, str(ckpt_path))

        result = validate_checkpoint(ckpt_path, model, verbose=False)
        assert not result.is_valid
        assert "extra.unexpected.key" in result.unexpected_keys

    def test_shape_mismatch(self, tmp_path):
        """Checkpoint with wrong shapes should report mismatches."""
        config = Config.from_name("pythia-14m")
        with torch.device("meta"):
            model = GPT(config)
        real_sd = {k: torch.randn(v.shape) for k, v in model.state_dict().items()}
        # Corrupt a shape
        key = "lm_head.weight"
        real_sd[key] = torch.randn(10, 10)  # wrong shape
        ckpt_path = tmp_path / "lit_model.pth"
        torch.save(real_sd, str(ckpt_path))

        result = validate_checkpoint(ckpt_path, model, verbose=False)
        assert not result.is_valid
        assert any(key in m for m in result.shape_mismatches)

    def test_file_not_found(self, tmp_path):
        """Non-existent checkpoint should report an error."""
        config = Config.from_name("pythia-14m")
        with torch.device("meta"):
            model = GPT(config)
        result = validate_checkpoint(tmp_path / "nonexistent.pth", model, verbose=False)
        assert not result.is_valid
        assert any("not found" in e for e in result.errors)

    def test_corrupted_file(self, tmp_path):
        """A file that is not a valid PyTorch checkpoint should report an error."""
        config = Config.from_name("pythia-14m")
        with torch.device("meta"):
            model = GPT(config)
        ckpt_path = tmp_path / "corrupted.pth"
        ckpt_path.write_text("this is not a checkpoint")

        result = validate_checkpoint(ckpt_path, model, verbose=False)
        assert not result.is_valid
        assert any("Failed to load" in e for e in result.errors)

    def test_model_key_wrapper(self, tmp_path):
        """Checkpoint wrapped under a 'model' key should be unwrapped."""
        config = Config.from_name("pythia-14m")
        with torch.device("meta"):
            model = GPT(config)
        real_sd = {k: torch.randn(v.shape) for k, v in model.state_dict().items()}
        wrapped = {"model": real_sd}
        ckpt_path = tmp_path / "lit_model.pth"
        torch.save(wrapped, str(ckpt_path))

        result = validate_checkpoint(ckpt_path, model, verbose=False)
        assert result.is_valid

    def test_summary_format(self):
        """Summary strings should be well-formed."""
        result = CheckpointValidationResult(
            is_valid=True, missing_keys=[], unexpected_keys=[], shape_mismatches=[], errors=[]
        )
        assert result.summary() == "Checkpoint validation passed."

        result = CheckpointValidationResult(
            is_valid=False,
            missing_keys=["a", "b"],
            unexpected_keys=["c"],
            shape_mismatches=["d: model=(2,3), checkpoint=(4,5)"],
            errors=[],
        )
        summary = result.summary()
        assert "Missing keys" in summary
        assert "Unexpected keys" in summary
        assert "Shape mismatches" in summary


# ---------------------------------------------------------------------------
# estimate_model_memory tests
# ---------------------------------------------------------------------------


class TestEstimateModelMemory:
    """Tests for the estimate_model_memory utility."""

    def test_basic_estimation(self):
        """Should return reasonable estimates for a known config."""
        config = Config.from_name("pythia-14m")
        result = estimate_model_memory(config, dtype=torch.float32, training=False)
        assert result["param_count"] > 0
        assert result["param_memory_gb"] > 0
        assert result["estimated_total_gb"] > 0
        # pythia-14m is ~14M params → ~0.05 GB in fp32
        assert result["param_memory_gb"] < 1.0

    def test_training_multiplier(self):
        """Training should use ~4× multiplier (params + gradients + Adam optimizer states)."""
        config = Config.from_name("pythia-14m")
        inference = estimate_model_memory(config, dtype=torch.float32, training=False)
        training = estimate_model_memory(config, dtype=torch.float32, training=True)
        assert training["estimated_total_gb"] > inference["estimated_total_gb"]
        # Should be approximately 4× (params + gradients + Adam optimizer states).
        # Bounds are loose (3.5–4.5) to absorb rounding from the two round() calls in the function.
        ratio = training["estimated_total_gb"] / inference["estimated_total_gb"]
        assert 3.5 < ratio < 4.5

    def test_dtype_affects_memory(self):
        """Half precision should use ~half the parameter memory."""
        config = Config.from_name("pythia-14m")
        fp32 = estimate_model_memory(config, dtype=torch.float32, training=False)
        fp16 = estimate_model_memory(config, dtype=torch.float16, training=False)
        assert fp16["param_memory_gb"] < fp32["param_memory_gb"]
        # Should be approximately double (exact ratio depends on estimate granularity)
        ratio = fp32["param_memory_gb"] / fp16["param_memory_gb"]
        assert 1.5 < ratio < 2.5

    def test_gpu_fields(self):
        """GPU-related fields should be None when no GPU is available."""
        config = Config.from_name("pythia-14m")
        with mock.patch("litgpt.utils.torch.cuda.is_available", return_value=False):
            result = estimate_model_memory(config, dtype=torch.float32)
        assert result["available_gpu_memory_gb"] is None
        assert result["fits_in_memory"] is None


# ---------------------------------------------------------------------------
# Tokenizer JSON warning test
# ---------------------------------------------------------------------------


def test_tokenizer_json_warning(tmp_path):
    """Tokenizer should emit a warning when generation_config.json has invalid JSON."""
    # Set up a minimal tokenizer directory with an HF tokenizer
    checkpoint_dir = tmp_path / "test_model"
    checkpoint_dir.mkdir()

    # We need tokenizer.json for the HF path and a valid tokenizer_config.json
    # Use a minimal invalid generation_config.json
    invalid_json = '{\n  "bos_token_id": 1,\n  "eos_token_id": 2,\n}'  # trailing comma

    (checkpoint_dir / "generation_config.json").write_text(invalid_json)
    (checkpoint_dir / "tokenizer_config.json").write_text(json.dumps({"tokenizer_class": "GPT2Tokenizer"}))

    # Create a minimal tokenizer.json that the HF tokenizer can load
    minimal_tokenizer_json = {
        "version": "1.0",
        "truncation": None,
        "padding": None,
        "added_tokens": [],
        "normalizer": None,
        "pre_tokenizer": None,
        "post_processor": None,
        "decoder": None,
        "model": {"type": "BPE", "vocab": {"<s>": 0, "</s>": 1}, "merges": []},
    }
    (checkpoint_dir / "tokenizer.json").write_text(json.dumps(minimal_tokenizer_json))

    from litgpt.tokenizer import Tokenizer

    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        tokenizer = Tokenizer(checkpoint_dir)

    # Check that the warning was raised
    assert len(w) == 1
    assert "invalid JSON" in str(w[0].message)
    # Check that the fix worked
    assert tokenizer.bos_id == 1
    assert tokenizer.eos_id == 2


# ---------------------------------------------------------------------------
# Validate script integration test
# ---------------------------------------------------------------------------


class TestValidateScript:
    """Integration tests for the validate CLI script."""

    @staticmethod
    def _make_checkpoint_dir(tmp_path: Path, config_name: str = "pythia-14m") -> Path:
        """Create a fake but structurally valid checkpoint directory."""
        checkpoint_dir = tmp_path / "checkpoints" / "test"
        checkpoint_dir.mkdir(parents=True)

        config = Config.from_name(config_name)
        config_dict = asdict(config)
        with open(checkpoint_dir / "model_config.yaml", "w") as f:
            yaml.dump(config_dict, f)

        # Create a real checkpoint
        with torch.device("meta"):
            model = GPT(config)
        real_sd = {k: torch.randn(v.shape) for k, v in model.state_dict().items()}
        torch.save(real_sd, str(checkpoint_dir / "lit_model.pth"))

        # Create minimal tokenizer files
        (checkpoint_dir / "tokenizer_config.json").write_text(json.dumps({"tokenizer_class": "GPT2Tokenizer"}))
        (checkpoint_dir / "tokenizer.json").write_text("{}")

        return checkpoint_dir

    def test_validate_missing_dir(self, tmp_path, capsys):
        """validate_setup should fail for a non-existent directory."""
        from litgpt.scripts.validate import validate_setup

        with pytest.raises(SystemExit):
            validate_setup(checkpoint_dir=tmp_path / "nonexistent")

    def test_validate_missing_model_file(self, tmp_path, capsys):
        """validate_setup should fail when checkpoint file is missing."""
        checkpoint_dir = tmp_path / "test"
        checkpoint_dir.mkdir()
        (checkpoint_dir / "model_config.yaml").write_text(yaml.dump(asdict(Config.from_name("pythia-14m"))))
        (checkpoint_dir / "tokenizer_config.json").write_text(json.dumps({"tokenizer_class": "GPT2Tokenizer"}))
        (checkpoint_dir / "tokenizer.json").write_text("{}")
        # No lit_model.pth

        from litgpt.scripts.validate import validate_setup

        with pytest.raises(SystemExit):
            validate_setup(checkpoint_dir=checkpoint_dir)
