import msgpack
from importlib.resources import files

MODEL_FORMAT_VERSION = 1


def normalize_model(model):
    """Normalize legacy/new model payloads to a single runtime schema."""
    normalized = dict(model)

    if "bigram_next" not in normalized and "bigram" in normalized:
        normalized["bigram_next"] = normalized["bigram"]

    if "trigram_topk" not in normalized and "trigram" in normalized:
        normalized["trigram_topk"] = normalized["trigram"]

    normalized.setdefault("model_format_version", MODEL_FORMAT_VERSION)
    normalized.setdefault("chargram", {})

    required = ("unigram", "bigram_next", "trigram_topk")
    missing = [key for key in required if key not in normalized]
    if missing:
        raise ValueError(
            f"Invalid model file, missing keys: {', '.join(missing)}"
        )

    return normalized


def load_model(model_path=None):
    if model_path:
        with open(model_path, "rb") as f:
            raw = msgpack.unpackb(f.read(), raw=False)
        return normalize_model(raw)

    package_root = files("khmer_keyboard")
    candidates = (
        package_root / "model.msgpack",
        package_root / "data" / "model.msgpack",  # backward compatibility
    )

    for model_file in candidates:
        try:
            with model_file.open("rb") as f:
                raw = msgpack.unpackb(f.read(), raw=False)
            return normalize_model(raw)
        except FileNotFoundError:
            continue

    raise FileNotFoundError(
        "Could not find bundled model.msgpack in package. "
        "Expected one of: khmer_keyboard/model.msgpack or "
        "khmer_keyboard/data/model.msgpack."
    )
