"""
OralGuard AI - Model Utilities for Hugging Face Spaces
Handles Google Drive download and on-demand model loading
"""

import torch
import torch.nn as nn
import timm
from efficientnet_pytorch import EfficientNet
from torchvision import models
from pathlib import Path
import gdown
import os
import streamlit as st

class OralGuardEfficientNet(nn.Module):
    """EfficientNet-B5 based model for oral cancer classification"""
    def __init__(self, num_classes=2, pretrained=False):
        super(OralGuardEfficientNet, self).__init__()
        
        if pretrained:
            self.backbone = EfficientNet.from_pretrained('efficientnet-b5')
        else:
            self.backbone = EfficientNet.from_name('efficientnet-b5')
        
        num_features = self.backbone._fc.in_features
        
        self.backbone._fc = nn.Sequential(
            nn.Dropout(0.3),
            nn.Linear(num_features, 512),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(512, num_classes)
        )
    
    def forward(self, x):
        return self.backbone(x)


class OralGuardVGG16(nn.Module):
    """VGG16 with Batch Normalization for oral cancer classification"""
    def __init__(self, num_classes=2, pretrained=False):
        super(OralGuardVGG16, self).__init__()
        
        self.backbone = models.vgg16_bn(pretrained=pretrained)
        num_features = self.backbone.classifier[6].in_features
        
        self.backbone.classifier[6] = nn.Sequential(
            nn.Dropout(0.3),
            nn.Linear(num_features, 512),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(512, num_classes)
        )
    
    def forward(self, x):
        return self.backbone(x)


