from __future__ import annotations

import urllib.request
from pathlib import Path

MODELS_DIR = Path(__file__).resolve().parent.parent / "models"
MODELS_DIR.mkdir(exist_ok=True)

MODELS = {
    "qwen2.5-1.5b-instruct-q4_k_m.gguf": (
        "https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct-GGUF/resolve/main/qwen2.5-1.5b-instruct-q4_k_m.gguf"
    ),
    "tr_ocr_base_handwritten.onnx": (
        "https://huggingface.co/Xenova/trocr-base-handwritten/resolve/main/onnx/model.onnx"
    ),
    "tr_ocr_tokenizer.json": (
        "https://huggingface.co/Xenova/trocr-base-handwritten/resolve/main/tokenizer.json"
    ),
    "all-MiniLM-L6-v2.onnx": (
        "https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/onnx/model.onnx"
    ),
    "tokenizer.json": (
        "https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/tokenizer.json"
    ),
    "ggml-tiny.en-q5_0.bin": (
        "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-tiny.en-q5_0.bin"
    ),
}


def download(url: str, dest: Path) -> None:
    if dest.exists() and dest.stat().st_size > 1000:
        print(f"  Already exists: {dest.name} ({dest.stat().st_size // 1024 // 1024} MB)")
        return
    print(f"  Downloading {dest.name}...")
    req = urllib.request.Request(
        url, headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"}
    )
    try:
        with urllib.request.urlopen(req, timeout=30) as response, open(dest, "wb") as out_file:
            total = int(response.headers.get("Content-Length", 0))
            downloaded = 0
            while True:
                chunk = response.read(1024 * 1024)
                if not chunk:
                    break
                out_file.write(chunk)
                downloaded += len(chunk)
                if total:
                    pct = downloaded * 100 // total
                    print(f"\r  {dest.name}: {pct}% ({downloaded // 1024 // 1024}/{total // 1024 // 1024} MB)", end="")
            print(f"\r  Saved: {dest.name} ({dest.stat().st_size // 1024 // 1024} MB)")
    except Exception as e:
        print(f"  FAILED: {dest.name} — {e}")


def check_updates() -> dict[str, bool]:
    results = {}
    for name, url in MODELS.items():
        dest = MODELS_DIR / name
        if dest.exists() and dest.stat().st_size > 1000:
            results[name] = True
        else:
            results[name] = False
    return results


def main() -> None:
    print("Fetching models...")
    for name, url in MODELS.items():
        download(url, MODELS_DIR / name)
    print("Done. Models cached in", MODELS_DIR)


if __name__ == "__main__":
    main()