#!/usr/bin/env python3
"""
BLUESCARF AI HR Assistant - Automated Setup and Validation Script
Provides comprehensive setup, validation, and deployment assistance.
"""

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

class Colors:
    """ANSI color codes for terminal output."""
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

class SetupManager:
    """Comprehensive setup and validation manager for BLUESCARF AI HR Assistant."""
    
    def __init__(self):
        self.project_root = Path(__file__).parent
        self.requirements_file = self.project_root / "requirements.txt"
        self.config_file = self.project_root / "config.py"
        self.logo_file = self.project_root / "logo.png"
        
    def print_header(self):
        """Print application header with branding."""
        print(f"{Colors.HEADER}{Colors.BOLD}")
        print("=" * 60)
        print("   BLUESCARF ARTIFICIAL INTELLIGENCE")
        print("   HR Assistant Setup & Validation")
        print("   Version 1.0.0")
        print("=" * 60)
        print(f"{Colors.ENDC}")
    
    def check_python_version(self) -> bool:
        """Validate Python version compatibility."""
        print(f"{Colors.OKBLUE}Checking Python version...{Colors.ENDC}")
        
        version = sys.version_info
        min_version = (3, 8)
        
        if version >= min_version:
            print(f"{Colors.OKGREEN}✓ Python {version.major}.{version.minor}.{version.micro} (Compatible){Colors.ENDC}")
            return True
        else:
            print(f"{Colors.FAIL}✗ Python {version.major}.{version.minor}.{version.micro} (Requires 3.8+){Colors.ENDC}")
            return False
    
    def check_dependencies(self) -> Tuple[bool, List[str]]:
        """Check if all required dependencies are available."""
        print(f"{Colors.OKBLUE}Checking dependencies...{Colors.ENDC}")
        
        if not self.requirements_file.exists():
            print(f"{Colors.FAIL}✗ requirements.txt not found{Colors.ENDC}")
            return False, ["requirements.txt missing"]
        
        # Read requirements
        with open(self.requirements_file, 'r') as f:
            requirements = [line.strip() for line in f if line.strip() and not line.startswith('#')]
        
        missing_packages = []
        
        for requirement in requirements:
            package_name = requirement.split('==')[0].split('>=')[0].split('~=')[0]
            try:
                __import__(package_name.replace('-', '_'))
                print(f"{Colors.OKGREEN}✓ {package_name}{Colors.ENDC}")
            except ImportError:
                print(f"{Colors.WARNING}! {package_name} (not installed){Colors.ENDC}")
                missing_packages.append(package_name)
        
        if missing_packages:
            return False, missing_packages
        else:
            print(f"{Colors.OKGREEN}✓ All dependencies satisfied{Colors.ENDC}")
            return True, []
    
    def install_dependencies(self) -> bool:
        """Install missing dependencies using pip."""
        print(f"{Colors.OKBLUE}Installing dependencies...{Colors.ENDC}")
        
        try:
            subprocess.check_call([
                sys.executable, "-m", "pip", "install", "-r", str(self.requirements_file)
            ])
            print(f"{Colors.OKGREEN}✓ Dependencies installed successfully{Colors.ENDC}")
            return True
        except subprocess.CalledProcessError as e:
            print(f"{Colors.FAIL}✗ Failed to install dependencies: {e}{Colors.ENDC}")
            return False
    
    def validate_project_structure(self) -> Tuple[bool, List[str]]:
        """Validate that all required project files exist."""
        print(f"{Colors.OKBLUE}Validating project structure...{Colors.ENDC}")
        
        required_files = [
            "app.py",
            "document_processor.py",
            "vector_store.py",
            "admin.py",
            "config.py",
            "utils.py",
            "requirements.txt"
        ]
        
        missing_files = []
        
        for file_name in required_files:
            file_path = self.project_root / file_name
            if file_path.exists():
                print(f"{Colors.OKGREEN}✓ {file_name}{Colors.ENDC}")
            else:
                print(f"{Colors.FAIL}✗ {file_name} (missing){Colors.ENDC}")
                missing_files.append(file_name)
        
        # Check for logo
        if self.logo_file.exists():
            print(f"{Colors.OKGREEN}✓ logo.png (company logo found){Colors.ENDC}")
        else:
            print(f"{Colors.WARNING}! logo.png (add your company logo){Colors.ENDC}")
        
        if missing_files:
            return False, missing_files
        else:
            print(f"{Colors.OKGREEN}✓ Project structure is valid{Colors.ENDC}")
            return True, []
    
    def setup_directories(self) -> bool:
        """Create necessary directories for the application."""
        print(f"{Colors.OKBLUE}Setting up directories...{Colors.ENDC}")
        
        directories = [
            "vector_db",
            "logs",
            "temp",
            "data",
            "data/vector_db",
            "data/logs"
        ]
        
        try:
            for directory in directories:
                dir_path = self.project_root / directory
                dir_path.mkdir(parents=True, exist_ok=True)
                print(f"{Colors.OKGREEN}✓ Created {directory}/{Colors.ENDC}")
            
            return True
        except Exception as e:
            print(f"{Colors.FAIL}✗ Failed to create directories: {e}{Colors.ENDC}")
            return False
    
    def create_env_file(self) -> bool:
        """Create .env file from template if it doesn't exist."""
        print(f"{Colors.OKBLUE}Setting up environment configuration...{Colors.ENDC}")
        
        env_file = self.project_root / ".env"
        env_example = self.project_root / ".env.example"
        
        if env_file.exists():
            print(f"{Colors.OKGREEN}✓ .env file already exists{Colors.ENDC}")
            return True
        
        if env_example.exists():
            try:
                shutil.copy(env_example, env_file)
                print(f"{Colors.OKGREEN}✓ Created .env from .env.example{Colors.ENDC}")
                print(f"{Colors.WARNING}! Please review and customize .env file{Colors.ENDC}")
                return True
            except Exception as e:
                print(f"{Colors.FAIL}✗ Failed to create .env file: {e}{Colors.ENDC}")
                return False
        else:
            print(f"{Colors.WARNING}! .env.example not found, skipping .env creation{Colors.ENDC}")
            return True
    
    def validate_streamlit_config(self) -> bool:
        """Validate Streamlit configuration."""
        print(f"{Colors.OKBLUE}Validating Streamlit configuration...{Colors.ENDC}")
        
        try:
            import streamlit as st
            print(f"{Colors.OKGREEN}✓ Streamlit is available{Colors.ENDC}")
            return True
        except ImportError:
            print(f"{Colors.FAIL}✗ Streamlit not available{Colors.ENDC}")
            return False
    
    def test_api_imports(self) -> Dict[str, bool]:
        """Test critical API imports."""
        print(f"{Colors.OKBLUE}Testing critical imports...{Colors.ENDC}")
        
        import_tests = {
            "Google AI": ("google.generativeai", "google-generativeai"),
            "ChromaDB": ("chromadb", "chromadb"),
            "Sentence Transformers": ("sentence_transformers", "sentence-transformers"),
            "PyPDF2": ("PyPDF2", "PyPDF2"),
            "Pandas": ("pandas", "pandas"),
            "NumPy": ("numpy", "numpy")
        }
        
        results = {}
        
        for name, (module, package) in import_tests.items():
            try:
                __import__(module)
                print(f"{Colors.OKGREEN}✓ {name}{Colors.ENDC}")
                results[name] = True
            except ImportError:
                print(f"{Colors.FAIL}✗ {name} (install with: pip install {package}){Colors.ENDC}")
                results[name] = False
        
        return results
    
    def generate_deployment_summary(self) -> Dict[str, any]:
        """Generate comprehensive deployment summary."""
        print(f"{Colors.OKBLUE}Generating deployment summary...{Colors.ENDC}")
        
        summary = {
            "timestamp": time.time(),
            "python_version": f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}",
            "project_path": str(self.project_root),
            "files_present": [],
            "directories_created": [],
            "configuration_status": "pending"
        }
        
        # Check files
        for file_path in self.project_root.glob("*.py"):
            summary["files_present"].append(file_path.name)
        
        # Check directories
        for dir_path in ["vector_db", "logs", "temp"]:
            if (self.project_root / dir_path).exists():
                summary["directories_created"].append(dir_path)
        
        return summary
    
    def provide_next_steps(self):
        """Provide clear next steps for deployment."""
        print(f"\n{Colors.HEADER}{Colors.BOLD}NEXT STEPS:{Colors.ENDC}")
        print(f"{Colors.OKBLUE}1. Get Google Gemini API Key:{Colors.ENDC}")
        print("   → Visit: https://makersuite.google.com/app/apikey")
        print("   → Create or use existing API key")
        
        print(f"\n{Colors.OKBLUE}2. Add Company Logo:{Colors.ENDC}")
        print("   → Replace 'logo.png' with your company logo")
        print("   → Recommended size: 200x200 pixels")
        
        print(f"\n{Colors.OKBLUE}3. Upload Initial Documents:{Colors.ENDC}")
        print("   → Run the application: streamlit run app.py")
        print("   → Access admin panel with password: bluescarf_admin_2024")
        print("   → Upload HR policies, handbooks, procedures")
        
        print(f"\n{Colors.OKBLUE}4. Test the System:{Colors.ENDC}")
        print("   → Enter your API key in the application")
        print("   → Ask test questions about uploaded documents")
        print("   → Verify responses are accurate and relevant")
        
        print(f"\n{Colors.OKBLUE}5. Deploy to Production:{Colors.ENDC}")
        print("   → For Hugging Face Spaces: Upload all files")
        print("   → For Docker: Use provided Dockerfile")
        print("   → For cloud: Follow platform-specific guides")
        
        print(f"\n{Colors.WARNING}IMPORTANT SECURITY NOTES:{Colors.ENDC}")
        print("   → Change default admin password immediately")
        print("   → Keep API keys secure and never commit to git")
        print("   → Review uploaded documents for sensitive information")
        
        print(f"\n{Colors.OKGREEN}Ready for deployment! 🚀{Colors.ENDC}")
    
    def run_comprehensive_setup(self) -> bool:
        """Run complete setup and validation process."""
        self.print_header()
        
        success = True
        
        # 1. Check Python version
        if not self.check_python_version():
            print(f"{Colors.FAIL}Setup failed: Incompatible Python version{Colors.ENDC}")
            return False
        
        # 2. Validate project structure
        structure_valid, missing_files = self.validate_project_structure()
        if not structure_valid:
            print(f"{Colors.FAIL}Setup failed: Missing files: {missing_files}{Colors.ENDC}")
            return False
        
        # 3. Check dependencies
        deps_valid, missing_deps = self.check_dependencies()
        if not deps_valid:
            print(f"{Colors.WARNING}Installing missing dependencies...{Colors.ENDC}")
            if not self.install_dependencies():
                print(f"{Colors.FAIL}Setup failed: Could not install dependencies{Colors.ENDC}")
                return False
        
        # 4. Setup directories
        if not self.setup_directories():
            print(f"{Colors.FAIL}Setup failed: Could not create directories{Colors.ENDC}")
            return False
        
        # 5. Create environment file
        if not self.create_env_file():
            print(f"{Colors.WARNING}Environment file setup incomplete{Colors.ENDC}")
        
        # 6. Validate Streamlit
        if not self.validate_streamlit_config():
            print(f"{Colors.FAIL}Setup failed: Streamlit configuration issue{Colors.ENDC}")
            return False
        
        # 7. Test imports
        import_results = self.test_api_imports()
        if not all(import_results.values()):
            print(f"{Colors.WARNING}Some imports failed, but setup can continue{Colors.ENDC}")
        
        # 8. Generate summary
        summary = self.generate_deployment_summary()
        
        print(f"\n{Colors.OKGREEN}{Colors.BOLD}✓ SETUP COMPLETED SUCCESSFULLY!{Colors.ENDC}")
        
        # 9. Provide next steps
        self.provide_next_steps()
        
        return True

def main():
    """Main setup function."""
    setup_manager = SetupManager()
    
    if len(sys.argv) > 1:
        command = sys.argv[1]
        
        if command == "validate":
            # Quick validation only
            setup_manager.check_python_version()
            setup_manager.validate_project_structure()
            setup_manager.check_dependencies()
            
        elif command == "install":
            # Install dependencies only
            setup_manager.install_dependencies()
            
        elif command == "structure":
            # Setup directories only
            setup_manager.setup_directories()
            
        elif command == "test":
            # Test imports only
            setup_manager.test_api_imports()
            
        else:
            print(f"Unknown command: {command}")
            print("Available commands: validate, install, structure, test")
    else:
        # Run complete setup
        success = setup_manager.run_comprehensive_setup()
        sys.exit(0 if success else 1)

if __name__ == "__main__":
    main()