from typer.testing import CliRunner
from mint.cli import app
from mint import __version__

from tests.utils.download_model import download_model, MODEL_PATH

import torch
from safetensors.torch import load_file, save_file

DOWNLOADED = download_model()


def get_model(tmp_path):
    if DOWNLOADED:
        state_dict = load_file(str(MODEL_PATH))
        for key, value in state_dict.items():
            if "embedding" in key or "embed_tokens" in key:
                return MODEL_PATH, value
        raise RuntimeError("No embedding tensor found in downloaded model")
    embeddings = torch.arange(6, dtype=torch.float32).reshape(2, 3)
    model_file = tmp_path / "model.safetensors"
    save_file({"embed_tokens.weight": embeddings}, str(model_file))
    return model_file, embeddings


def test_cli_main():
    runner = CliRunner()
    result = runner.invoke(app)
    assert result.exit_code == 0
    assert "Usage:" in result.stdout


def test_cli_version():
    runner = CliRunner()
    result = runner.invoke(app, ["--version"])
    assert result.exit_code == 0
    assert __version__ in result.stdout


def test_cli_extract(tmp_path):
    runner = CliRunner()
    model_file, expected = get_model(tmp_path)

    out_file = tmp_path / "embeddings.safetensors"
    result = runner.invoke(app, ["extract", str(model_file), str(out_file)])
    assert result.exit_code == 0
    saved = load_file(str(out_file))["embedding"]
    assert torch.equal(saved, expected)


def test_cli_extract_duplicate_embeddings(tmp_path):
    runner = CliRunner()
    emb1 = torch.zeros(2, 3)
    emb2 = torch.ones(2, 3)
    model_file = tmp_path / "model.safetensors"
    save_file({"embedding.weight": emb1, "embed_tokens.weight": emb2}, str(model_file))

    out_file = tmp_path / "embeddings.safetensors"
    result = runner.invoke(app, ["extract", str(model_file), str(out_file)])
    assert result.exit_code != 0
    assert "Expected one embedding tensor" in str(result.exception)


def test_cli_extract_creates_subdirs(tmp_path):
    runner = CliRunner()
    model_file, expected = get_model(tmp_path)

    out_file = tmp_path / "out" / "embeddings.safetensors"
    result = runner.invoke(app, ["extract", str(model_file), str(out_file)])
    assert result.exit_code == 0
    assert out_file.exists()
    saved = load_file(str(out_file))["embedding"]
    assert torch.equal(saved, expected)


def test_cli_extract_pt(tmp_path):
    runner = CliRunner()
    embeddings = torch.arange(6, dtype=torch.float32).reshape(2, 3)
    model_file = tmp_path / "model.pt"
    torch.save({"embed_tokens.weight": embeddings}, model_file)

    out_file = tmp_path / "embeddings.safetensors"
    result = runner.invoke(app, ["extract", str(model_file), str(out_file)])
    assert result.exit_code == 0
    saved = load_file(str(out_file))["embedding"]
    assert torch.equal(saved, embeddings)
