"""Unit tests: formats, estimator, compatibility engine, OpenAI helpers."""

from __future__ import annotations

from pathlib import Path

from llm_api.models import compat, formats
from llm_api.models.estimator import estimate, kv_bytes_per_token, recommended_context
from llm_api.worker.openai_types import StopMatcher, ThinkSplitter, parse_tool_calls

from .conftest import make_gguf, make_mlx_model


def test_parse_param_count():
    assert formats.parse_param_count_from_name("Qwen3-30B-A3B-4bit") == (30_000_000_000, 3_000_000_000)
    assert formats.parse_param_count_from_name("gemma-3-4b-it-qat-4bit")[0] == 4_000_000_000
    assert formats.parse_param_count_from_name("Qwen3-Embedding-0.6B-8bit")[0] == 600_000_000
    assert formats.parse_param_count_from_name("all-MiniLM-L6-v2-4bit") == (None, None)
    assert formats.parse_param_count_from_name("embeddinggemma-300m-4bit")[0] == 300_000_000


def test_parse_quant():
    assert formats.parse_quant_from_name("Qwen3-4B-Instruct-2507-4bit") == ("4bit", 4.0)
    assert formats.parse_quant_from_name("gpt-oss-20b-MXFP4-Q8")[0] == "MXFP4"
    q, b = formats.parse_quant_from_name("gemma-3-1b-it-Q4_K_M")
    assert q == "Q4_K_M" and b == 4.5
    assert formats.parse_quant_from_name("Qwen3-Embedding-0.6B-4bit-DWQ") == ("DWQ-4bit", 4.0)
    assert formats.parse_quant_from_name("Qwen3.8-27B-bf16") == ("bf16", 16)


def test_family():
    assert formats.guess_family("Qwen3.6-35B-A3B-4bit") == "qwen"
    assert formats.guess_family("gemma-4-26b-a4b-it-4bit") == "gemma"
    assert formats.guess_family("Devstral-Small-2-24B") == "mistral"
    assert formats.guess_family("gpt-oss-20b-MXFP4-Q8") == "gpt-oss"
    assert formats.guess_family("Something-Unknown", "phi3") == "phi"


def test_size_class():
    assert formats.size_class(3, 45) == "TINY"
    assert formats.size_class(7, 45) == "SMALL"
    assert formats.size_class(15, 45) == "MEDIUM"
    assert formats.size_class(30, 45) == "LARGE"
    assert formats.size_class(40, 45) == "XL"
    assert formats.size_class(50, 45) == "TOO_LARGE"


def test_kv_and_estimate():
    kv = kv_bytes_per_token(36, 8, 128, 16)
    assert kv == 2 * 36 * 8 * 128 * 2
    est = estimate(2_400_000_000, "mlx", kv, 32768)
    assert 9 < est.total_gb < 11
    ctx, e2 = recommended_context(2_400_000_000, "mlx", kv, 262144, 45)
    assert ctx == 32768  # capped at the preferred context
    ctx3, _ = recommended_context(40 * 1024**3, "mlx", kv * 4, 131072, 45)
    assert ctx3 is not None and ctx3 <= 8192


