"""Download manager for model artifacts.

Handles downloading model weights and other artifacts from various sources:
- Direct HTTP/HTTPS URLs
- Google Drive URLs (via gdown)
- Hugging Face Hub (single files and full snapshots)

Supports setting environment variables to control cache directories.
"""

import json
import shutil
import sys
import tarfile
import urllib.request
import zipfile
from pathlib import Path
from typing import Any, Dict, List, Optional


def _looks_like_html_file(path: Path) -> bool:
    """Detect common HTML placeholders/interstitial pages."""
    if not path.exists() or path.is_dir() or path.stat().st_size == 0:
        return False
    try:
        sample = path.read_bytes()[:512].lstrip().lower()
    except Exception:
        return False
    return sample.startswith(b"<!doctype html") or sample.startswith(b"<html")


def _is_expected_destination_ready(path: Path) -> bool:
    """Return True when a destination artifact already looks usable."""
    if not path.exists():
        return False
    if path.is_dir():
        try:
            return any(path.iterdir())
        except Exception:
            return False
    return not _looks_like_html_file(path)


def _is_google_drive_url(url: str) -> bool:
    value = url.strip().lower()
    return "drive.google.com" in value or "drive.usercontent.google.com" in value


def _normalize_patterns(value: Any) -> Optional[List[str]]:
    if value is None:
        return None
    if isinstance(value, list):
        return [str(item) for item in value if str(item).strip()]
    if isinstance(value, str) and value.strip():
        return [value.strip()]
    return None


def download_google_drive(url: str, dest_path: Path) -> None:
    """Download a file from Google Drive using gdown."""
    try:
        import gdown
    except ImportError as exc:
        raise ImportError(
            "gdown is required for Google Drive downloads. "
            "Install it with: pip install gdown"
        ) from exc

    print(f"Downloading Google Drive file: {url}")
    print(f"     -> {dest_path}")
    dest_path.parent.mkdir(parents=True, exist_ok=True)
    output = gdown.download(url=url, output=str(dest_path), quiet=False, fuzzy=True)
    if not output or not dest_path.exists():
        raise RuntimeError(f"gdown failed to create expected file: {dest_path}")

    if _looks_like_html_file(dest_path):
        raise RuntimeError(
            "Google Drive download returned HTML instead of binary artifact bytes. "
            f"URL: {url}"
        )

    file_size = dest_path.stat().st_size
    print(f"Downloaded {dest_path.name} ({file_size:,} bytes)")


def download_direct(url: str, dest_path: Path) -> None:
    """Download a file directly from a URL."""
    if _is_google_drive_url(url):
        download_google_drive(url, dest_path)
        return

    print(f"Downloading {url}")
    print(f"     -> {dest_path}")

    dest_path.parent.mkdir(parents=True, exist_ok=True)

    try:
        urllib.request.urlretrieve(url, dest_path)
        if _looks_like_html_file(dest_path):
            raise RuntimeError(
                "Downloaded payload appears to be HTML instead of artifact bytes: "
                f"{dest_path}"
            )
        file_size = dest_path.stat().st_size
        print(f"Downloaded {dest_path.name} ({file_size:,} bytes)")
    except Exception as exc:
        raise RuntimeError(f"Failed to download {url}: {exc}") from exc


def download_huggingface(repo_id: str, filename: str, dest_path: Path, cache_dir: str = None) -> None:
    """Download a single file from Hugging Face Hub."""
    try:
        from huggingface_hub import hf_hub_download
    except ImportError as exc:
        raise ImportError(
            "huggingface_hub is not installed. "
            "Install it with: pip install huggingface-hub"
        ) from exc

    print(f"Downloading from Hugging Face: {repo_id}/{filename}")
    print(f"     -> {dest_path}")

    dest_path.parent.mkdir(parents=True, exist_ok=True)

    try:
        hf_path = hf_hub_download(repo_id=repo_id, filename=filename, cache_dir=cache_dir)
        shutil.copy2(hf_path, dest_path)
        if _looks_like_html_file(dest_path):
            raise RuntimeError(
                "Downloaded Hugging Face artifact appears to be HTML instead of bytes: "
                f"{dest_path}"
            )
        file_size = dest_path.stat().st_size
        print(f"Downloaded {dest_path.name} ({file_size:,} bytes)")
    except Exception as exc:
        raise RuntimeError(f"Failed to download from Hugging Face: {exc}") from exc


