#!/usr/bin/env python3
"""
Resource Recovery Utility for PyInstaller Applications
Handles recovery of bundled resources when temp folders are deleted
"""

import os
import sys
import shutil
import platform
import time
from pathlib import Path
from typing import Optional, List, Dict
from ws_package.log.writer import LogWriter

logger = LogWriter.file(__name__)

class ResourceRecovery:
    """
    Handles recovery of bundled resources when PyInstaller temp folders are deleted.
    """
    
    def __init__(self):
        self.platform = platform.system().lower()
        self._recovery_cache = {}
        self._validation_cache = {}  # Cache validation results
    
    def get_permanent_dir(self) -> Path:
        """Get the permanent directory for storing recovered resources"""
        if self.platform == "windows":
            return Path(os.environ.get('APPDATA', '')) / "Workstatus"
        elif self.platform == "darwin":
            return Path.home() / "Library" / "Application Support" / "Workstatus"
        else:
            return Path.home() / ".workstatus"
    
    def recover_resource(self, resource_name: str, bundle_path: str, 
                        permanent_subdir: str = None, optional: bool = False) -> Optional[Path]:
        """
        Recover a bundled resource to a permanent location.
        
        Args:
            resource_name: Name of the resource file/directory
            bundle_path: Path within the PyInstaller bundle
            permanent_subdir: Subdirectory in permanent location (optional)
            optional: If True, log warnings instead of errors for missing resources
        
        Returns:
            Path to the recovered resource, or None if recovery failed
        """
        if not getattr(sys, 'frozen', False):
            if optional:
                logger.debug("Resource recovery only works in PyInstaller bundles (optional resource)")
            else:
                logger.warning("Resource recovery only works in PyInstaller bundles")
            return None
        
        # Check cache first
        cache_key = f"{resource_name}_{bundle_path}"
        if cache_key in self._recovery_cache:
            cached_path = self._recovery_cache[cache_key]
            if cached_path.exists():
                return cached_path
        
        try:
            # Determine permanent location
            permanent_dir = self.get_permanent_dir()
            if permanent_subdir:
                permanent_dir = permanent_dir / permanent_subdir
            else:
                permanent_dir = permanent_dir / resource_name
            
            permanent_dir.mkdir(parents=True, exist_ok=True)
            
            # Try to find the resource in the bundle
            bundle_resource_path = Path(sys._MEIPASS) / bundle_path
            if not bundle_resource_path.exists():
                if optional:
                    logger.debug("Optional bundled resource not found (this is normal if not needed): %s", bundle_resource_path)
                else:
                    logger.error("Bundled resource not found: %s", bundle_resource_path)
                return None
            
            # Copy the resource
            if bundle_resource_path.is_file():
                permanent_path = permanent_dir / resource_name
                shutil.copy2(bundle_resource_path, permanent_path)
                logger.info("Recovered file %s to: %s", resource_name, permanent_path)
            elif bundle_resource_path.is_dir():
                permanent_path = permanent_dir
                if permanent_path.exists():
                    shutil.rmtree(permanent_path)
                shutil.copytree(bundle_resource_path, permanent_path)
                logger.info("Recovered directory %s to: %s", resource_name, permanent_path)
            else:
                logger.error("Unknown resource type: %s", bundle_resource_path)
                return None
            
            # Cache the result
            self._recovery_cache[cache_key] = permanent_path
            return permanent_path
            
        except Exception as e:
            logger.error("Failed to recover resource %s: %s", resource_name, e)
            return None
    
    def _validate_resource(self, path: Path, resource_name: str) -> bool:
        """
        Validate that a resource is complete and usable.
        
        Args:
            path: Path to the resource
            resource_name: Name of the resource for validation logic
            
        Returns:
            True if resource is valid, False otherwise
        """
        if not path or not path.exists():
            return False
        
        # Check cache first
        cache_key = f"validate_{resource_name}_{path}"
        if cache_key in self._validation_cache:
            cached_result, cached_time = self._validation_cache[cache_key]
            # Cache validation for 1 hour
            if time.time() - cached_time < 3600:
                return cached_result
        
        try:
            # For directories, check if they're not empty
            if path.is_dir():
                try:
                    has_files = any(path.iterdir())
                    self._validation_cache[cache_key] = (has_files, time.time())
                    return has_files
                except Exception:
                    self._validation_cache[cache_key] = (False, time.time())
                    return False
            
            # For files, check if readable
            if path.is_file():
                try:
                    readable = path.stat().st_size > 0
                    self._validation_cache[cache_key] = (readable, time.time())
                    return readable
                except Exception:
                    self._validation_cache[cache_key] = (False, time.time())
                    return False
            
            self._validation_cache[cache_key] = (True, time.time())
            return True
            
        except Exception as e:
            logger.debug("Resource validation error for %s: %s", resource_name, e)
            self._validation_cache[cache_key] = (False, time.time())
            return False
    
    def get_resource_path(self, resource_name: str, bundle_path: str, 
                         permanent_subdir: str = None) -> Optional[Path]:
        """
        Get resource path with AppData-first approach (most stable).
        Checks AppData first, then bundle, then recovers if needed.
        
        Args:
            resource_name: Name of the resource file/directory
            bundle_path: Path within the PyInstaller bundle
            permanent_subdir: Subdirectory in permanent location (optional)
            
        Returns:
            Path to the resource, or None if not available
        """
        # 1. Check AppData (permanent location) - most stable
        permanent_dir = self.get_permanent_dir()
        if permanent_subdir:
            permanent_dir = permanent_dir / permanent_subdir
        else:
            permanent_dir = permanent_dir / resource_name
        
        if permanent_dir.exists():
            if self._validate_resource(permanent_dir, resource_name):
                logger.debug("Using AppData resource: %s", permanent_dir)
                return permanent_dir
            else:
                logger.warning("AppData resource invalid, will recover: %s", permanent_dir)
        
        # 2. Check bundle (if still available) and copy to AppData
        if getattr(sys, 'frozen', False):
            try:
                bundle_path_full = Path(sys._MEIPASS) / bundle_path
                if bundle_path_full.exists():
                    logger.info("Bundle resource found, copying to AppData: %s", bundle_path_full)
                    # Copy to AppData for future use
                    recovered = self.recover_resource(
                        resource_name, 
                        bundle_path, 
                        permanent_subdir,
                        optional=True
                    )
                    if recovered and self._validate_resource(recovered, resource_name):
                        return recovered
            except Exception as e:
                logger.debug("Bundle check failed: %s", e)
        
        # 3. Try recovery if we have a cached path
        cache_key = f"{resource_name}_{bundle_path}"
        if cache_key in self._recovery_cache:
            cached_path = self._recovery_cache[cache_key]
            if cached_path.exists() and self._validate_resource(cached_path, resource_name):
                return cached_path
        
        return None
    
    def get_validated_resource(self, resource_name: str, bundle_path: str,
                               permanent_subdir: str = None) -> Optional[Path]:
        """
        Get resource with validation and automatic recovery.
        This is the recommended method to use when accessing resources.
        
        Args:
            resource_name: Name of the resource file/directory
            bundle_path: Path within the PyInstaller bundle
            permanent_subdir: Subdirectory in permanent location (optional)
            
        Returns:
            Path to validated resource, or None if not available
        """
        # Try to get existing resource
        resource_path = self.get_resource_path(resource_name, bundle_path, permanent_subdir)
        
        if resource_path and self._validate_resource(resource_path, resource_name):
            return resource_path
        
        # Resource not found or invalid, try recovery
        logger.info("Resource %s not found or invalid, attempting recovery...", resource_name)
        recovered = self.recover_resource(resource_name, bundle_path, permanent_subdir, optional=True)
        
        if recovered and self._validate_resource(recovered, resource_name):
            return recovered
        
        return None
    
    def pre_copy_critical_resources(self) -> Dict[str, bool]:
        """
        Proactively copy critical resources to AppData on startup.
        This prevents issues if temp folder is deleted later.

        ONEDIR FAST-PATH:
          In onedir mode (_MEIPASS == <exe_dir>/_internal) the bundle is
          already on disk permanently next to the EXE.  There is nothing
          to pre-copy — and proactively copying _internal / AWS
          resources to AppData would re-create exactly the legacy
          AppData clutter that _quarantine_stale_appdata_install() works
          hard to remove.  Return an empty result early.

        Returns:
            Dictionary mapping resource names to copy success status
        """
        results = {}

        if not getattr(sys, 'frozen', False):
            logger.debug("Pre-copy only works in PyInstaller bundles")
            return results

        # Onedir detection: skip the whole legacy AppData-recovery path.
        if hasattr(sys, '_MEIPASS'):
            try:
                meipass = Path(sys._MEIPASS)
                exe_dir = Path(sys.executable).parent
                if meipass.parent == exe_dir:
                    logger.debug(
                        "Onedir mode: skipping pre_copy_critical_resources "
                        "(bundle is already permanent next to the EXE)."
                    )
                    return results
            except (OSError, ValueError):
                pass

        logger.info("Pre-copying critical resources to AppData...")
        
        # Pre-copy _internal folder (contains Python DLLs) - CRITICAL for Windows
        if platform.system() == "Windows":
            internal_path = self.recover_internal_folder()
            results['_internal'] = internal_path is not None and self._validate_resource(
                internal_path, "_internal"
            ) if internal_path else False
            if results['_internal']:
                logger.info("_internal folder pre-copied successfully")
            else:
                logger.warning("_internal folder pre-copy failed - Python DLL may not be found")
        
        # Pre-copy AWS resources
        aws_path = self.recover_aws_resources()
        results['aws_resources'] = aws_path is not None and self._validate_resource(
            aws_path, "aws_resources"
        ) if aws_path else False
        if results['aws_resources']:
            logger.info("AWS resources pre-copied successfully")
        else:
            logger.debug("AWS resources pre-copy skipped (optional resource)")
        
        return results
    
    def recover_internal_folder(self) -> Optional[Path]:
        """
        Recover _internal folder (contains Python DLLs) to permanent location.

        In onedir mode (PyInstaller ≥ 6) _MEIPASS IS _internal next to the EXE,
        so it is already permanent — no copy is needed.  We return _MEIPASS
        directly and, as a side-effect, clean up any stale AppData/_internal copy
        left over from a pre-onedir install so it never confuses SSL path resolution.
        """
        if platform.system() != "Windows":
            logger.debug("_internal folder recovery is Windows-specific")
            return None

        # Win11 22H2+ Smart App Control (SAC) can block unsigned binaries.
        # We detect it early and emit an actionable warning so the admin knows
        # to either sign the build or disable SAC for this install path.
        try:
            import winreg
            _sac_key = r"SYSTEM\CurrentControlSet\Control\CI\Policy"
            with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, _sac_key) as _k:
                _sac_state, _ = winreg.QueryValueEx(_k, "VerifiedAndReputablePolicyState")
            if _sac_state == 1:
                logger.warning(
                    "Windows Smart App Control (SAC) is ENABLED on this machine. "
                    "Unsigned binaries in _internal may be blocked. "
                    "Sign the installer and all bundled EXEs, or ask the user to "
                    "disable SAC in Windows Security > App & browser control."
                )
        except (FileNotFoundError, OSError, KeyError):
            pass  # SAC registry key absent (Win10 or SAC never activated)
        except Exception:
            pass

        try:
            exe_dir = Path(sys.executable).parent
            meipass = Path(sys._MEIPASS) if hasattr(sys, '_MEIPASS') else None

            # --- ONEDIR DETECTION ---
            # In onedir, PyInstaller sets _MEIPASS = <exe_dir>/_internal.
            # We detect this by checking whether _MEIPASS's parent is the exe dir.
            if meipass and meipass.parent == exe_dir:
                logger.debug("Onedir mode detected: _internal is permanent at %s", meipass)

                # Clean up any stale AppData/_internal left by a pre-onedir install.
                # A stale copy there confuses certifi / SSL path resolution.
                appdata_internal = self.get_permanent_dir() / "_internal"
                version_marker = self.get_permanent_dir() / ".ws_internal_ver"
                if appdata_internal.exists():
                    try:
                        shutil.rmtree(appdata_internal)
                        logger.info("Removed stale AppData/_internal: %s", appdata_internal)
                    except Exception as e:
                        logger.warning("Could not remove stale AppData/_internal: %s", e)
                if version_marker.exists():
                    try:
                        version_marker.unlink()
                    except Exception:
                        pass

                return meipass  # already the right permanent dir

            # Onedir is the only supported build mode. If _MEIPASS is not a
            # child of the exe directory, the build is misconfigured.
            logger.error(
                "Unexpected _MEIPASS location %s — expected under %s. "
                "Ensure the build uses onedir mode (Workstatus-Windows.spec).",
                meipass, exe_dir,
            )
            return None

        except Exception as e:
            logger.error("Failed to recover _internal folder: %s", e)
            return None
    
    def recover_aws_resources(self) -> Optional[Path]:
        """
        Recover AWS-related resources to permanent location.
        Note: This resource may not exist in all builds.
        """
        return self.recover_resource(
            "aws_resources",
            "ws_package/aws_resources",
            "aws_resources",
            optional=True  # Optional because it may not exist in all builds
        )
    
    def check_ssl_certificates(self) -> Dict[str, bool]:
        """
        Check if SSL certificates are available.
        
        Returns:
            Dictionary mapping certificate types to availability status
        """
        import ssl
        import os
        
        certificates = {}
        
        try:
            # Try to create default SSL context
            ssl_context = ssl.create_default_context()
            certificates['default_context'] = True
            logger.info("Default SSL context created successfully")
        except (FileNotFoundError, OSError) as e:
            certificates['default_context'] = False
            logger.error("Default SSL context creation failed: %s", e)
        
        # Check for common certificate file locations
        cert_locations = [
            '/etc/ssl/certs/ca-certificates.crt',  # Ubuntu/Debian
            '/etc/ssl/certs/ca-bundle.crt',         # CentOS/RHEL
            '/etc/pki/tls/certs/ca-bundle.crt',     # Fedora
            '/usr/local/share/ca-certificates/',   # Custom certs
            '/etc/ssl/cert.pem',                    # macOS
        ]
        
        for location in cert_locations:
            if os.path.exists(location):
                certificates[f'cert_file_{os.path.basename(location)}'] = True
                logger.info("Found certificate file: %s", location)
            else:
                certificates[f'cert_file_{os.path.basename(location)}'] = False
        
        return certificates
    
    def recover_ssl_certificates(self) -> bool:
        """
        Attempt to recover SSL certificates by downloading them.
        This is a fallback when bundled certificates are missing.
        
        Returns:
            True if recovery was successful, False otherwise
        """
        try:
            import ssl
            import urllib.request
            import tempfile
            import os
            
            logger.info("Attempting to recover SSL certificates...")
            
            logger.warning("SSL certificates missing — attempting recovery via default context")
            
            # Try to create a temporary certificate bundle
            try:
                # Use Python's built-in certificate handling
                ssl_context = ssl.create_default_context()
                logger.info("SSL certificates recovered using default context")
                return True
            except Exception as e:
                logger.warning("Default SSL context failed: %s", e)
            
            # Try to download certificates from a trusted source
            try:
                # Download Mozilla's certificate bundle
                cert_url = "https://curl.se/ca/cacert.pem"
                with tempfile.NamedTemporaryFile(delete=False, suffix='.pem') as temp_cert:
                    urllib.request.urlretrieve(cert_url, temp_cert.name)
                    
                    # Set environment variable for SSL certificate bundle
                    os.environ['SSL_CERT_FILE'] = temp_cert.name
                    os.environ['REQUESTS_CA_BUNDLE'] = temp_cert.name
                    
                    logger.info("Downloaded SSL certificates to: %s", temp_cert.name)
                    return True
                    
            except Exception as e:
                logger.error("Failed to download SSL certificates: %s", e)
                return False
                
        except Exception as e:
            logger.error("Error during SSL certificates recovery: %s", e)
            return False
    
    def check_aws_dependencies(self) -> Dict[str, bool]:
        """
        Check if AWS dependencies (boto3, botocore) are available.
        
        Returns:
            Dictionary mapping dependency names to availability status
        """
        dependencies = {}
        
        try:
            import boto3
            dependencies['boto3'] = True
            logger.info("boto3 module is available")
        except ImportError as e:
            dependencies['boto3'] = False
            logger.error("boto3 module not available: %s", e)
        
        try:
            import botocore
            dependencies['botocore'] = True
            logger.info("botocore module is available")
        except ImportError as e:
            dependencies['botocore'] = False
            logger.error("botocore module not available: %s", e)
        
        try:
            from botocore.exceptions import ClientError
            dependencies['botocore.exceptions'] = True
            logger.info("botocore.exceptions module is available")
        except ImportError as e:
            dependencies['botocore.exceptions'] = False
            logger.error("botocore.exceptions module not available: %s", e)
        
        return dependencies
    
    def recover_aws_dependencies(self) -> bool:
        """
        Attempt to recover AWS dependencies by reinstalling them.
        This is a fallback when bundled dependencies are missing.
        
        Returns:
            True if recovery was successful, False otherwise
        """
        try:
            import subprocess
            import sys
            
            logger.info("Attempting to recover AWS dependencies...")
            
            # Try to install boto3 and botocore
            packages = ['boto3', 'botocore']
            for package in packages:
                try:
                    logger.info("Installing %s...", package)
                    subprocess.check_call([
                        sys.executable, '-m', 'pip', 'install', 
                        '--user', '--upgrade', package
                    ], timeout=300)  # 5 minute timeout
                    logger.info("Successfully installed %s", package)
                except subprocess.TimeoutExpired:
                    logger.error("Timeout installing %s", package)
                    return False
                except subprocess.CalledProcessError as e:
                    logger.error("Failed to install %s: %s", package, e)
                    return False
            
            # Verify installation
            dependencies = self.check_aws_dependencies()
            if all(dependencies.values()):
                logger.info("AWS dependencies recovery successful")
                return True
            else:
                logger.error("AWS dependencies recovery failed - some modules still missing")
                return False
                
        except Exception as e:
            logger.error("Error during AWS dependencies recovery: %s", e)
            return False
    
    def check_and_recover_all(self) -> Dict[str, bool]:
        """
        Check for missing resources and attempt recovery.
        
        Returns:
            Dictionary mapping resource names to recovery success status
        """
        results = {}

        # Check AWS resources
        aws_path = self.recover_aws_resources()
        results['aws_resources'] = aws_path is not None
        
        # Check AWS dependencies
        aws_deps = self.check_aws_dependencies()
        results['aws_dependencies'] = all(aws_deps.values())
        
        # If AWS dependencies are missing, try to recover them
        if not results['aws_dependencies']:
            logger.warning("AWS dependencies missing, attempting recovery...")
            recovery_success = self.recover_aws_dependencies()
            results['aws_dependencies_recovery'] = recovery_success
            
            # Re-check after recovery attempt
            if recovery_success:
                aws_deps_after = self.check_aws_dependencies()
                results['aws_dependencies'] = all(aws_deps_after.values())
        
        # Check SSL certificates
        ssl_certs = self.check_ssl_certificates()
        results['ssl_certificates'] = ssl_certs.get('default_context', False)
        
        # If SSL certificates are missing, try to recover them
        if not results['ssl_certificates']:
            logger.warning("SSL certificates missing, attempting recovery...")
            recovery_success = self.recover_ssl_certificates()
            results['ssl_certificates_recovery'] = recovery_success
            
            # Re-check after recovery attempt
            if recovery_success:
                ssl_certs_after = self.check_ssl_certificates()
                results['ssl_certificates'] = ssl_certs_after.get('default_context', False)
        
        # Log results (only log failures for non-optional resources)
        optional_resources = {'aws_resources'}
        for resource, success in results.items():
            if success:
                logger.info("Resource recovery successful: %s", resource)
            elif resource in optional_resources:
                # Don't log warnings for optional resources that are expected to be missing
                logger.debug("Optional resource not available (this is normal): %s", resource)
            else:
                logger.warning("Resource recovery failed: %s", resource)
        
        return results
    
    def cleanup_old_recoveries(self, max_age_days: int = 30):
        """
        Clean up old recovered resources to prevent disk space issues.
        
        Args:
            max_age_days: Maximum age in days for recovered resources
        """
        try:
            permanent_dir = self.get_permanent_dir()
            if not permanent_dir.exists():
                return
            
            import time
            current_time = time.time()
            max_age_seconds = max_age_days * 24 * 60 * 60
            
            for item in permanent_dir.iterdir():
                if item.is_file() or item.is_dir():
                    item_age = current_time - item.stat().st_mtime
                    if item_age > max_age_seconds:
                        if item.is_file():
                            item.unlink()
                        else:
                            shutil.rmtree(item)
                        logger.info("Cleaned up old resource: %s", item)
                        
        except Exception as e:
            logger.error("Error during cleanup: %s", e)

# Global instance
resource_recovery = ResourceRecovery()
