from pathlib import Path

import scripts.quant.build_model as bm


def test_baseline_and_online_need_no_build(tmp_path):
    assert bm.needs_build("bf16") is False
    assert bm.needs_build("fp8") is False
    assert bm.needs_build("awq") is True
    assert bm.needs_build("bnb4") is True


def test_build_is_idempotent_when_output_exists(tmp_path, monkeypatch):
    out = tmp_path / "awq"
    out.mkdir()
    (out / "config.json").write_text("{}", encoding="utf-8")
    called = {"n": 0}

    def fake_compressor(*a, **k):
        called["n"] += 1

    monkeypatch.setattr(bm, "_build_compressor", fake_compressor)
    result = bm.build_model("awq", base_model="datalab-to/surya-ocr-2", out_dir=out, calib_images=[])
    assert result == out
    assert called["n"] == 0  # skipped because config.json already present


def test_build_dispatches_to_compressor(tmp_path, monkeypatch):
    out = tmp_path / "gptq"
    seen = {}

    def fake_compressor(method, base_model, out_dir, calib_images):
        seen["method"] = method
        out_dir.mkdir(parents=True, exist_ok=True)
        (out_dir / "config.json").write_text("{}", encoding="utf-8")

    monkeypatch.setattr(bm, "_build_compressor", fake_compressor)
    bm.build_model("gptq", base_model="b", out_dir=out, calib_images=[])
    assert seen["method"] == "gptq"