def test_compat_statuses(monkeypatch):
    monkeypatch.setattr(compat, "mlx_available", lambda: True)
    monkeypatch.setattr(compat, "mlx_lm_model_types", lambda: {"qwen3"})
    kv = kv_bytes_per_token(36, 8, 128)
    ok = compat.evaluate(runtime="mlx", weights_bytes=2_400_000_000, kv_per_token=kv, max_context=131072, model_type="qwen3",
                         architecture="Qwen3ForCausalLM", vision=False, embedding=False, reranker=False, budget_gb=45, absolute_gb=50,
                         llamacpp_available=True)
    assert ok.status == compat.COMPATIBLE and ok.recommended_context == 32768
    too_big = compat.evaluate(runtime="mlx", weights_bytes=60 * 1024**3, kv_per_token=kv, max_context=131072, model_type="qwen3",
                              architecture=None, vision=False, embedding=False, reranker=False, budget_gb=45, absolute_gb=50,
                              llamacpp_available=True)
    assert too_big.status == compat.INCOMPATIBLE and not too_big.compatible
    swap = compat.evaluate(runtime="mlx", weights_bytes=44 * 1024**3, kv_per_token=kv, max_context=131072, model_type="qwen3",
                           architecture=None, vision=False, embedding=False, reranker=False, budget_gb=45, absolute_gb=50,
                           llamacpp_available=True)
    assert swap.status == compat.NOT_RECOMMENDED
    unknown = compat.evaluate(runtime="mlx", weights_bytes=1e9, kv_per_token=kv, max_context=8192, model_type="mystery",
                              architecture=None, vision=False, embedding=False, reranker=False, budget_gb=45, absolute_gb=50,
                              llamacpp_available=True)
    assert unknown.status == compat.INCOMPATIBLE
    no_llama = compat.evaluate(runtime="llamacpp", weights_bytes=1e9, kv_per_token=kv, max_context=8192, model_type="llama",
                               architecture="llama", vision=False, embedding=False, reranker=False, budget_gb=45, absolute_gb=50,
                               llamacpp_available=False, weights_file="x.gguf")
    assert no_llama.status == compat.INCOMPATIBLE
    exp = compat.evaluate(runtime="llamacpp", weights_bytes=1e9, kv_per_token=kv, max_context=8192, model_type="weird",
                          architecture="weird", vision=False, embedding=False, reranker=False, budget_gb=45, absolute_gb=50,
                          llamacpp_available=True, weights_file="x.gguf")
    assert exp.status == compat.EXPERIMENTAL


def test_gguf_reader(tmp_path: Path):
    p = tmp_path / "m-Q4_K_M.gguf"
    make_gguf(p, arch="llama", layers=8, kv_heads=4, head_dim=64, ctx=4096)
    info = formats.read_gguf(p)
    assert info.architecture == "llama" and info.n_layers == 8 and info.n_kv_heads == 4 and info.head_dim == 64
    assert info.context_length == 4096 and info.file_type_label == "Q4_K_M"
    assert info.param_count == 1_000_000


def test_safetensors_params(tmp_path: Path):
    d = tmp_path / "m"
    make_mlx_model(d, layers=2, hidden=256, bits=4)
    params, nbytes = formats.count_safetensors_params([d / "model.safetensors"], 4)
    # 2 layers * (1024 x 256) + embeddings 1000*256
    assert params == 2 * 1024 * 256 + 1000 * 256
    assert nbytes > 0


def test_hf_config_parse():
    cfg = {"model_type": "qwen3", "architectures": ["Qwen3ForCausalLM"], "num_hidden_layers": 36, "num_attention_heads": 32,
           "num_key_value_heads": 8, "hidden_size": 2560, "head_dim": 128, "max_position_embeddings": 262144,
           "quantization": {"bits": 4, "group_size": 64}}
    p = formats.parse_hf_config(cfg)
    assert p["n_layers"] == 36 and p["n_kv_heads"] == 8 and p["head_dim"] == 128 and p["quant_bits"] == 4 and not p["vision"]
    v = formats.parse_hf_config({"model_type": "qwen3_vl", "text_config": {"num_hidden_layers": 2}, "vision_config": {}})
    assert v["vision"] and v["n_layers"] == 2


def test_stop_matcher():
    sm = StopMatcher(["</s>", "User:"])
    out = sm.feed("Hello <")
    assert out == "Hello " and not sm.done
    out += sm.feed("/s> tail")
    assert out == "Hello " and sm.done
    sm2 = StopMatcher(["STOP"])
    assert sm2.feed("abc ST") == "abc "
    assert sm2.feed("art") == "STart"
    assert sm2.flush() == ""


def test_think_splitter():
    ts = ThinkSplitter("<think>", "</think>")
    r, c = ts.feed("<think>reasoning here</thi")
    assert r == "reasoning here" and c == ""
    r2, c2 = ts.feed("nk>\n\nanswer")
    assert r2 == "" and c2 == "answer"


def test_tool_calls():
    rest, calls = parse_tool_calls('Sure.\n<tool_call>\n{"name": "get_weather", "arguments": {"city": "Montreal"}}\n</tool_call>')
    assert rest == "Sure." and len(calls) == 1 and calls[0]["function"]["name"] == "get_weather"
    assert '"city": "Montreal"' in calls[0]["function"]["arguments"]