class OralGuardViT(nn.Module):
    """Vision Transformer (ViT-B/16) for oral cancer classification"""
    def __init__(self, num_classes=2, pretrained=False):
        super(OralGuardViT, self).__init__()
        
        self.backbone = timm.create_model(
            'vit_base_patch16_224',
            pretrained=pretrained,
            num_classes=0
        )
        
        embed_dim = self.backbone.embed_dim
        
        self.classifier = nn.Sequential(
            nn.Linear(embed_dim, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(256, num_classes)
        )
    
    def forward(self, x):
        features = self.backbone(x)
        return self.classifier(features)


class GoogleDriveModelLoader:
    """
    Handles Google Drive model downloading and loading with caching
    """
    def __init__(self, cache_dir="./models_cache"):
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.loaded_models = {}
        self.class_names = ['Normal', 'OSCC']
        
        # Google Drive file IDs - will be set from Hugging Face secrets or config
        self.gdrive_file_ids = self._get_gdrive_ids()
    
    def _get_gdrive_ids(self):
        """
        Get Google Drive file IDs from Hugging Face secrets or environment
        Format: https://drive.google.com/file/d/FILE_ID/view?usp=sharing
        """
        # Try to get from Streamlit secrets first, fallback to env vars
        try:
            return {
                'efficientnet_100x': st.secrets.get("EFFICIENTNET_100X_ID", os.getenv("EFFICIENTNET_100X_ID")),
                'efficientnet_400x': st.secrets.get("EFFICIENTNET_400X_ID", os.getenv("EFFICIENTNET_400X_ID")),
                'vgg16_100x': st.secrets.get("VGG16_100X_ID", os.getenv("VGG16_100X_ID")),
                'vgg16_400x': st.secrets.get("VGG16_400X_ID", os.getenv("VGG16_400X_ID")),
                'vit_100x': st.secrets.get("VIT_100X_ID", os.getenv("VIT_100X_ID")),
                'vit_400x': st.secrets.get("VIT_400X_ID", os.getenv("VIT_400X_ID")),
            }
        except:
            # Return empty dict if secrets not configured
            return {}

    def download_model_from_gdrive(self, file_id, model_name, magnification):
        """
        Download model from Google Drive if not already cached
        """
        # Ensure cache directory exists
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        
        # Create filename
        filename = f"{model_name}_{magnification}.pth"
        filepath = self.cache_dir / filename
        
        # Check if already downloaded
        if filepath.exists():
            st.info(f"✅ Using cached model: {filename}")
            return filepath
        
        # Check file_id
        if not file_id:
            raise ValueError(f"Google Drive file ID not configured for {model_name}_{magnification}")
        
        # Show download info
        download_msg = st.empty()
        download_msg.info(f"⬇️ Downloading {filename} from Google Drive... This may take a few minutes.")
        
        try:
            # Construct download URL
            url = f"https://drive.google.com/uc?id={file_id}"
            
            # Create output path string
            output = str(filepath.absolute())
            
            # Download with gdown (quiet mode to reduce jitter)
            import gdown
            #gdown.download(url, output, quiet=True)
            # Download with gdown WITH output to see what's happening
            st.write(f"📥 Downloading from: {url}")
            st.write(f"📂 Saving to: {output}")
            st.write(f"📁 Cache dir: {self.cache_dir.absolute()}")
            
            # Download with verbose output
            import sys
            from io import StringIO
            old_stdout = sys.stdout
            sys.stdout = buffer = StringIO()
            
            try:
                result = gdown.download(url, output, quiet=False)
                download_output = buffer.getvalue()
                st.code(download_output)
                st.write(f"✅ gdown returned: {result}")
            finally:
                sys.stdout = old_stdout
            
            # NOW check what actually happened
            st.write(f"📂 Checking: {output}")
            st.write(f"  - Exists: {os.path.exists(output)}")
            if os.path.exists(output):
                st.write(f"  - Size: {os.path.getsize(output)} bytes")
                
            # Also check the cache directory
            st.write(f"📁 Cache directory contents:")
            for f in self.cache_dir.glob("*"):
                st.write(f"  - {f.name}: {f.stat().st_size} bytes")
            
            # Verify file was downloaded
            if not filepath.exists():
                raise FileNotFoundError(f"Download completed but file not found at {filepath}")
            
            download_msg.success(f"✅ Downloaded {filename} successfully!")
            
            # Clear message after 2 seconds
            import time
            time.sleep(2)
            download_msg.empty()
            
            return filepath
            
        except Exception as e:
            download_msg.error(f"❌ Error downloading {filename}: {str(e)}")
            # Clean up partial download
            if filepath.exists():
                filepath.unlink()
            raise
    
    def load_model(self, model_type, magnification):
        """
        Load a specific model (downloads if necessary)
        """
        model_key = f"{model_type}_{magnification}"
        
        # Return if already loaded
        if model_key in self.loaded_models:
            return self.loaded_models[model_key]
        
        # Get file ID
        file_id = self.gdrive_file_ids.get(model_key)
        
        # Download model if needed
        model_path = self.download_model_from_gdrive(file_id, model_type, magnification)
        
        # CREATE MODEL INSTANCE FIRST
        if model_type == 'efficientnet':
            model = OralGuardEfficientNet(num_classes=2, pretrained=False)
        elif model_type == 'vgg16':
            model = OralGuardVGG16(num_classes=2, pretrained=False)
        elif model_type == 'vit':
            model = OralGuardViT(num_classes=2, pretrained=False)
        else:
            raise ValueError(f"Unknown model type: {model_type}")
        
        # CRITICAL FIX: Convert Path to absolute string
        model_path_str = str(model_path.absolute())
        
        # Load weights
        st.info(f"🔄 Loading {model_key} into memory...")
        
        try:
            checkpoint = torch.load(model_path_str, map_location=self.device)
            
            # Handle different checkpoint formats
            if isinstance(checkpoint, dict):
                if 'model_state_dict' in checkpoint:
                    model.load_state_dict(checkpoint['model_state_dict'])
                elif 'state_dict' in checkpoint:
                    model.load_state_dict(checkpoint['state_dict'])
                else:
                    model.load_state_dict(checkpoint)
            else:
                model.load_state_dict(checkpoint)
            
            # Move to device and set to eval mode
            model = model.to(self.device)
            model.eval()
            
            # Cache the loaded model
            self.loaded_models[model_key] = model
            st.success(f"✅ {model_key} loaded successfully!")
            
            return model
            
        except Exception as e:
            st.error(f"❌ Failed to load model: {str(e)}")
            st.error(f"Model path: {model_path_str}")
            st.error(f"File exists: {os.path.exists(model_path_str)}")
            raise

        
    def get_model_info(self):
        """Get information about loaded models"""
        info = {}
        for key, model in self.loaded_models.items():
            total_params = sum(p.numel() for p in model.parameters())
            trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
            
            info[key] = {
                'total_parameters': f"{total_params:,}",
                'trainable_parameters': f"{trainable_params:,}",
                'device': str(next(model.parameters()).device),
                'loaded': True
            }
        
        return info
    
    def clear_cache(self):
        """Clear loaded models from memory"""
        self.loaded_models.clear()
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
        st.success("✅ Model cache cleared!")


def get_model_display_name(model_type):
    """Get display name for model type"""
    names = {
        'efficientnet': 'EfficientNet-B5',
        'vgg16': 'VGG16-BN',
        'vit': 'ViT-B/16',
        'ensemble': 'Ensemble (All Models)'
    }
    return names.get(model_type, model_type)


def get_magnification_display_name(mag):
    """Get display name for magnification"""
    names = {
        '100x': '100× Magnification',
        '400x': '400× Magnification'
    }
    return names.get(mag, mag)


def get_model_size_info():
    """Get expected model sizes"""
    return {
        'efficientnet': '340 MB per model',
        'vgg16': '1.6 GB per model',
        'vit': '1.0 GB per model',
        'total_per_magnification': '~2.94 GB',
        'total_all': '~5.88 GB'
    }