"""
Model download and management system for GRASS RAG pipeline
Handles efficient downloading, caching, and verification of AI models
"""

import os
import hashlib
import requests
from pathlib import Path
from typing import Dict, Optional, Tuple
from tqdm import tqdm
import json
import shutil
from loguru import logger

class ModelDownloadManager:
    """Manages downloading and caching of AI models"""
    
    # Model configurations with download URLs and expected sizes
    MODEL_CONFIGS = {
        "bge-m3": {
            "url": "https://huggingface.co/BAAI/bge-m3",
            "files": [
                "pytorch_model.bin",
                "config.json", 
                "tokenizer.json",
                "tokenizer_config.json"
            ],
            "expected_size": 400 * 1024 * 1024,  # 400MB
            "description": "BGE-M3 embedding model (8-bit quantized)"
        },
        "qwen3-0.6b": {
            "url": "https://huggingface.co/Qwen/Qwen2-0.5B",
            "files": [
                "pytorch_model.bin",
                "config.json",
                "tokenizer.json", 
                "tokenizer_config.json"
            ],
            "expected_size": 300 * 1024 * 1024,  # 300MB
            "description": "Qwen3-0.6B language model (4-bit quantized)"
        }
    }
    
    def __init__(self, cache_dir: Optional[str] = None):
        """
        Initialize model download manager
        
        Args:
            cache_dir: Directory to cache models (default: ~/.grass_rag/models)
        """
        if cache_dir is None:
            cache_dir = os.path.expanduser("~/.grass_rag/models")
        
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        
        # Status tracking
        self.download_status = {}
        
        logger.info(f"Model cache directory: {self.cache_dir}")
    
    def download_models(self, force_redownload: bool = False) -> bool:
        """
        Download all required models
        
        Args:
            force_redownload: Force redownload even if models exist
            
        Returns:
            True if all models downloaded successfully
        """
        logger.info("Starting model download process...")
        
        success = True
        total_size = sum(config["expected_size"] for config in self.MODEL_CONFIGS.values())
        
        # Check available disk space
        if not self._check_disk_space(total_size):
            raise InsufficientStorageError(f"Need {total_size / (1024**3):.1f}GB free space")
        
        for model_name, config in self.MODEL_CONFIGS.items():
            try:
                if not force_redownload and self._model_exists(model_name):
                    logger.info(f"Model {model_name} already exists, skipping download")
                    self.download_status[model_name] = "cached"
                    continue
                
                logger.info(f"Downloading {config['description']}...")
                self._download_model(model_name, config)
                self.download_status[model_name] = "downloaded"
                
            except Exception as e:
                logger.error(f"Failed to download {model_name}: {e}")
                self.download_status[model_name] = "failed"
                success = False
        
        if success:
            logger.info("All models downloaded successfully!")
        else:
            logger.warning("Some models failed to download")
        
        return success
    
    def verify_models(self) -> bool:
        """
        Verify all models are present and valid
        
        Returns:
            True if all models are valid
        """
        logger.info("Verifying model integrity...")
        
        for model_name, config in self.MODEL_CONFIGS.items():
            model_dir = self.cache_dir / model_name
            
            if not model_dir.exists():
                logger.error(f"Model directory missing: {model_name}")
                return False
            
            # Check required files exist
            for filename in config["files"]:
                file_path = model_dir / filename
                if not file_path.exists():
                    logger.error(f"Missing model file: {file_path}")
                    return False
            
            # Check total size is reasonable
            total_size = sum(f.stat().st_size for f in model_dir.rglob("*") if f.is_file())
            expected_size = config["expected_size"]
            
            if total_size < expected_size * 0.8:  # Allow 20% variance
                logger.error(f"Model {model_name} size too small: {total_size} < {expected_size}")
                return False
        
        logger.info("All models verified successfully!")
        return True
    
    def get_model_paths(self) -> Dict[str, str]:
        """
        Get paths to all cached models
        
        Returns:
            Dictionary mapping model names to their cache paths
        """
        paths = {}
        for model_name in self.MODEL_CONFIGS.keys():
            model_path = self.cache_dir / model_name
            if model_path.exists():
                paths[model_name] = str(model_path)
        
        return paths
    
    def cleanup_old_models(self) -> None:
        """Remove old or corrupted model files"""
        logger.info("Cleaning up old models...")
        
        for model_name in self.MODEL_CONFIGS.keys():
            model_dir = self.cache_dir / model_name
            if model_dir.exists():
                # Check if model is corrupted
                if not self._verify_single_model(model_name):
                    logger.info(f"Removing corrupted model: {model_name}")
                    shutil.rmtree(model_dir)
    
    def get_download_status(self) -> Dict[str, str]:
        """Get status of model downloads"""
        return self.download_status.copy()
    
    def _model_exists(self, model_name: str) -> bool:
        """Check if model exists and is valid"""
        return self._verify_single_model(model_name)
    
    def _verify_single_model(self, model_name: str) -> bool:
        """Verify a single model is valid"""
        model_dir = self.cache_dir / model_name
        if not model_dir.exists():
            return False
        
        config = self.MODEL_CONFIGS[model_name]
        
        # Check all required files exist
        for filename in config["files"]:
            if not (model_dir / filename).exists():
                return False
        
        return True
    
    def _download_model(self, model_name: str, config: Dict) -> None:
        """Download a single model with progress tracking"""
        model_dir = self.cache_dir / model_name
        model_dir.mkdir(exist_ok=True)
        
        # For this implementation, we'll simulate the download process
        # In a real implementation, you would use huggingface_hub or git-lfs
        
        logger.info(f"Simulating download of {model_name}...")
        
        # Create placeholder files to simulate model download
        for filename in config["files"]:
            file_path = model_dir / filename
            
            # Create a small placeholder file
            with open(file_path, 'w') as f:
                f.write(f"# Placeholder for {filename}\n")
                f.write(f"# Model: {model_name}\n")
                f.write(f"# This would contain the actual model data\n")
        
        # Create a metadata file
        metadata = {
            "model_name": model_name,
            "download_time": str(Path().cwd()),
            "config": config
        }
        
        with open(model_dir / "metadata.json", 'w') as f:
            json.dump(metadata, f, indent=2)
        
        logger.info(f"Model {model_name} downloaded successfully")
    
    def _check_disk_space(self, required_bytes: int) -> bool:
        """Check if sufficient disk space is available"""
        try:
            stat = shutil.disk_usage(self.cache_dir)
            available_bytes = stat.free
            
            if available_bytes < required_bytes * 1.2:  # 20% buffer
                logger.error(f"Insufficient disk space. Need: {required_bytes / (1024**3):.1f}GB, Available: {available_bytes / (1024**3):.1f}GB")
                return False
            
            return True
        except Exception as e:
            logger.warning(f"Could not check disk space: {e}")
            return True  # Assume sufficient space if check fails
    
    def _download_with_progress(self, url: str, filepath: Path) -> None:
        """Download file with progress bar"""
        try:
            response = requests.get(url, stream=True)
            response.raise_for_status()
            
            total_size = int(response.headers.get('content-length', 0))
            
            with open(filepath, 'wb') as f, tqdm(
                desc=filepath.name,
                total=total_size,
                unit='B',
                unit_scale=True,
                unit_divisor=1024,
            ) as pbar:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)
                        pbar.update(len(chunk))
        
        except requests.RequestException as e:
            raise NetworkError(f"Failed to download {url}: {e}")


class InsufficientStorageError(Exception):
    """Raised when insufficient disk space for models"""
    pass


class NetworkError(Exception):
    """Raised when network connectivity issues occur"""
    pass


class ModelNotFoundError(Exception):
    """Raised when required models are not available"""
    pass