#!/usr/bin/env python3

from __future__ import annotations

import argparse
import json
import os
import re
from pathlib import Path
from typing import Any

from huggingface_hub import hf_hub_download, snapshot_download


DEFAULT_ROOT = "/mnt/llm_storage"


def safe_name(value: str) -> str:
    """
    Keep names readable while avoiding problematic path characters.
    Allows repo-like names to become nested folders only when intentional elsewhere.
    """
    value = value.strip()
    value = value.replace("\\", "/")
    value = re.sub(r"[^A-Za-z0-9._/\-]+", "_", value)
    value = re.sub(r"_+", "_", value)
    return value.strip("_/")


def split_hub_id(hub_id: str) -> tuple[str, str | None]:
    """
    Hugging Face repo IDs normally look like:
        owner/repo

    Sometimes manifests accidentally store a subfolder inside hub_id:
        owner/repo/subfolder/path

    This converts that into:
        repo_id = owner/repo
        subfolder = subfolder/path
    """
    parts = hub_id.strip("/").split("/")
    if len(parts) <= 2:
        return hub_id, None

    repo_id = "/".join(parts[:2])
    subfolder = "/".join(parts[2:])
    return repo_id, subfolder


def route_destination(root: str, model: dict[str, Any]) -> Path:
    name = safe_name(model.get("hub_id") or model.get("name") or model.get("folder"))
    
    framework = model.get("framework", "").strip()
    task = safe_name(model.get("primary_task", "misc"))

    if task == "dataset":
        return root,"datasets",name

    if framework == "llama_cpp":
        return root,"models","gguf",task,name

    if framework == "transformers":
        return root,"models","transformers",task,name

    return root,"models","misc",task,name


def download_one(key: str, model: dict[str, Any], root: str, dry_run: bool = False) -> None:
    hub_id = model.get("hub_id")
    framework = model.get("framework")
    filename = model.get("filename")
    include = model.get("include")

    if not hub_id:
        print(f"[SKIP] {key}: missing hub_id")
        return
    destination = route_destination(root, model)
    repo_id, subfolder_from_hub = split_hub_id(hub_id)

    print()
    print(f"[MODEL] {key}")
    print(f"  framework: {framework}")
    print(f"  task:      {model.get('primary_task')}")
    print(f"  repo_id:   {repo_id}")
    if subfolder_from_hub:
        print(f"  subfolder: {subfolder_from_hub}")
    print(f"  dest:      {destination}")

    if dry_run:
        return

    destination.mkdir(parents=True, exist_ok=True)

    # GGUF/single-file style download.
    if framework == "llama_cpp":
        if filename:
            hf_hub_download(
                repo_id=repo_id,
                filename=filename,
                subfolder=subfolder_from_hub,
                local_dir=destination,
                local_dir_use_symlinks=False,
            )
            print(f"  downloaded file: {filename}")
            return

        # Pattern-based GGUF download.
        if include:
            snapshot_download(
                repo_id=repo_id,
                allow_patterns=include,
                local_dir=destination,
                local_dir_use_symlinks=False,
            )
            print(f"  downloaded pattern: {include}")
            return

        # Fallback: download repo snapshot.
        snapshot_download(
            repo_id=repo_id,
            local_dir=destination,
            local_dir_use_symlinks=False,
        )
        print("  downloaded full snapshot")
        return

    # Dataset-style snapshot.
    if model.get("task") == "dataset":
        snapshot_download(
            repo_id=repo_id,
            repo_type="dataset",
            local_dir=destination,
            local_dir_use_symlinks=False,
        )
        print("  downloaded dataset snapshot")
        return

    # Transformers/full repo download.
    allow_patterns = include if include else None

    snapshot_download(
        repo_id=repo_id,
        allow_patterns=allow_patterns,
        local_dir=destination,
        local_dir_use_symlinks=False,
    )

    if allow_patterns:
        print(f"  downloaded snapshot with pattern: {allow_patterns}")
    else:
        print("  downloaded full snapshot")


def download_dict_models() -> None:
    parser = argparse.ArgumentParser(description="Auto-distribute LLM downloads by framework/task.")
    parser.add_argument(
        "manifest",
        type=str,
        help="Path to model manifest JSON.",
    )
    parser.add_argument(
        "--root",
        type=str,
        default=DEFAULT_ROOT,
        help="LLM storage root.",
    )
    parser.add_argument(
        "--only",
        nargs="*",
        default=None,
        help="Optional model keys to download.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print destinations without downloading.",
    )

    args = parser.parse_args()

    os.environ.setdefault("HF_HOME", os.path.join(args.root,"cache","huggingface"))
    os.environ.setdefault("HF_HUB_CACHE", os.path.join(args.root,"cache","huggingface","hub"))
    os.environ.setdefault("TORCH_HOME", os.path.join(args.root,"cache","torch"))
    os.environ.setdefault("PIP_CACHE_DIR", os.path.join(args.root,"cache","pip"))

    with args.manifest.open("r", encoding="utf-8") as f:
        manifest = json.load(f)

    for key, model in manifest.items():
        if args.only and key not in args.only:
            continue

        try:
            download_one(key, model, args.root, dry_run=args.dry_run)
        except Exception as exc:
            print(f"[ERROR] {key}: {exc}")