def download_huggingface_snapshot(
    repo_id: str,
    dest_dir: Path,
    cache_dir: str = None,
    revision: str = None,
    allow_patterns: Optional[List[str]] = None,
    ignore_patterns: Optional[List[str]] = None,
) -> None:
    """Download an entire Hugging Face repo snapshot into a local directory."""
    try:
        from huggingface_hub import snapshot_download
    except ImportError as exc:
        raise ImportError(
            "huggingface_hub is not installed. "
            "Install it with: pip install huggingface-hub"
        ) from exc

    print(f"Downloading Hugging Face snapshot: {repo_id}")
    print(f"     -> {dest_dir}")

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

    kwargs: Dict[str, Any] = {
        "repo_id": repo_id,
        "cache_dir": cache_dir,
        "local_dir": str(dest_dir),
        "local_dir_use_symlinks": False,
    }
    if revision:
        kwargs["revision"] = revision
    if allow_patterns:
        kwargs["allow_patterns"] = allow_patterns
    if ignore_patterns:
        kwargs["ignore_patterns"] = ignore_patterns

    try:
        snapshot_download(**kwargs)
        if not any(dest_dir.iterdir()):
            raise RuntimeError(f"Snapshot download completed but produced an empty directory: {dest_dir}")
        print(f"Downloaded snapshot to {dest_dir}")
    except Exception as exc:
        raise RuntimeError(f"Failed to download Hugging Face snapshot {repo_id}: {exc}") from exc


def load_download_config(config_path: Path) -> Dict[str, Any]:
    """Load download configuration from JSON."""
    with config_path.open("r", encoding="utf-8") as handle:
        return json.load(handle)


def normalize_file_entries(is_hf: int, files: List[Any]) -> List[Dict[str, Any]]:
    """Normalize and validate structured file entries.

    Required schema for each item:
    - {"path": "...", "filename": "..."}

    For direct downloads (isHF == 0), path is a URL.
    For HF downloads (isHF != 0):
    - file mode: path = "repo_id:filename"
    - snapshot mode: path = "repo_id" and set snapshot: 1 (or omit ':' in path)
    """
    del is_hf
    entries: List[Dict[str, Any]] = []

    for item in files:
        if isinstance(item, dict):
            source = item.get("path")
            filename = item.get("filename")
            if not source or not filename:
                raise ValueError(f"File entries must include 'path' and 'filename', got: {item}")

            entry: Dict[str, Any] = {
                "source": str(source),
                "filename": str(filename),
            }

            if "snapshot" in item:
                entry["snapshot"] = bool(item.get("snapshot"))
            if item.get("revision"):
                entry["revision"] = str(item.get("revision"))

            allow_patterns = _normalize_patterns(item.get("allow_patterns"))
            if allow_patterns:
                entry["allow_patterns"] = allow_patterns

            ignore_patterns = _normalize_patterns(item.get("ignore_patterns"))
            if ignore_patterns:
                entry["ignore_patterns"] = ignore_patterns

            entries.append(entry)
            continue

        raise ValueError(
            "Legacy string entries are no longer supported in 'files'. "
            "Use objects like {'path': '...', 'filename': '...'}"
        )

    return entries


def expected_destination_paths(is_hf: int, destination_base: Path, files: List[Any]) -> List[Path]:
    """Compute expected destination file paths from model_cache config."""
    entries = normalize_file_entries(is_hf, files)
    return [destination_base / entry["filename"] for entry in entries]


