from typing import Optional, Union, Tuple
from pathlib import Path

from huggingface_hub import snapshot_download

from emissary_ml_toolkit import get_logger
from emissary_ml_toolkit.storage.manager import StorageManager
from emissary_ml_toolkit.config import BaseModelConfig, FoundationModel
from emissary_ml_toolkit.config.storage import StorageConfig

logger = get_logger(__name__)

class ModelDownloader:
    """Handles downloading models from HuggingFace or S3"""

    def __init__(self, storage_manager: Optional[StorageManager] = None):
        self._storage_manager = storage_manager

    def download(
            self,
            base_model_config: BaseModelConfig,
            local_dir: Union[str, Path] = "./base_model"
    ) -> Tuple[str, Optional[str]]:
        """
        Download foundation model and optional LoRA adapter.

        Returns:
            Tuple of (model_path, lora_adapter_path)
        """
        local_dir = Path(local_dir)

        # Download foundation model
        model_path = self._download_foundation_model(
            base_model_config.foundation_model,
            local_dir
        )

        # Download LoRA adapter if present in configuration
        lora_path = None
        if base_model_config.lora_adapter:
            lora_path = self._download_lora_adapter(
                base_model_config.lora_adapter,
                local_dir / "lora"
            )

        return str(model_path), lora_path

    def _download_foundation_model(self, foundation_model: FoundationModel, local_dir: Path) -> Path:
        """Download foundation model from configured source."""
        logger.info(f"Downloading foundation model from {foundation_model.source}")

        if foundation_model.source == "huggingface":
            repo_id = foundation_model.huggingface.repo_id
            logger.info(f"Downloading from HuggingFace: {repo_id}")
            snapshot_download(repo_id=repo_id)
            return Path(repo_id)

        elif foundation_model.source == "storage":
            if not self._storage_manager:
                raise ValueError(
                    "Storage manager required for downloading from storage source."
                    "Please configure storage_config in configuration file."
                )

            local_dir.mkdir(parents=True, exist_ok=True)
            storage = self._storage_manager.get_storage(
                foundation_model.storage.storage_ref
            )
            storage.download_directory(
                prefix=foundation_model.storage.storage_path,
                local_dir=local_dir
            )
            return local_dir

        else:
            raise ValueError(f"Unsupported model source: {foundation_model.source}")

    def _download_lora_adapter(self, lora_adapter: StorageConfig, local_dir: Path) -> str:
        """Download LoRA adapter from storage."""
        if not self._storage_manager:
            raise ValueError(
                "Storage manager required for downloading from storage source."
                "Please configure storage_config in configuration file."
            )

        logger.info(f"Downloading LoRA adapter to {local_dir}")
        local_dir.mkdir(parents=True, exist_ok=True)

        storage = self._storage_manager.get_storage(lora_adapter.storage_ref)
        storage.download_file(
            storage_path=lora_adapter.storage_path,
            local_path=str(local_dir)
        )

        return str(local_dir)
