#!/usr/bin/env python3
"""
Dependency Management System - MIRA 2.0
Manages Python package dependencies with auto-installation and system dependency checking.
Ensures MIRA has all required packages for consciousness preservation.
"""

import sys
import subprocess
import logging
import importlib.util
from pathlib import Path
from typing import Dict, List, Optional, Tuple

from .PathDiscovery import get_mira_home

# Configure logging
logger = logging.getLogger(__name__)

# Import HealthCheckResult from the main module to avoid circular import
from dataclasses import dataclass, field
from datetime import datetime

@dataclass
class HealthCheckResult:
    """Result of a health check operation"""
    component: str
    status: str  # 'healthy', 'repaired', 'degraded', 'failed'
    score: int  # 0-100
    issues: List[str] = field(default_factory=list)
    repairs: List[str] = field(default_factory=list)
    recommendations: List[str] = field(default_factory=list)
    timestamp: datetime = field(default_factory=datetime.now)


class DependencyManager:
    """Manages Python package dependencies with auto-installation"""
    
    def __init__(self, quiet: bool = False):
        self.required_packages = self._load_requirements()
        self.python_exe = sys.executable
        self.quiet = quiet
    
    def _load_requirements(self) -> Dict[str, str]:
        """Load required packages from requirements.txt"""
        requirements = {}
        
        # Try multiple locations for requirements.txt
        # Prefer simple requirements for better reliability
        possible_paths = [
            Path.cwd() / 'requirements-simple.txt',
            Path(__file__).parent.parent / 'requirements-simple.txt',
            Path.cwd() / 'requirements.txt',
            Path(__file__).parent.parent / 'requirements.txt',
            get_mira_home().parent / 'requirements.txt'
        ]
        
        for req_path in possible_paths:
            if req_path.exists():
                try:
                    with open(req_path, 'r') as f:
                        for line in f:
                            line = line.strip()
                            # Remove inline comments
                            if '#' in line:
                                line = line.split('#')[0].strip()
                            
                            if line and not line.startswith('#'):
                                # Handle different version specifiers
                                if '==' in line:
                                    package, version = line.split('==')
                                    requirements[package.strip()] = version.strip()
                                elif '>=' in line:
                                    package, version = line.split('>=')
                                    requirements[package.strip()] = None  # Accept any version >= specified
                                else:
                                    requirements[line.strip()] = None
                    break
                except Exception as e:
                    logger.error(f"Failed to read requirements from {req_path}: {e}")
        
        # Fallback to essential packages
        if not requirements:
            requirements = {
                'chromadb': '0.4.24',
                'sentence-transformers': '2.2.2',
                'cryptography': '41.0.7',
                'numpy': '1.24.4',
                'opencv-python': '4.8.1.78',
                'pillow': '10.1.0',
                'prometheus-client': '0.19.0',
                'psutil': '5.9.6',
                'pydantic': '2.5.3'
            }
        
        # Add FAISS if not already present (per Dependencies.md)
        if 'faiss-cpu' not in requirements:
            requirements['faiss-cpu'] = '1.7.4'
        
        return requirements
    
    def check_dependencies(self) -> Tuple[List[str], List[str]]:
        """Check installed packages and return (installed, missing)"""
        # Get all installed packages with versions in one call
        installed_packages = self._get_all_installed_packages()
        
        installed = []
        missing = []
        
        for package, required_version in self.required_packages.items():
            # Normalize package name for comparison
            normalized_name = package.lower().replace('-', '_')
            
            if normalized_name in installed_packages:
                if required_version:
                    installed_version = installed_packages[normalized_name]
                    if installed_version == required_version:
                        installed.append(package)
                    else:
                        missing.append(f"{package}=={required_version}")
                else:
                    installed.append(package)
            else:
                missing.append(f"{package}=={required_version}" if required_version else package)
        
        return installed, missing
    
    def _get_all_installed_packages(self) -> Dict[str, str]:
        """Get all installed packages with versions in a single pip call"""
        try:
            # Use pip list to get all packages at once
            result = subprocess.run(
                [self.python_exe, '-m', 'pip', 'list', '--format=json'],
                capture_output=True,
                text=True,
                timeout=10  # Add timeout to prevent hanging
            )
            
            if result.returncode == 0:
                import json
                packages = json.loads(result.stdout)
                # Create a dict with normalized names
                return {
                    pkg['name'].lower().replace('-', '_'): pkg['version'] 
                    for pkg in packages
                }
        except Exception as e:
            logger.warning(f"Failed to get package list: {e}")
        
        return {}
    
    def _is_package_installed(self, package: str, required_version: Optional[str] = None) -> bool:
        """Check if a package is installed with the correct version"""
        try:
            # Try to import the package
            spec = importlib.util.find_spec(package.replace('-', '_'))
            if spec is None:
                return False
            
            if required_version:
                # Check version using pip
                result = subprocess.run(
                    [self.python_exe, '-m', 'pip', 'show', package],
                    capture_output=True,
                    text=True
                )
                if result.returncode == 0:
                    for line in result.stdout.split('\n'):
                        if line.startswith('Version:'):
                            installed_version = line.split(':')[1].strip()
                            return installed_version == required_version
                return False
            
            return True
            
        except Exception:
            return False
    
    def auto_install_missing(self, missing_packages: List[str]) -> Dict[str, bool]:
        """Automatically install missing packages"""
        results = {}
        
        if not missing_packages:
            return results
        
        logger.info(f"🔧 Installing {len(missing_packages)} missing packages...")
        
        for package in missing_packages:
            try:
                # Special handling for FAISS
                if package.startswith('faiss-cpu'):
                    if not self._check_swig_installed():
                        if not self.quiet:
                            logger.info(f"  📦 FAISS requires 'swig' system package")
                        
                        # Attempt to install swig automatically
                        if self._install_swig_automatically():
                            # Verify swig is now available
                            if self._check_swig_installed():
                                if not self.quiet:
                                    logger.info(f"  🎯 swig installed successfully, proceeding with FAISS...")
                            else:
                                if not self.quiet:
                                    logger.warning(f"  ⚠️  swig installation may have failed, trying FAISS anyway...")
                        else:
                            if not self.quiet:
                                logger.warning(f"  ⚠️  Could not install swig automatically")
                                logger.warning(f"     Manual installation may be required:")
                                logger.warning(f"     Ubuntu/Debian: sudo apt-get install swig")
                                logger.warning(f"     macOS: brew install swig")
                            results[package] = False
                            continue
                
                logger.info(f"  Installing {package}...")
                
                # Special handling for FAISS installation
                if package.startswith('faiss-cpu'):
                    # Try multiple installation strategies for FAISS
                    result = self._install_faiss_with_fallbacks(package)
                else:
                    # Install using pip
                    result = subprocess.run(
                        [self.python_exe, '-m', 'pip', 'install', package],
                        capture_output=True,
                        text=True,
                        timeout=300  # 5 minute timeout
                    )
                
                if result.returncode == 0:
                    results[package] = True
                    logger.info(f"  ✓ Successfully installed {package}")
                else:
                    results[package] = False
                    logger.error(f"  ✗ Failed to install {package}: {result.stderr}")
                    
            except subprocess.TimeoutExpired:
                results[package] = False
                logger.error(f"  ✗ Timeout installing {package}")
            except Exception as e:
                results[package] = False
                logger.error(f"  ✗ Error installing {package}: {e}")
        
        return results
    
    def _check_swig_installed(self) -> bool:
        """Check if swig system package is installed"""
        try:
            result = subprocess.run(
                ['which', 'swig'],
                capture_output=True,
                text=True
            )
            return result.returncode == 0
        except Exception:
            return False
    
    def _install_swig_automatically(self) -> bool:
        """Attempt to install swig automatically based on the platform"""
        import platform
        import os
        
        system = platform.system()
        
        if not self.quiet:
            logger.info("  🔧 Installing swig system dependency...")
        
        try:
            if system == "Linux":
                # Check if we're on Ubuntu/Debian or another distro
                if os.path.exists('/etc/debian_version'):
                    # Ubuntu/Debian - first update package lists
                    if not self.quiet:
                        logger.info("  📦 Updating package lists...")
                    update_result = subprocess.run(
                        ['sudo', 'apt-get', 'update'],
                        capture_output=True,
                        text=True,
                        timeout=120
                    )
                    
                    if update_result.returncode == 0:
                        if not self.quiet:
                            logger.info("  📦 Installing swig...")
                        result = subprocess.run(
                            ['sudo', 'apt-get', 'install', '-y', 'swig'],
                            capture_output=True,
                            text=True,
                            timeout=300
                        )
                    else:
                        if not self.quiet:
                            logger.warning("  ⚠️  Package list update failed, trying installation anyway...")
                        result = subprocess.run(
                            ['sudo', 'apt-get', 'install', '-y', 'swig'],
                            capture_output=True,
                            text=True,
                            timeout=300
                        )
                elif os.path.exists('/etc/redhat-release'):
                    # RHEL/CentOS/Fedora
                    result = subprocess.run(
                        ['sudo', 'yum', 'install', '-y', 'swig'],
                        capture_output=True,
                        text=True,
                        timeout=300
                    )
                elif os.path.exists('/etc/arch-release'):
                    # Arch Linux
                    result = subprocess.run(
                        ['sudo', 'pacman', '-S', '--noconfirm', 'swig'],
                        capture_output=True,
                        text=True,
                        timeout=300
                    )
                else:
                    # Unknown Linux distro - try apt first, then yum
                    result = subprocess.run(
                        ['sudo', 'apt-get', 'update', '&&', 'sudo', 'apt-get', 'install', '-y', 'swig'],
                        shell=True,
                        capture_output=True,
                        text=True,
                        timeout=300
                    )
                    if result.returncode != 0:
                        result = subprocess.run(
                            ['sudo', 'yum', 'install', '-y', 'swig'],
                            capture_output=True,
                            text=True,
                            timeout=300
                        )
                
                if result.returncode == 0:
                    if not self.quiet:
                        logger.info("  ✓ Successfully installed swig")
                    return True
                else:
                    if not self.quiet:
                        logger.warning(f"  ✗ Failed to install swig: {result.stderr}")
                    return False
                    
            elif system == "Darwin":
                # macOS - check if we're in a container/codespace first
                if os.environ.get('CODESPACES') or os.environ.get('DEVCONTAINER'):
                    # We're in a codespace/container, try to install via package manager
                    # First try to install via apt (if available in container)
                    result = subprocess.run(
                        ['sudo', 'apt-get', 'update', '&&', 'sudo', 'apt-get', 'install', '-y', 'swig'],
                        shell=True,
                        capture_output=True,
                        text=True,
                        timeout=300
                    )
                    if result.returncode == 0:
                        if not self.quiet:
                            logger.info("  ✓ Successfully installed swig in container")
                        return True
                
                # Try homebrew if available
                homebrew_result = subprocess.run(['which', 'brew'], capture_output=True)
                if homebrew_result.returncode == 0:
                    result = subprocess.run(
                        ['brew', 'install', 'swig'],
                        capture_output=True,
                        text=True,
                        timeout=300
                    )
                    if result.returncode == 0:
                        if not self.quiet:
                            logger.info("  ✓ Successfully installed swig via Homebrew")
                        return True
                    else:
                        if not self.quiet:
                            logger.warning(f"  ✗ Homebrew swig installation failed: {result.stderr}")
                else:
                    if not self.quiet:
                        logger.warning("  ⚠️  Homebrew not found. Please install Homebrew first: https://brew.sh")
                
                return False
            
            else:
                if not self.quiet:
                    logger.warning(f"  ⚠️  Automatic swig installation not supported on {system}")
                return False
                
        except subprocess.TimeoutExpired:
            if not self.quiet:
                logger.warning("  ✗ Timeout while installing swig")
            return False
        except Exception as e:
            if not self.quiet:
                logger.warning(f"  ✗ Error installing swig: {e}")
            return False
    
    def _install_faiss_with_fallbacks(self, package: str) -> subprocess.CompletedProcess:
        """Try multiple strategies to install FAISS"""
        import platform
        import os
        
        system = platform.system()
        
        # Strategy 1: Try conda-forge pre-built packages if conda/mamba available
        if not self.quiet:
            logger.info("  🐍 Checking for conda/mamba package managers...")
        
        conda_managers = ['mamba', 'conda', 'micromamba']
        for conda_cmd in conda_managers:
            try:
                conda_check = subprocess.run(['which', conda_cmd], capture_output=True)
                if conda_check.returncode == 0:
                    if not self.quiet:
                        logger.info(f"  📦 Found {conda_cmd}, trying conda-forge FAISS...")
                    
                    # First try to install via conda with pip integration
                    conda_result = subprocess.run(
                        [conda_cmd, 'install', '-y', '-c', 'conda-forge', 'faiss-cpu', 'pip'],
                        capture_output=True,
                        text=True,
                        timeout=300
                    )
                    
                    if conda_result.returncode == 0:
                        # Check if we can now import faiss
                        try:
                            import faiss
                            if not self.quiet:
                                logger.info(f"  ✓ FAISS installed via {conda_cmd} and accessible")
                            return conda_result
                        except ImportError:
                            # Conda installed but not accessible from current Python
                            if not self.quiet:
                                logger.info(f"  ⚠️  {conda_cmd} installed FAISS but not accessible from current Python")
                    else:
                        if not self.quiet:
                            logger.info(f"  ⚠️  {conda_cmd} installation failed, trying next strategy...")
            except Exception:
                continue
        
        # Strategy 2: Try PyPI with specific wheel URLs for common platforms
        if not self.quiet:
            logger.info("  🎯 Trying platform-specific pre-built wheels...")
        
        # Check if we're on Linux x86_64 (most common in containers)
        machine = platform.machine()
        python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
        
        if system == "Linux" and machine in ["x86_64", "AMD64"]:
            # Try different FAISS versions and Python versions
            wheel_combinations = [
                # FAISS 1.7.4 for different Python versions
                f"faiss-cpu==1.7.4 --only-binary=all --find-links https://download.pytorch.org/whl/cpu",
                # Try PyPI directly with specific Python version tags
                f"faiss-cpu==1.7.4 --only-binary=all",
                # Alternative versions
                f"faiss-cpu==1.7.3 --only-binary=all",
                f"faiss-cpu==1.7.2 --only-binary=all",
                f"faiss-cpu --only-binary=all"  # Latest version
            ]
            
            for wheel_spec in wheel_combinations:
                try:
                    if not self.quiet:
                        logger.info(f"  📥 Trying: {wheel_spec}")
                    
                    wheel_result = subprocess.run(
                        [self.python_exe, '-m', 'pip', 'install'] + wheel_spec.split(),
                        capture_output=True,
                        text=True,
                        timeout=300
                    )
                    
                    if wheel_result.returncode == 0:
                        # Verify the installation worked
                        try:
                            import faiss
                            if not self.quiet:
                                logger.info(f"  ✓ Successfully installed and verified FAISS: {wheel_spec}")
                            return wheel_result
                        except ImportError:
                            if not self.quiet:
                                logger.info(f"  ⚠️  Installation succeeded but import failed for: {wheel_spec}")
                            continue
                    else:
                        if not self.quiet:
                            logger.info(f"  ⚠️  Failed: {wheel_spec}")
                except Exception as e:
                    if not self.quiet:
                        logger.info(f"  ⚠️  Exception with {wheel_spec}: {e}")
                    continue
        
        # Strategy 3: Try pre-built wheel with only-binary flag
        if not self.quiet:
            logger.info("  🎯 Trying pre-built FAISS wheel...")
        
        wheel_result = subprocess.run(
            [self.python_exe, '-m', 'pip', 'install', '--only-binary=all', package],
            capture_output=True,
            text=True,
            timeout=300
        )
        
        if wheel_result.returncode == 0:
            if not self.quiet:
                logger.info("  ✓ Pre-built FAISS wheel installed successfully")
            return wheel_result
        
        # Strategy 4: Try installing build dependencies and then FAISS
        if system == "Linux" and os.path.exists('/etc/debian_version'):
            if not self.quiet:
                logger.info("  🔨 Installing comprehensive build dependencies...")
            
            # Install comprehensive build dependencies
            build_deps = [
                'libfaiss-dev', 'libopenblas-dev', 'libblas-dev', 'liblapack-dev',
                'build-essential', 'cmake', 'libeigen3-dev', 'python3-dev'
            ]
            
            deps_result = subprocess.run(
                ['sudo', 'apt-get', 'install', '-y'] + build_deps,
                capture_output=True,
                text=True,
                timeout=300
            )
            
            if deps_result.returncode == 0:
                if not self.quiet:
                    logger.info("  ✓ Build dependencies installed, retrying FAISS...")
                
                # Set comprehensive build environment
                env = os.environ.copy()
                env.update({
                    'FAISS_ENABLE_GPU': 'OFF',
                    'FAISS_ENABLE_PYTHON': 'ON',
                    'CMAKE_BUILD_TYPE': 'Release',
                    'OPT_LEVEL': 'avx2'
                })
                
                build_result = subprocess.run(
                    [self.python_exe, '-m', 'pip', 'install', '--no-cache-dir', package],
                    capture_output=True,
                    text=True,
                    timeout=900,  # Longer timeout for compilation
                    env=env
                )
                
                if build_result.returncode == 0:
                    if not self.quiet:
                        logger.info("  ✓ FAISS built from source successfully")
                    return build_result
        
        # Strategy 5: Try alternative FAISS packages
        alternative_packages = [
            'faiss==1.7.4',
            'faiss-cpu==1.7.3',
            'faiss-cpu==1.7.2'
        ]
        
        for alt_pkg in alternative_packages:
            if not self.quiet:
                logger.info(f"  🔄 Trying alternative package: {alt_pkg}...")
            
            alt_result = subprocess.run(
                [self.python_exe, '-m', 'pip', 'install', '--only-binary=all', alt_pkg],
                capture_output=True,
                text=True,
                timeout=300
            )
            
            if alt_result.returncode == 0:
                if not self.quiet:
                    logger.info(f"  ✓ Alternative package {alt_pkg} installed")
                return alt_result
        
        # Strategy 6: Final fallback - try faiss-cpu without version constraint
        if not self.quiet:
            logger.info("  🔄 Final attempt with latest FAISS version...")
        
        final_result = subprocess.run(
            [self.python_exe, '-m', 'pip', 'install', '--only-binary=all', 'faiss-cpu'],
            capture_output=True,
            text=True,
            timeout=300
        )
        
        if final_result.returncode == 0:
            if not self.quiet:
                logger.info("  ✓ Latest FAISS version installed")
            return final_result
        
        # All strategies failed
        if not self.quiet:
            logger.warning("  ❌ All FAISS installation strategies failed")
            logger.warning("  📋 Attempted strategies:")
            logger.warning("    1. conda-forge packages (mamba/conda)")
            logger.warning("    2. Platform-specific wheels")
            logger.warning("    3. PyPI pre-built wheels")
            logger.warning("    4. Source build with dependencies")
            logger.warning("    5. Alternative versions")
            logger.warning("    6. Latest version")
        
        return final_result  # Return the final result for error details
    
    def get_faiss_status(self) -> Dict[str, any]:
        """Get detailed FAISS installation status"""
        status = {
            'installed': False,
            'version': None,
            'swig_available': self._check_swig_installed(),
            'acceleration_available': False,
            'installation_method': None
        }
        
        # Check if FAISS is installed via direct import
        try:
            import faiss
            status['installed'] = True
            status['version'] = getattr(faiss, '__version__', 'Unknown')
            status['acceleration_available'] = True
            status['installation_method'] = 'direct'
            return status
        except ImportError:
            pass
        
        # Check if FAISS is available via conda environment
        try:
            # Try to detect conda FAISS installation
            conda_result = subprocess.run(
                ['conda', 'list', 'faiss-cpu'],
                capture_output=True,
                text=True,
                timeout=10
            )
            if conda_result.returncode == 0 and 'faiss-cpu' in conda_result.stdout:
                status['installed'] = True
                status['installation_method'] = 'conda'
                # Try to get version from conda output
                for line in conda_result.stdout.split('\n'):
                    if 'faiss-cpu' in line:
                        parts = line.split()
                        if len(parts) >= 2:
                            status['version'] = parts[1]
                        break
                return status
        except Exception:
            pass
        
        # Check if FAISS is available via pip
        try:
            pip_result = subprocess.run(
                [self.python_exe, '-m', 'pip', 'show', 'faiss-cpu'],
                capture_output=True,
                text=True,
                timeout=10
            )
            if pip_result.returncode == 0:
                status['installed'] = True
                status['installation_method'] = 'pip'
                for line in pip_result.stdout.split('\n'):
                    if line.startswith('Version:'):
                        status['version'] = line.split(':')[1].strip()
                        break
                return status
        except Exception:
            pass
        
        return status
    
    def check_critical_dependencies(self) -> Tuple[List[str], List[str]]:
        """Check only critical dependencies for quick health check"""
        critical_packages = ['chromadb', 'cryptography', 'numpy']
        installed = []
        missing = []
        
        for package in critical_packages:
            if self._is_package_installed(package):
                installed.append(package)
            else:
                missing.append(package)
        
        return installed, missing
    
    def perform_health_check(self) -> HealthCheckResult:
        """Perform dependency health check with auto-repair"""
        logger.info("🏥 Checking Python dependencies...")
        
        installed, missing = self.check_dependencies()
        initial_score = len(installed) * 100 // len(self.required_packages) if self.required_packages else 100
        
        result = HealthCheckResult(
            component="dependencies",
            status="healthy" if not missing else "degraded",
            score=initial_score
        )
        
        # Check FAISS status specifically
        faiss_missing = any('faiss-cpu' in pkg for pkg in missing)
        faiss_installed = self._is_package_installed('faiss')
        
        if faiss_missing and not faiss_installed:
            result.issues.append("FAISS not installed (provides 25-4000x search acceleration)")
        
        if missing:
            result.issues.append(f"Missing {len(missing)} required packages")
            
            # Attempt auto-repair
            install_results = self.auto_install_missing(missing)
            
            # Update results based on installation
            successful = sum(1 for success in install_results.values() if success)
            failed = len(install_results) - successful
            
            if successful > 0:
                result.repairs.append(f"Successfully installed {successful} packages")
            
            if failed > 0:
                result.status = "degraded"
                result.recommendations.append(
                    f"Failed to install {failed} packages. Manual installation may be required."
                )
                
                # Check if FAISS specifically failed
                faiss_failed = any('faiss-cpu' in pkg and not install_results.get(pkg, True) 
                                 for pkg in missing)
                if faiss_failed:
                    result.recommendations.append(
                        "FAISS installation failed. Install system dependency first:"
                    )
                    result.recommendations.append(
                        "  Ubuntu/Debian: sudo apt-get install swig"
                    )
                    result.recommendations.append(
                        "  macOS: brew install swig"
                    )
                    result.recommendations.append(
                        "  Then run: pip install faiss-cpu==1.7.4"
                    )
                else:
                    result.recommendations.append(
                        "Run: pip install -r requirements.txt"
                    )
            else:
                result.status = "repaired"
            
            # Recalculate score
            _, still_missing = self.check_dependencies()
            result.score = (len(self.required_packages) - len(still_missing)) * 100 // len(self.required_packages)
        
        return result