def extract_archives(destination_base: Path, archive_paths: List[Path], archive_format: str) -> None:
    """Extract archives into destination_base and keep idempotent markers."""
    fmt = (archive_format or "").strip().lower()
    if not fmt:
        return

    for archive_path in archive_paths:
        if not archive_path.exists():
            raise RuntimeError(f"Archive not found for extraction: {archive_path}")

        marker = destination_base / f".extract_complete_{archive_path.name}"
        if marker.exists():
            print(f"Extraction already completed for {archive_path.name}, skipping.")
            continue

        print(f"Extracting {archive_path.name} ({archive_format}) into {destination_base}")
        destination_base.mkdir(parents=True, exist_ok=True)

        if fmt == "zip":
            with zipfile.ZipFile(archive_path, "r") as zip_ref:
                zip_ref.extractall(destination_base)
        elif fmt in {"tar", "tar.gz", "tgz"}:
            mode = "r:gz" if fmt in {"tar.gz", "tgz"} else "r"
            with tarfile.open(archive_path, mode) as tar_ref:
                tar_ref.extractall(destination_base)
        else:
            raise ValueError(f"Unsupported archive format '{archive_format}' in model_cache.json")

        marker.write_text("ok\n", encoding="utf-8")
        print(f"Extracted {archive_path.name}")


def download_model_artifacts(
    config_path: Path,
    workspace_root: Path,
    env_vars: Dict[str, str] = None,
    skip_if_complete: bool = True,
) -> None:
    """Download all model artifacts specified in config."""
    import os

    config = load_download_config(config_path)
    is_hf = config.get("isHF", 0)
    archive_format = config.get("format", "")
    destination = config.get("destination", "")
    files = config.get("files", [])

    if not destination or not files:
        print("No files to download (empty config).")
        return

    if env_vars:
        for key, value in env_vars.items():
            if value.startswith("./") or value.startswith(".\\"):
                value = str(workspace_root / value[2:])
            os.environ[key] = value
            print(f"Set {key}={value}")

    dest_base = workspace_root / destination
    entries = normalize_file_entries(is_hf, files)
    expected_paths = [dest_base / entry["filename"] for entry in entries]
    all_downloads_present = bool(expected_paths) and all(
        _is_expected_destination_ready(path) for path in expected_paths
    )

    if skip_if_complete and all_downloads_present:
        print(f"All expected artifacts already exist in {dest_base}. Skipping download phase.")

    if not all_downloads_present:
        if is_hf == 0:
            for entry in entries:
                url = entry["source"]
                dest_path = dest_base / entry["filename"]

                if dest_path.exists():
                    if _looks_like_html_file(dest_path):
                        print(f"{dest_path} is an HTML placeholder. Re-downloading.")
                        dest_path.unlink()
                    else:
                        print(f"{dest_path} already exists, skipping.")
                        continue

                download_direct(url, dest_path)
        else:
            for entry in entries:
                file_spec = entry["source"]
                dest_path = dest_base / entry["filename"]
                hf_cache = os.environ.get("HF_HOME")

                snapshot_mode = bool(entry.get("snapshot")) or (":" not in file_spec)

                if snapshot_mode:
                    if ":" in file_spec:
                        raise ValueError(
                            "HF snapshot entries must use a repo_id without ':filename'. "
                            f"Got: {file_spec}"
                        )

                    if _is_expected_destination_ready(dest_path):
                        print(f"{dest_path} already exists, skipping snapshot download.")
                        continue

                    download_huggingface_snapshot(
                        repo_id=file_spec,
                        dest_dir=dest_path,
                        cache_dir=hf_cache,
                        revision=entry.get("revision"),
                        allow_patterns=entry.get("allow_patterns"),
                        ignore_patterns=entry.get("ignore_patterns"),
                    )
                    continue

                if ":" not in file_spec:
                    raise ValueError(f"HF file spec must be 'repo_id:filename', got: {file_spec}")

                repo_id, source_filename = file_spec.split(":", 1)

                if dest_path.exists():
                    if _looks_like_html_file(dest_path):
                        print(f"{dest_path} is an HTML placeholder. Re-downloading.")
                        dest_path.unlink()
                    else:
                        print(f"{dest_path} already exists, skipping.")
                        continue

                download_huggingface(repo_id, source_filename, dest_path, cache_dir=hf_cache)

    if archive_format:
        extract_archives(dest_base, expected_paths, archive_format)


if __name__ == "__main__":
    if len(sys.argv) > 1:
        config_file = Path(sys.argv[1])
        workspace = Path.cwd()
        download_model_artifacts(config_file, workspace)
    else:
        print("Usage: python downloads_manager.py <config_path>")
