"""Download SAM2 ONNX model weights from HuggingFace."""

import urllib.request
from collections.abc import Callable
from functools import partial
from pathlib import Path

from samgis_core import app_logger
from samgis_core.prediction_api.model_registry import (
    MODELS,
    DEFAULT_MODEL,
    get_model_dir,
    verify_download,
)


def _report_progress(
    block_num: int,
    block_size: int,
    total_size: int,
    *,
    file_idx: int,
    file_count: int,
    on_progress: Callable[[int, str], None],
) -> None:
    """urlretrieve reporthook that translates per-file progress to overall %."""
    if total_size > 0:
        file_pct = min(block_num * block_size * 100 // total_size, 100)
    else:
        file_pct = 0
    overall = (file_idx * 100 + file_pct) // file_count
    on_progress(min(overall, 99), "")


def download_model(
    variant: str = DEFAULT_MODEL,
    *,
    target_dir: Path | None = None,
    on_progress: Callable[[int, str], None] | None = None,
) -> Path:
    """Download model weights for a variant, skipping files that already exist.

    Args:
        variant: model variant key (must be in MODELS).
        target_dir: override directory. Defaults to get_model_dir(variant).
        on_progress: optional callback(percent, message) for progress updates.

    Returns:
        Path to the model directory.

    Raises:
        KeyError: unknown variant.
        urllib.error.URLError / OSError: download failure.
    """
    model_info = MODELS[variant]
    model_dir = target_dir or get_model_dir(variant)
    model_dir.mkdir(parents=True, exist_ok=True)

    files = [
        (remote, local, model_dir / local)
        for remote, local in model_info["files"].items()
    ]
    total = len(files)

    for idx, (remote, local, target_path) in enumerate(files):
        if target_path.is_file():
            app_logger.debug("skipping existing file", file=local)
            if on_progress:
                on_progress((idx + 1) * 100 // total, f"Skipped {local}")
            continue

        url = model_info["urls"][remote]
        app_logger.info("downloading", url=url, target=str(target_path))
        if on_progress:
            on_progress(
                idx * 100 // total, f"Downloading {local} ({idx + 1}/{total})..."
            )

        hook = (
            partial(
                _report_progress,
                file_idx=idx,
                file_count=total,
                on_progress=on_progress,
            )
            if on_progress
            else None
        )
        urllib.request.urlretrieve(url, str(target_path), hook)

    if on_progress:
        on_progress(100, "Verifying checksums...")

    failures = verify_download(variant, model_dir=model_dir)
    if failures:
        msg = f"SHA-256 mismatch for: {', '.join(failures)}"
        app_logger.error(
            "checksum verification failed", variant=variant, files=failures
        )
        raise OSError(msg)

    if on_progress:
        on_progress(100, "Done")

    app_logger.info("model ready", variant=variant, dir=str(model_dir))
    return model_dir
