import argparse
import hashlib
import os
import sys
import urllib.request
from pathlib import Path
from typing import Optional

from .paths import resolve_base_dir, ensure_app_dirs, resolve_dir

DEFAULT_SAM_URL = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
DEFAULT_SAM_NAME = "sam_vit_b.pth"
DEFAULT_YOLO_URL = "https://raw.githubusercontent.com/Nathanielneil/smpas-models/main/yolo_mushroom.pt"
DEFAULT_YOLO_NAME = "yolo_mushroom.pt"


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _download(url: str, dest: Path, expected_sha256: Optional[str], force: bool) -> None:
    if dest.exists() and not force:
        print(f"Skip (exists): {dest}")
        return

    dest.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = dest.with_suffix(dest.suffix + ".tmp")

    print(f"Downloading: {url}")
    with urllib.request.urlopen(url) as response, tmp_path.open("wb") as handle:
        total = response.getheader("Content-Length")
        total_size = int(total) if total else None
        downloaded = 0
        while True:
            chunk = response.read(1024 * 1024)
            if not chunk:
                break
            handle.write(chunk)
            downloaded += len(chunk)
            if total_size:
                percent = downloaded / total_size * 100
                sys.stdout.write(f"\r  {dest.name}: {percent:.1f}%")
                sys.stdout.flush()
        if total_size:
            sys.stdout.write("\n")

    tmp_path.replace(dest)

    if expected_sha256:
        actual = _sha256(dest)
        if actual.lower() != expected_sha256.lower():
            raise ValueError(
                f"SHA256 mismatch for {dest.name}: expected {expected_sha256}, got {actual}"
            )

    print(f"Saved: {dest}")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Download SMPAS model files.")
    parser.add_argument("--base-dir", default=os.getenv("SMPAS_HOME", os.getenv("MUSHROOM_HOME")))
    parser.add_argument("--models-dir", default=os.getenv("SMPAS_MODELS_DIR", os.getenv("MUSHROOM_MODELS_DIR")))
    parser.add_argument("--yolo-url", default=os.getenv("SMPAS_YOLO_URL", DEFAULT_YOLO_URL))
    parser.add_argument("--sam-url", default=os.getenv("SMPAS_SAM_URL", DEFAULT_SAM_URL))
    parser.add_argument("--yolo-name", default=DEFAULT_YOLO_NAME)
    parser.add_argument("--sam-name", default=DEFAULT_SAM_NAME)
    parser.add_argument("--yolo-sha256", default=os.getenv("SMPAS_YOLO_SHA256"))
    parser.add_argument("--sam-sha256", default=os.getenv("SMPAS_SAM_SHA256"))
    parser.add_argument("--no-yolo", action="store_true")
    parser.add_argument("--no-sam", action="store_true")
    parser.add_argument("--force", action="store_true")
    return parser


def main() -> None:
    args = build_parser().parse_args()

    base_dir = resolve_base_dir(args.base_dir)
    ensure_app_dirs(base_dir)
    models_dir = resolve_dir(base_dir, args.models_dir, "SMPAS_MODELS_DIR", "models")
    models_dir.mkdir(parents=True, exist_ok=True)

    if not args.no_yolo:
        yolo_dest = models_dir / args.yolo_name
        _download(args.yolo_url, yolo_dest, args.yolo_sha256, args.force)

    if not args.no_sam:
        if not args.sam_url:
            print("ERROR: 缺少 SAM 模型下载地址，请使用 --sam-url 或设置 SMPAS_SAM_URL")
            sys.exit(2)
        sam_dest = models_dir / args.sam_name
        _download(args.sam_url, sam_dest, args.sam_sha256, args.force)

    print("Done.")
    print(f"Models dir: {models_dir}")


if __name__ == "__main__":
    main()
