#!/usr/bin/env python3
"""
Deployment Helper Script for Power Systems Mini-Consultant
Automates the setup and deployment process for Hugging Face Spaces
"""

import os
import json
import shutil
import subprocess
import sys
from pathlib import Path
from typing import Dict, List, Optional

class PowerSystemsDeployer:
    """
    Automated deployment helper for Power Systems Mini-Consultant
    """
    
    def __init__(self):
        self.project_root = Path.cwd()
        self.required_files = [
            'app.py',
            'requirements.txt',
            'data/knowledge_base.json',
            'utils/__init__.py',
            'utils/rag_system.py',
            'utils/diagram_generator.py'
        ]
        self.optional_files = [
            'README.md',
            '.gitignore',
            '.env.example'
        ]
        
    def check_requirements(self) -> bool:
        """Check if all required files exist"""
        print("🔍 Checking project requirements...")
        
        missing_files = []
        for file_path in self.required_files:
            if not (self.project_root / file_path).exists():
                missing_files.append(file_path)
        
        if missing_files:
            print("❌ Missing required files:")
            for file in missing_files:
                print(f"   - {file}")
            return False
        
        print("✅ All required files found!")
        return True
    
    def validate_knowledge_base(self) -> bool:
        """Validate the knowledge base JSON structure"""
        print("🔍 Validating knowledge base...")
        
        try:
            kb_path = self.project_root / 'data' / 'knowledge_base.json'
            with open(kb_path, 'r', encoding='utf-8') as f:
                kb_data = json.load(f)
            
            # Check for essential sections
            required_sections = ['faults', 'protection', 'standards', 'formulas']
            missing_sections = [s for s in required_sections if s not in kb_data]
            
            if missing_sections:
                print(f"⚠️  Missing knowledge base sections: {missing_sections}")
                print("   The app will still work but with limited knowledge.")
            else:
                print("✅ Knowledge base structure validated!")
            
            return True
            
        except json.JSONDecodeError as e:
            print(f"❌ Invalid JSON in knowledge base: {e}")
            return False
        except Exception as e:
            print(f"❌ Error validating knowledge base: {e}")
            return False
    
    def setup_environment(self) -> bool:
        """Set up the development environment"""
        print("🔧 Setting up environment...")
        
        try:
            # Create virtual environment if it doesn't exist
            venv_path = self.project_root / 'venv'
            if not venv_path.exists():
                print("Creating virtual environment...")
                subprocess.run([sys.executable, '-m', 'venv', 'venv'], check=True)
            
            # Determine the correct pip path
            if os.name == 'nt':  # Windows
                pip_path = venv_path / 'Scripts' / 'pip'
            else:  # Unix/Linux/MacOS
                pip_path = venv_path / 'bin' / 'pip'
            
            # Install requirements
            print("Installing requirements...")
            subprocess.run([str(pip_path), 'install', '-r', 'requirements.txt'], check=True)
            
            print("✅ Environment setup complete!")
            return True
            
        except subprocess.CalledProcessError as e:
            print(f"❌ Error setting up environment: {e}")
            return False
        except Exception as e:
            print(f"❌ Unexpected error: {e}")
            return False
    
    def test_local_deployment(self) -> bool:
        """Test the application locally"""
        print("🧪 Testing local deployment...")
        
        try:
            # Check if .env file exists
            env_path = self.project_root / '.env'
            if not env_path.exists():
                print("⚠️  No .env file found. Creating template...")
                self.create_env_template()
                print("   Please add your GROQ_API_KEY to .env file and run again.")
                return False
            
            # Try importing the main application
            import importlib.util
            spec = importlib.util.spec_from_file_location("app", self.project_root / "app.py")
            app_module = importlib.util.module_from_spec(spec)
            
            # This will test if the imports work
            spec.loader.exec_module(app_module)
            
            print("✅ Local deployment test passed!")
            return True
            
        except ImportError as e:
            print(f"❌ Import error: {e}")
            print("   Please check your dependencies in requirements.txt")
            return False
        except Exception as e:
            print(f"❌ Error testing deployment: {e}")
            return False
    
    def create_env_template(self):
        """Create .env template file"""
        env_template = """# Environment Variables for Power Systems Mini-Consultant

# Required: Groq API Key (Get from https://console.groq.com)
GROQ_API_KEY=your_groq_api_key_here

# Optional: Debug settings
DEBUG=False
LOG_LEVEL=INFO

# Optional: Model settings
MODEL_NAME=mixtral-8x7b-32768
MAX_TOKENS=2000
TEMPERATURE=0.7
"""
        
        with open(self.project_root / '.env', 'w') as f:
            f.write(env_template)
    
    def prepare_huggingface_deployment(self) -> bool:
        """Prepare files for Hugging Face Spaces deployment"""
        print("🚀 Preparing Hugging Face deployment...")
        
        try:
            # Create deployment directory
            deploy_dir = self.project_root / 'hf_deployment'
            if deploy_dir.exists():
                shutil.rmtree(deploy_dir)
            deploy_dir.mkdir()
            
            # Copy required files
            for file_path in self.required_files:
                src = self.project_root / file_path
                dst = deploy_dir / file_path
                
                # Create parent directories if needed
                dst.parent.mkdir(parents=True, exist_ok=True)
                shutil.copy2(src, dst)
            
            # Copy optional files if they exist
            for file_path in self.optional_files:
                src = self.project_root / file_path
                if src.exists():
                    dst = deploy_dir / file_path
                    shutil.copy2(src, dst)
            
            # Create HuggingFace README
            self.create_hf_readme(deploy_dir)
            
            print(f"✅ Deployment files prepared in: {deploy_dir}")
            print("\n📋 Next steps for Hugging Face deployment:")
            print("   1. Create a new Space on Hugging Face")
            print("   2. Choose 'Gradio' as the SDK")
            print("   3. Upload all files from the hf_deployment folder")
            print("   4. Add GROQ_API_KEY as a secret in Space settings")
            print("   5. Your Space will automatically build and deploy!")
            
            return True
            
        except Exception as e:
            print(f"❌ Error preparing deployment: {e}")
            return False
    
    def create_hf_readme(self, deploy_dir: Path):
        """Create Hugging Face specific README"""
        hf_readme = """---
title: Power Systems Mini-Consultant
emoji: ⚡
colorFrom: blue
colorTo: cyan
sdk: gradio
sdk_version: 4.44.0
app_file: app.py
pinned: false
license: mit
short_description: AI-powered assistant for power systems engineering
tags:
- power-systems
- electrical-engineering
- fault-analysis
- protection-systems
- education
- RAG
- groq
---

# ⚡ Power Systems Mini-Consultant

An advanced AI-powered chatbot for Power Systems engineering education and professional support.

## Setup Instructions

1. **Configure API Key**: Add your GROQ_API_KEY as a secret in the Space settings
2. **Launch**: The app will automatically start once configured

## Features

- 💬 **AI Consultant Chat**: Technical assistance for power systems
- 📚 **Practice Pack Generator**: Custom exam preparation materials
- 📋 **Standards Explorer**: IEEE and IEC standards guidance
- 📖 **Study Resources**: Formulas and reference materials

## Usage

Simply start chatting with the AI consultant about power systems topics, or use the specialized tools for practice questions and standards exploration.

---

*For complete documentation, visit the project repository.*
"""
        
        with open(deploy_dir / 'README.md', 'w', encoding='utf-8') as f:
            f.write(hf_readme)
    
    def generate_deployment_report(self) -> Dict:
        """Generate a deployment report"""
        report = {
            "timestamp": str(Path.cwd()),
            "files_checked": len(self.required_files + self.optional_files),
            "status": "ready" if self.check_requirements() else "incomplete",
            "recommendations": []
        }
        
        # Check file sizes
        large_files = []
        for file_path in self.required_files:
            file_full_path = self.project_root / file_path
            if file_full_path.exists():
                size_mb = file_full_path.stat().st_size / (1024 * 1024)
                if size_mb > 10:  # Files larger than 10MB
                    large_files.append(f"{file_path} ({size_mb:.1f}MB)")
        
        if large_files:
            report["recommendations"].append(f"Large files detected: {large_files}")
        
        # Check for sensitive information
        env_file = self.project_root / '.env'
        if env_file.exists():
            report["recommendations"].append("Ensure .env file is in .gitignore")
        
        return report
    
    def run_deployment_checklist(self):
        """Run complete deployment checklist"""
        print("🚀 Power Systems Mini-Consultant Deployment Checklist")
        print("=" * 60)
        
        steps = [
            ("Check requirements", self.check_requirements),
            ("Validate knowledge base", self.validate_knowledge_base),
            ("Setup environment", self.setup_environment),
            ("Test local deployment", self.test_local_deployment),
            ("Prepare HF deployment", self.prepare_huggingface_deployment)
        ]
        
        passed_steps = 0
        total_steps = len(steps)
        
        for i, (step_name, step_func) in enumerate(steps, 1):
            print(f"\n[{i}/{total_steps}] {step_name}...")
            if step_func():
                passed_steps += 1
            else:
                print(f"❌ Step {i} failed. Please fix the issues and try again.")
                if i < 4:  # Don't break on HF deployment prep failure
                    break
        
        print(f"\n🎯 Deployment Summary: {passed_steps}/{total_steps} steps completed")
        
        if passed_steps == total_steps:
            print("🎉 All checks passed! Your app is ready for deployment.")
        elif passed_steps >= 3:
            print("⚠️  Basic requirements met. Review warnings and deploy when ready.")
        else:
            print("❌ Critical issues found. Please fix them before deploying.")
        
        # Generate report
        report = self.generate_deployment_report()
        report_path = self.project_root / 'deployment_report.json'
        with open(report_path, 'w') as f:
            json.dump(report, f, indent=2)
        print(f"\n📄 Detailed report saved to: {report_path}")

def main():
    """Main deployment script"""
    print("🔧 Power Systems Mini-Consultant Deployment Helper")
    print("=" * 50)
    
    deployer = PowerSystemsDeployer()
    
    if len(sys.argv) > 1:
        command = sys.argv[1].lower()
        
        if command == 'check':
            deployer.check_requirements()
        elif command == 'setup':
            deployer.setup_environment()
        elif command == 'test':
            deployer.test_local_deployment()
        elif command == 'prepare':
            deployer.prepare_huggingface_deployment()
        elif command == 'full':
            deployer.run_deployment_checklist()
        else:
            print(f"Unknown command: {command}")
            print("Available commands: check, setup, test, prepare, full")
    else:
        # Run full checklist by default
        deployer.run_deployment_checklist()

if __name__ == "__main__":
    main()