"""Model download utilities for yakit."""

import hashlib
import os
import sys
from pathlib import Path
from typing import Optional
from urllib.request import urlretrieve
from urllib.error import URLError, HTTPError

# Default model settings
DEFAULT_MODEL_NAME = "yak_word2vec.bin"
DEFAULT_TRAINING_DATA_NAME = "train_pairs.txt"

# Hugging Face Hub settings
HF_REPO_ID = "loalkota/yakut-normalizer"
HF_MODEL_FILENAME = "yak_word2vec.bin"
HF_TRAINING_DATA_FILENAME = "train_pairs.txt"

# Alternative: direct URL download (GitHub releases, etc.)
# Set these if you want to use direct URL instead of HF Hub
DIRECT_MODEL_URL: Optional[str] = None
DIRECT_TRAINING_DATA_URL: Optional[str] = None

# Expected checksums (optional, for verification)
MODEL_SHA256: Optional[str] = None
TRAINING_DATA_SHA256: Optional[str] = None


def get_cache_dir() -> Path:
    """
    Get the cache directory for yakit models.
    
    Uses XDG_CACHE_HOME on Linux, ~/Library/Caches on macOS,
    and %LOCALAPPDATA% on Windows. Falls back to ~/.cache/yakit.
    
    Returns:
        Path to cache directory
    """
    if sys.platform == "win32":
        base = Path(os.environ.get("LOCALAPPDATA", Path.home() / "AppData" / "Local"))
    elif sys.platform == "darwin":
        base = Path.home() / "Library" / "Caches"
    else:
        base = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache"))
    
    cache_dir = base / "yakit"
    cache_dir.mkdir(parents=True, exist_ok=True)
    return cache_dir


def get_model_path(model_name: str = DEFAULT_MODEL_NAME) -> Path:
    """
    Get the path where a model should be stored.
    
    Args:
        model_name: Name of the model file
        
    Returns:
        Path to model file in cache directory
    """
    return get_cache_dir() / model_name


