"""
Kaggle Model Upload Service

Handles uploading trained models to Kaggle Model Hub with model cards.
Uses kagglehub Python library for modern, reliable uploads.
"""

import os
import json
import logging
from pathlib import Path
from typing import Dict, Any, Optional, List
from datetime import datetime

logger = logging.getLogger(__name__)


class KaggleModelUploader:
    """Upload trained models to Kaggle Model Hub"""
    
    def __init__(self, username: str, api_key: str):
        """
        Initialize Kaggle uploader
        
        Args:
            username: Kaggle username
            api_key: Kaggle API key
        """
        self.username = username
        self.api_key = api_key
        
        # Set Kaggle credentials as environment variables
        os.environ['KAGGLE_USERNAME'] = username
        os.environ['KAGGLE_KEY'] = api_key
    
    def create_model_card(
        self,
        model_slug: str,
        model_variant: str,
        training_config: Dict[str, Any],
        metrics: Optional[Dict[str, Any]] = None,
        dataset_info: Optional[Dict[str, Any]] = None
    ) -> str:
        """
        Generate a model card in Markdown format
        
        Args:
            model_slug: Model identifier (e.g., "my-yolo-model")
            model_variant: YOLO variant (e.g., "yolo11n")
            training_config: Training configuration dict
            metrics: Training metrics (mAP, precision, recall, etc.)
            dataset_info: Dataset information
            
        Returns:
            Model card content as Markdown string
        """
        # Extract key training parameters
        epochs = training_config.get('epochs', 'N/A')
        batch_size = training_config.get('batch_size', 'N/A')
        img_size = training_config.get('img_size', 640)
        learning_rate = training_config.get('learning_rate', 0.01)
        
        # Format metrics if available
        metrics_section = ""
        if metrics:
            metrics_section = f"""
## 📊 Performance Metrics

| Metric | Value |
|--------|-------|
| mAP50 | {metrics.get('mAP50', 'N/A')} |
| mAP50-95 | {metrics.get('mAP50_95', 'N/A')} |
| Precision | {metrics.get('precision', 'N/A')} |
| Recall | {metrics.get('recall', 'N/A')} |
"""
        
        # Format dataset info if available
        dataset_section = ""
        if dataset_info:
            dataset_section = f"""
## 📦 Dataset Information

- **Name**: {dataset_info.get('name', 'Custom Dataset')}
- **Classes**: {dataset_info.get('num_classes', 'N/A')}
- **Training Images**: {dataset_info.get('train_images', 'N/A')}
- **Validation Images**: {dataset_info.get('val_images', 'N/A')}
"""
        
        # Generate model card
        model_card = f"""# {model_slug}

## 🎯 Model Overview

This is a **{model_variant}** model trained using the Aegis AI platform. The model is optimized for object detection tasks and can be used with the Ultralytics YOLO framework.

## 🚀 Quick Start

```python
from ultralytics import YOLO

# Load the model
model = YOLO('path/to/best.pt')

# Run inference
results = model('image.jpg')

# Display results
results.show()
```

## ⚙️ Training Configuration

| Parameter | Value |
|-----------|-------|
| Model Variant | {model_variant} |
| Epochs | {epochs} |
| Batch Size | {batch_size} |
| Image Size | {img_size}x{img_size} |
| Learning Rate | {learning_rate} |
| Optimizer | {training_config.get('optimizer', 'auto')} |
| Device | {training_config.get('device', 'auto')} |
{metrics_section}{dataset_section}
## 📝 Training Details

- **Framework**: Ultralytics YOLO
- **Training Platform**: Aegis AI
- **Training Date**: {datetime.utcnow().strftime('%Y-%m-%d')}
- **Augmentation**: Enabled
- **Early Stopping**: Patience {training_config.get('early_stopping', {}).get('patience', 50)}

## 📥 Model Files

- `best.pt` - Best model checkpoint
- `last.pt` - Last epoch checkpoint
- `best.onnx` - ONNX export (if available)
- Additional export formats may be included

## 🔧 Usage with Aegis AI

This model was trained using [Aegis AI](https://github.com/your-repo/aegis-ai), an advanced AI-powered security and monitoring platform.

## 📄 License

Please refer to the model's license file for usage terms and conditions.

## 🤝 Citation

If you use this model in your research or project, please cite:

```bibtex
@misc{{{model_slug},
  author = {{{self.username}}},
  title = {{{model_slug}}},
  year = {{{datetime.utcnow().year}}},
  publisher = {{Kaggle}},
  howpublished = {{\\url{{https://www.kaggle.com/models/{self.username}/{model_slug}}}}}
}}
```

---

*Generated by Aegis AI - Automated Model Training & Deployment*
"""
        return model_card
    
    def create_model_metadata(
        self,
        model_slug: str,
        model_variant: str,
        training_config: Dict[str, Any],
        framework: str = "PyTorch"
    ) -> Dict[str, Any]:
        """
        Create Kaggle model metadata
        
        Args:
            model_slug: Model identifier
            model_variant: YOLO variant
            training_config: Training configuration
            framework: Model framework (default: PyTorch)
            
        Returns:
            Metadata dictionary for Kaggle API
        """
        metadata = {
            "title": f"{model_slug} - {model_variant}",
            "slug": model_slug,
            "subtitle": f"YOLO {model_variant} trained with Aegis AI",
            "isPrivate": False,
            "description": f"Object detection model trained using {model_variant}",
            "publishDate": datetime.utcnow().isoformat(),
            "framework": framework,
            "instanceType": "GPU" if training_config.get('requiresGPU', False) else "CPU",
            "licenseName": "Apache 2.0",
            "modelInstanceSlug": "1",  # Version number
            "files": []  # Will be populated during upload
        }
        return metadata
    
    def upload_model(
        self,
        model_dir: Path,
        model_slug: str,
        model_variant: str,
        training_config: Dict[str, Any],
        metrics: Optional[Dict[str, Any]] = None,
        dataset_info: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """
        Upload trained model to Kaggle Model Hub using kagglehub library
        
        Args:
            model_dir: Directory containing model files
            model_slug: Model identifier (e.g., "my-yolo-model")
            model_variant: YOLO variant (e.g., "yolo11n")
            training_config: Training configuration dict
            metrics: Training metrics
            dataset_info: Dataset information
            
        Returns:
            Upload result dictionary
        """
        try:
            logger.info(f"📤 Preparing to upload model to Kaggle: {model_slug}")
            
            # Validate model directory
            if not model_dir.exists():
                raise ValueError(f"Model directory not found: {model_dir}")
            
            # Check for required model files
            best_model = model_dir / "best.pt"
            if not best_model.exists():
                raise ValueError(f"best.pt not found in {model_dir}")
            
            # Create model card
            logger.info("📝 Generating model card...")
            model_card = self.create_model_card(
                model_slug=model_slug,
                model_variant=model_variant,
                training_config=training_config,
                metrics=metrics,
                dataset_info=dataset_info
            )
            
            # Save model card as README.md (required by Kaggle)
            model_card_path = model_dir / "README.md"
            with open(model_card_path, 'w') as f:
                f.write(model_card)
            logger.info(f"✅ Model card saved to {model_card_path}")
            
            # ✅ Filter to include ONLY model files (exclude batch images and metrics visualizations)
            # Model file extensions: .pt (PyTorch), .torchscript, .onnx, .mlpackage, .mlmodel, .engine (TensorRT), 
            # .tflite, .pb (TensorFlow), .xml/.bin (OpenVINO), .md (README), .txt/.yaml (config)
            ALLOWED_EXTENSIONS = {
                '.pt',  # PyTorch
                '.pth',  # PyTorch alternative
                '.torchscript',  # TorchScript (intermediate format for CoreML/OpenVINO)
                '.onnx',  # ONNX
                '.engine',  # TensorRT
                '.tflite',  # TensorFlow Lite
                '.pb',  # TensorFlow
                '.xml',  # OpenVINO
                '.bin',  # OpenVINO weights
                '.md',  # Model card
                '.txt',  # Labels/config files
                '.yaml',  # Config files
                '.yml',  # Config files
            }
            
            ALLOWED_DIRS = {
                '.mlpackage',  # CoreML package (directory)
            }
            
            # File patterns to EXPLICITLY EXCLUDE (no images, no visualizations)
            EXCLUDE_PATTERNS = [
                'train_batch',  # Training batch images
                'val_batch',  # Validation batch images  
                'confusion_matrix',  # Confusion matrix PNG
                'F1_curve',  # F1 curve PNG
                'P_curve',  # Precision curve PNG
                'R_curve',  # Recall curve PNG
                'PR_curve',  # Precision-Recall curve PNG
                'labels',  # Labels correlation PNG (labels.jpg or labels_correlogram.jpg)
                'results.png',  # Results visualization
                'results.csv',  # Results CSV (available on Wandb)
            ]
            
            # Excluded file extensions (images and visualizations)
            EXCLUDE_EXTENSIONS = {
                '.jpg',
                '.jpeg',
                '.png',
                '.gif',
                '.bmp',
                '.csv',  # CSVs are on Wandb
            }
            
            # List model files to upload (filtered)
            model_files = []
            for file_path in model_dir.iterdir():
                if file_path.name.startswith('.'):
                    continue  # Skip hidden files
                
                # Check if filename matches exclusion pattern
                should_exclude = False
                for pattern in EXCLUDE_PATTERNS:
                    if pattern in file_path.name:
                        logger.debug(f"🚫 Excluding {file_path.name} (matches pattern: {pattern})")
                        should_exclude = True
                        break
                
                if should_exclude:
                    continue
                
                # Check if extension is explicitly excluded (images, CSVs)
                if file_path.is_file() and file_path.suffix.lower() in EXCLUDE_EXTENSIONS:
                    logger.debug(f"🚫 Excluding {file_path.name} (image/visualization file)")
                    continue
                
                # Include directories only if they're allowed (e.g., .mlpackage)
                if file_path.is_dir():
                    if file_path.suffix in ALLOWED_DIRS or file_path.name in ALLOWED_DIRS:
                        model_files.append(file_path.name)
                        logger.info(f"📦 Including directory: {file_path.name}")
                    else:
                        logger.debug(f"🚫 Excluding directory: {file_path.name}")
                # Include files only if extension is allowed
                elif file_path.is_file():
                    if file_path.suffix in ALLOWED_EXTENSIONS:
                        model_files.append(file_path.name)
                        logger.info(f"📦 Including file: {file_path.name}")
                    else:
                        logger.debug(f"🚫 Excluding file: {file_path.name} (extension not allowed)")
            
            if not model_files:
                raise ValueError(f"No model files found to upload in {model_dir}")
            
            logger.info(f"📦 Uploading {len(model_files)} model files: {', '.join(model_files)}")
            logger.info(f"ℹ️  Excluded: batch images, metrics visualizations (available on Wandb)")
            
            # Create a temporary directory with ONLY the allowed files
            # kagglehub.model_upload() uploads the entire directory recursively, so we need to filter
            import tempfile
            import shutil
            
            with tempfile.TemporaryDirectory() as temp_dir:
                temp_path = Path(temp_dir)
                logger.info(f"📁 Creating filtered upload directory: {temp_path}")
                
                # Copy only allowed files to temp directory
                for file_name in model_files:
                    src = model_dir / file_name
                    dst = temp_path / file_name
                    
                    if src.is_dir():
                        # Copy directory recursively (e.g., .mlpackage)
                        shutil.copytree(src, dst)
                        logger.debug(f"  📂 Copied directory: {file_name}")
                    else:
                        # Copy file
                        shutil.copy2(src, dst)
                        logger.debug(f"  📄 Copied file: {file_name}")
                
                logger.info(f"✅ Filtered directory created with {len(model_files)} items")
                
                # Upload using kagglehub library (modern API)
                logger.info(f"🚀 Uploading model to Kaggle using kagglehub...")
                logger.info(f"   Username: {self.username}")
                logger.info(f"   Model: {model_slug}")
                logger.info(f"   Variant: {model_variant}")
                
                try:
                    import kagglehub
                    
                    # Construct model handle in format: username/model/framework/variation
                    # For YOLO models, we use PyTorch as framework and model_variant as variation
                    model_handle = f"{self.username}/{model_slug}/PyTorch/{model_variant}"
                    
                    logger.info(f"📍 Model handle: {model_handle}")
                    logger.info(f"📂 Uploading from: {temp_path}")
                    
                    # Upload model using kagglehub
                    # This will create a new variation or new version of existing variation
                    version_notes = f"Trained with Aegis AI - {datetime.utcnow().strftime('%Y-%m-%d %H:%M UTC')}"
                    if metrics:
                        version_notes += f" | mAP50: {metrics.get('mAP50', 'N/A')}"
                    
                    kagglehub.model_upload(
                        handle=model_handle,
                        local_model_dir=str(temp_path),  # Use filtered temp directory
                        version_notes=version_notes,
                        license_name='Apache 2.0'
                    )
                    
                    model_url = f"https://www.kaggle.com/models/{self.username}/{model_slug}/PyTorch/{model_variant}"
                    
                    logger.info("✅ Model uploaded successfully to Kaggle!")
                    logger.info(f"🔗 Model URL: {model_url}")
                    
                    return {
                        "success": True,
                        "model_url": model_url,
                        "model_handle": model_handle,
                        "files_uploaded": len(model_files),
                        "message": "Model uploaded successfully to Kaggle Model Hub"
                    }
                    
                except ImportError as ie:
                    logger.error(f"❌ kagglehub package not installed: {ie}")
                    return {
                        "success": False,
                        "error": "kagglehub package not installed. Install with: pip install kagglehub",
                        "message": "kagglehub package required for upload"
                    }
                except Exception as upload_error:
                    error_msg = str(upload_error)
                    logger.error(f"❌ Kaggle upload failed: {error_msg}")
                    
                    # Provide more helpful error messages
                    if "404" in error_msg or "not found" in error_msg.lower():
                        logger.error("💡 Model may not exist yet. Creating models via API requires the model to be created on Kaggle.com first.")
                        logger.error("💡 Please create the model manually at: https://www.kaggle.com/models/new")
                        logger.error(f"💡 Then use model slug: {model_slug}")
                    
                    return {
                        "success": False,
                        "error": error_msg,
                        "message": "Failed to upload model to Kaggle",
                        "help": "Ensure the model exists on Kaggle.com first, or create it at https://www.kaggle.com/models/new"
                    }
            
        except Exception as e:
            logger.error(f"❌ Failed to upload model: {e}")
            import traceback
            logger.error(f"Traceback: {traceback.format_exc()}")
            return {
                "success": False,
                "error": str(e),
                "message": f"Upload failed: {str(e)}"
            }
    
    def list_models(self) -> List[Dict[str, Any]]:
        """
        List all models for the authenticated user using kagglehub
        
        Returns:
            List of model dictionaries
        """
        try:
            import kagglehub
            
            # Note: kagglehub doesn't have a direct list models API yet
            # This is a placeholder for future functionality
            logger.warning("Model listing not yet implemented in kagglehub")
            logger.info(f"💡 View your models at: https://www.kaggle.com/{self.username}/models")
            return []
                
        except Exception as e:
            logger.error(f"Failed to list models: {e}")
            return []


def upload_trained_model(
    model_dir: str,
    model_slug: str,
    model_variant: str,
    training_config: Dict[str, Any],
    kaggle_username: str,
    kaggle_api_key: str,
    metrics: Optional[Dict[str, Any]] = None,
    dataset_info: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
    """
    Convenience function to upload a trained model to Kaggle
    
    Args:
        model_dir: Path to directory containing model files
        model_slug: Model identifier
        model_variant: YOLO variant
        training_config: Training configuration
        kaggle_username: Kaggle username
        kaggle_api_key: Kaggle API key
        metrics: Training metrics (optional)
        dataset_info: Dataset information (optional)
        
    Returns:
        Upload result dictionary
    """
    uploader = KaggleModelUploader(
        username=kaggle_username,
        api_key=kaggle_api_key
    )
    
    return uploader.upload_model(
        model_dir=Path(model_dir),
        model_slug=model_slug,
        model_variant=model_variant,
        training_config=training_config,
        metrics=metrics,
        dataset_info=dataset_info
    )