def _verify_checksum(file_path: Path, expected_sha256: str) -> bool:
    """Verify SHA256 checksum of a file."""
    sha256_hash = hashlib.sha256()
    with open(file_path, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            sha256_hash.update(chunk)
    return sha256_hash.hexdigest() == expected_sha256


def _download_with_progress(url: str, destination: Path, description: str = "Downloading") -> None:
    """Download a file with progress bar."""
    try:
        from tqdm import tqdm
        
        class DownloadProgressBar(tqdm):
            def update_to(self, b=1, bsize=1, tsize=None):
                if tsize is not None:
                    self.total = tsize
                self.update(b * bsize - self.n)
        
        with DownloadProgressBar(unit='B', unit_scale=True, miniters=1, desc=description) as t:
            urlretrieve(url, destination, reporthook=t.update_to)
    except ImportError:
        # tqdm not available, download without progress bar
        print(f"{description}...")
        urlretrieve(url, destination)
        print("Done!")


def download_from_huggingface(
    repo_id: str,
    filename: str,
    destination: Path,
    revision: str = "main",
) -> Path:
    """
    Download a file from Hugging Face Hub.
    
    Args:
        repo_id: Repository ID on Hugging Face (e.g., "yakit/yakut-normalizer")
        filename: Name of the file to download
        destination: Local path to save the file
        revision: Git revision (branch, tag, or commit hash)
        
    Returns:
        Path to downloaded file
        
    Raises:
        ImportError: If huggingface_hub is not installed
        Exception: If download fails
    """
    try:
        from huggingface_hub import hf_hub_download
    except ImportError:
        raise ImportError(
            "huggingface_hub is required to download models from Hugging Face. "
            "Install it with: pip install huggingface_hub"
        )
    
    print(f"Downloading {filename} from Hugging Face Hub ({repo_id})...")
    
    # Download to HF cache, then copy to our destination
    downloaded_path = hf_hub_download(
        repo_id=repo_id,
        filename=filename,
        revision=revision,
        local_dir=destination.parent,
        local_dir_use_symlinks=False,
    )
    
    # If the file was downloaded to a different location, move it
    downloaded_path = Path(downloaded_path)
    if downloaded_path != destination:
        import shutil
        shutil.move(str(downloaded_path), str(destination))
    
    return destination


def download_from_url(url: str, destination: Path, description: str = "Downloading") -> Path:
    """
    Download a file from a direct URL.
    
    Args:
        url: URL to download from
        destination: Local path to save the file
        description: Description for progress bar
        
    Returns:
        Path to downloaded file
        
    Raises:
        URLError: If download fails
    """
    destination.parent.mkdir(parents=True, exist_ok=True)
    
    try:
        _download_with_progress(url, destination, description)
    except (URLError, HTTPError) as e:
        raise RuntimeError(f"Failed to download from {url}: {e}") from e
    
    return destination


def download_model(
    model_name: str = DEFAULT_MODEL_NAME,
    force: bool = False,
    source: str = "auto",
) -> Path:
    """
    Download the Word2Vec model for Yakut language.
    
    Args:
        model_name: Name of the model file
        force: If True, re-download even if file exists
        source: Download source - "huggingface", "url", or "auto"
        
    Returns:
        Path to downloaded model file
        
    Raises:
        RuntimeError: If download fails from all sources
    """
    model_path = get_model_path(model_name)
    
    # Check if already downloaded
    if model_path.exists() and not force:
        # Optionally verify checksum
        if MODEL_SHA256 and not _verify_checksum(model_path, MODEL_SHA256):
            print(f"Model checksum mismatch, re-downloading...")
        else:
            return model_path
    
    model_path.parent.mkdir(parents=True, exist_ok=True)
    
    errors = []
    
    # Try Hugging Face first if source is auto or huggingface
    if source in ("auto", "huggingface"):
        try:
            return download_from_huggingface(
                repo_id=HF_REPO_ID,
                filename=HF_MODEL_FILENAME,
                destination=model_path,
            )
        except Exception as e:
            errors.append(f"Hugging Face: {e}")
            if source == "huggingface":
                raise RuntimeError(f"Failed to download model from Hugging Face: {e}") from e
    
    # Try direct URL if available
    if source in ("auto", "url") and DIRECT_MODEL_URL:
        try:
            return download_from_url(
                url=DIRECT_MODEL_URL,
                destination=model_path,
                description=f"Downloading {model_name}",
            )
        except Exception as e:
            errors.append(f"Direct URL: {e}")
            if source == "url":
                raise RuntimeError(f"Failed to download model from URL: {e}") from e
    
    # All sources failed
    error_msg = (
        f"Failed to download model '{model_name}'. Errors:\n" +
        "\n".join(f"  - {e}" for e in errors) +
        "\n\nPlease download the model manually and provide the path via word2vec_path parameter."
    )
    raise RuntimeError(error_msg)


def download_training_data(
    data_name: str = DEFAULT_TRAINING_DATA_NAME,
    force: bool = False,
    source: str = "auto",
) -> Path:
    """
    Download the training data for Yakut normalizer.
    
    Args:
        data_name: Name of the training data file
        force: If True, re-download even if file exists
        source: Download source - "huggingface", "url", or "auto"
        
    Returns:
        Path to downloaded training data file
        
    Raises:
        RuntimeError: If download fails from all sources
    """
    data_path = get_model_path(data_name)
    
    # Check if already downloaded
    if data_path.exists() and not force:
        if TRAINING_DATA_SHA256 and not _verify_checksum(data_path, TRAINING_DATA_SHA256):
            print(f"Training data checksum mismatch, re-downloading...")
        else:
            return data_path
    
    data_path.parent.mkdir(parents=True, exist_ok=True)
    
    errors = []
    
    # Try Hugging Face first
    if source in ("auto", "huggingface"):
        try:
            return download_from_huggingface(
                repo_id=HF_REPO_ID,
                filename=HF_TRAINING_DATA_FILENAME,
                destination=data_path,
            )
        except Exception as e:
            errors.append(f"Hugging Face: {e}")
            if source == "huggingface":
                raise RuntimeError(f"Failed to download training data from Hugging Face: {e}") from e
    
    # Try direct URL if available
    if source in ("auto", "url") and DIRECT_TRAINING_DATA_URL:
        try:
            return download_from_url(
                url=DIRECT_TRAINING_DATA_URL,
                destination=data_path,
                description=f"Downloading {data_name}",
            )
        except Exception as e:
            errors.append(f"Direct URL: {e}")
            if source == "url":
                raise RuntimeError(f"Failed to download training data from URL: {e}") from e
    
    error_msg = (
        f"Failed to download training data '{data_name}'. Errors:\n" +
        "\n".join(f"  - {e}" for e in errors) +
        "\n\nPlease download the data manually and provide the path via training_data_path parameter."
    )
    raise RuntimeError(error_msg)
