import sys
import importlib
from pathlib import Path
import zipfile
import tarfile
import hashlib
import filecmp
import mmap
import os
import tempfile
import shutil
from typing import Union, List, Dict, Set, Optional, Tuple, Any
from concurrent.futures import ThreadPoolExecutor, as_completed
from contextlib import contextmanager
import time
import logging
from dataclasses import dataclass
from enum import Enum
import fnmatch


class ComparisonMethod(Enum):
    """Methods available for file comparison."""
    HASH = "hash"
    SIZE = "size"
    CONTENT = "content"
    MTIME = "mtime"


@dataclass
class ComparisonStats:
    """Statistics about the comparison results."""
    total_module_files: int = 0
    total_other_files: int = 0
    common_files: int = 0
    different_files: int = 0
    module_only: int = 0
    other_only: int = 0
    comparison_time: float = 0.0
    memory_used_mb: float = 0.0


class FileComparisonError(Exception):
    """Custom exception for file comparison errors."""
    pass


class ModuleComparator:
    """
    Compare a Python module with another folder or archive.

    Features:
        - Multiple comparison methods: hash, size, content, mtime
        - Support for file filtering and patterns
        - Progress tracking and detailed logging
        - Memory-efficient large file handling
        - Checksum caching for repeated comparisons
        - Detailed reporting and export capabilities
    """

    def __init__(
        self,
        module_name: str,
        other_path: Union[str, Path],
        temp_dir: Optional[Path] = None,
        comparison_method: Union[ComparisonMethod, str] = ComparisonMethod.HASH,
        parallel_processing: bool = True,
        max_workers: Optional[int] = None,
        hash_algorithm: str = "md5",
        chunk_size: int = 8192,
        include_patterns: Optional[List[str]] = None,
        exclude_patterns: Optional[List[str]] = None,
        follow_symlinks: bool = False,
        enable_caching: bool = True,
        log_level: int = logging.INFO,
    ):
        self._start_time = time.time()
        self._setup_logging(log_level)
        
        self.module_name = module_name
        self.other_path = Path(other_path).resolve()
        self.temp_dir = temp_dir or Path(tempfile.gettempdir()) / "module_comparator"
        self.comparison_method = self._parse_comparison_method(comparison_method)
        self.parallel_processing = parallel_processing
        self.max_workers = max_workers or min(32, (os.cpu_count() or 1) + 4)
        self.hash_algorithm = hash_algorithm
        self.chunk_size = chunk_size
        self.include_patterns = include_patterns or ["*"]
        self.exclude_patterns = exclude_patterns or []
        self.follow_symlinks = follow_symlinks
        self.enable_caching = enable_caching
        
        # Cache for file hashes and metadata
        self._hash_cache: Dict[Path, str] = {}
        self._metadata_cache: Dict[Path, Dict[str, Any]] = {}
        
        # Initialize paths and perform comparison
        self._initialize_module()
        self._setup_temp_dir()
        self._extract_if_archive()
        
        # Collect and compare files
        self.mod_files = self._collect_files(self.module_path)
        self.other_files = self._collect_files(self.other_path)
        
        self._perform_comparison()
        self._cleanup_temp_dir()

    def _setup_logging(self, level: int):
        """Setup logging configuration."""
        self.logger = logging.getLogger(f"ModuleComparator_{id(self)}")
        if not self.logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            handler.setFormatter(formatter)
            self.logger.addHandler(handler)
        self.logger.setLevel(level)

    def _parse_comparison_method(self, method: Union[ComparisonMethod, str]) -> ComparisonMethod:
        """Parse and validate comparison method."""
        if isinstance(method, ComparisonMethod):
            return method
        try:
            return ComparisonMethod(method.lower())
        except ValueError:
            raise ValueError(
                f"Invalid comparison method: {method}. "
                f"Available methods: {[m.value for m in ComparisonMethod]}"
            )

    def _initialize_module(self):
        """Load and validate the module."""
        if self.module_name not in sys.modules:
            try:
                importlib.import_module(self.module_name)
            except ImportError as e:
                raise FileComparisonError(
                    f"Failed to import module '{self.module_name}': {e}"
                )
        
        mod = sys.modules[self.module_name]
        
        if not hasattr(mod, "__file__") or not mod.__file__:
            raise FileComparisonError(
                f"Cannot locate source for module '{self.module_name}'."
            )
        
        self.module_path = Path(mod.__file__).resolve().parent
        self.logger.info(f"Module path: {self.module_path}")

    def _setup_temp_dir(self):
        """Create and setup temporary directory."""
        self.temp_dir.mkdir(parents=True, exist_ok=True)
        self.logger.debug(f"Using temp directory: {self.temp_dir}")

    def _cleanup_temp_dir(self):
        """Cleanup temporary directory if it's empty."""
        try:
            if self.temp_dir.exists() and not any(self.temp_dir.iterdir()):
                self.temp_dir.rmdir()
        except OSError as e:
            self.logger.warning(f"Failed to cleanup temp directory: {e}")

    def _extract_if_archive(self):
        """Extract other_path if it's an archive file."""
        if self.other_path.is_file():
            archive_extensions = {
                '.zip', '.tar', '.gz', '.bz2', '.xz', '.tgz', '.tbz2'
            }
            if self.other_path.suffix.lower() in archive_extensions:
                extract_to = self.temp_dir / f"extracted_{self.other_path.stem}"
                self.logger.info(f"Extracting archive to: {extract_to}")
                self.other_path = self._extract_archive(self.other_path, extract_to)

    def _should_include_file(self, file_path: Path) -> bool:
        """Check if file should be included based on patterns."""
        rel_path = str(file_path)
        
        # Check include patterns
        included = any(fnmatch.fnmatch(rel_path, pattern) 
                      for pattern in self.include_patterns)
        if not included:
            return False
        
        # Check exclude patterns
        excluded = any(fnmatch.fnmatch(rel_path, pattern) 
                      for pattern in self.exclude_patterns)
        return not excluded

    def _collect_files(self, base: Path) -> Dict[Path, Path]:
        """Collect files with filtering and pattern matching."""
        files = {}
        
        try:
            for pattern in self.include_patterns:
                for file_path in base.rglob(pattern):
                    if (file_path.is_file() and 
                        not file_path.is_symlink() and  # Skip symlinks unless following
                        self._should_include_file(file_path.relative_to(base))):
                        try:
                            rel_path = file_path.relative_to(base)
                            files[rel_path] = file_path
                        except ValueError:
                            continue
        except Exception as e:
            self.logger.error(f"Error collecting files from {base}: {e}")
            raise FileComparisonError(f"Failed to collect files: {e}")
        
        self.logger.info(f"Collected {len(files)} files from {base}")
        return files

    def _extract_archive(self, archive_path: Path, extract_to: Path) -> Path:
        """Extract archive to temporary directory."""
        if extract_to.exists():
            shutil.rmtree(extract_to)
        
        extract_to.mkdir(parents=True, exist_ok=True)
        
        try:
            if archive_path.suffix == ".zip":
                with zipfile.ZipFile(archive_path, "r") as zf:
                    zf.extractall(extract_to)
            elif archive_path.suffix in (".tar", ".gz", ".bz2", ".xz", ".tgz", ".tbz2"):
                with tarfile.open(archive_path, "r:*") as tf:
                    tf.extractall(extract_to)
            else:
                raise ValueError(f"Unsupported archive format: {archive_path.suffix}")
        except Exception as e:
            self.logger.error(f"Failed to extract archive {archive_path}: {e}")
            raise FileComparisonError(f"Archive extraction failed: {e}")
        
        self.logger.info(f"Successfully extracted {archive_path} to {extract_to}")
        return extract_to

    def _perform_comparison(self):
        """Perform all comparisons and cache results."""
        mod_keys = set(self.mod_files.keys())
        other_keys = set(self.other_files.keys())

        common_keys = mod_keys & other_keys
        self._left_only_keys = mod_keys - other_keys
        self._right_only_keys = other_keys - mod_keys

        # Compare common files
        compare_start = time.time()
        
        if self.parallel_processing:
            self._diff_keys = self._compare_files_parallel(common_keys)
        else:
            self._diff_keys = self._compare_files_sequential(common_keys)

        self._common_keys = common_keys - self._diff_keys
        self._comparison_time = time.time() - compare_start
        
        self.logger.info(f"Comparison completed in {self._comparison_time:.2f} seconds")

    def _compare_files_sequential(self, common_keys: Set[Path]) -> Set[Path]:
        """Compare files sequentially."""
        diff_keys = set()
        total = len(common_keys)
        
        for i, key in enumerate(common_keys):
            if i % 100 == 0:
                self.logger.debug(f"Compared {i}/{total} files...")
            
            if not self._files_equal(self.mod_files[key], self.other_files[key]):
                diff_keys.add(key)
        
        return diff_keys

    def _compare_files_parallel(self, common_keys: Set[Path]) -> Set[Path]:
        """Compare files in parallel with progress tracking."""
        diff_keys = set()
        total = len(common_keys)
        
        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            # Submit all comparison tasks
            future_to_key = {
                executor.submit(
                    self._files_equal, 
                    self.mod_files[key], 
                    self.other_files[key]
                ): key for key in common_keys
            }
            
            # Process completed tasks
            for i, future in enumerate(as_completed(future_to_key)):
                if i % 100 == 0:
                    self.logger.debug(f"Compared {i}/{total} files...")
                
                key = future_to_key[future]
                try:
                    if not future.result():
                        diff_keys.add(key)
                except Exception as e:
                    self.logger.error(f"Error comparing {key}: {e}")
                    diff_keys.add(key)  # Treat errors as differences
        
        return diff_keys

    def _files_equal(self, file1: Path, file2: Path) -> bool:
        """Check if two files are equal using the configured method."""
        try:
            if self.comparison_method == ComparisonMethod.HASH:
                return self._get_file_hash(file1) == self._get_file_hash(file2)
            elif self.comparison_method == ComparisonMethod.SIZE:
                return self._get_file_size(file1) == self._get_file_size(file2)
            elif self.comparison_method == ComparisonMethod.CONTENT:
                return self._compare_files_content(file1, file2)
            elif self.comparison_method == ComparisonMethod.MTIME:
                return self._get_file_mtime(file1) == self._get_file_mtime(file2)
            else:
                raise ValueError(f"Unknown comparison method: {self.comparison_method}")
        except (OSError, IOError) as e:
            self.logger.warning(f"Error comparing {file1} and {file2}: {e}")
            return False

    def _get_file_hash(self, file_path: Path) -> str:
        """Get file hash with caching."""
        if self.enable_caching and file_path in self._hash_cache:
            return self._hash_cache[file_path]
        
        try:
            hasher = hashlib.new(self.hash_algorithm)
            with file_path.open("rb") as f:
                for chunk in iter(lambda: f.read(self.chunk_size), b""):
                    hasher.update(chunk)
            hash_value = hasher.hexdigest()
            
            if self.enable_caching:
                self._hash_cache[file_path] = hash_value
            
            return hash_value
        except (IOError, OSError) as e:
            self.logger.error(f"Error hashing file {file_path}: {e}")
            return ""

    def _get_file_size(self, file_path: Path) -> int:
        """Get file size."""
        try:
            return file_path.stat().st_size
        except (OSError, IOError):
            return -1

    def _get_file_mtime(self, file_path: Path) -> float:
        """Get file modification time."""
        try:
            return file_path.stat().st_mtime
        except (OSError, IOError):
            return -1

    def _compare_files_content(self, file1: Path, file2: Path) -> bool:
        """Compare files content using memory mapping for large files."""
        try:
            # Quick size check
            if self._get_file_size(file1) != self._get_file_size(file2):
                return False
            
            # Use filecmp for reliable comparison
            return filecmp.cmp(file1, file2, shallow=False)
        except (OSError, IOError):
            return False

    @contextmanager
    def _memory_map_file(self, file_path: Path):
        """Context manager for memory mapping large files."""
        try:
            with open(file_path, "rb") as f:
                with mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as mm:
                    yield mm
        except (OSError, IOError, ValueError) as e:
            self.logger.warning(f"Memory mapping failed for {file_path}: {e}")
            raise

    def _get_memory_usage(self) -> float:
        """Get current memory usage in MB."""
        try:
            import psutil
            process = psutil.Process()
            return process.memory_info().rss / 1024 / 1024
        except ImportError:
            return 0.0

    # Properties for comparison results
    @property
    def left_only(self) -> List[str]:
        """Files present in module but missing in other."""
        return [str(f) for f in sorted(self._left_only_keys)]

    @property
    def right_only(self) -> List[str]:
        """Files present in other but missing in module."""
        return [str(f) for f in sorted(self._right_only_keys)]

    @property
    def diff_files(self) -> List[str]:
        """Files present in both but different content."""
        return [str(f) for f in sorted(self._diff_keys)]

    @property
    def common_files(self) -> List[str]:
        """Files present in both with identical content."""
        return [str(f) for f in sorted(self._common_keys)]

    @property
    def all_files(self) -> Dict[str, List[str]]:
        """All comparison results in a dictionary."""
        return {
            "left_only": self.left_only,
            "right_only": self.right_only,
            "diff_files": self.diff_files,
            "common_files": self.common_files,
        }

    @property
    def stats(self) -> ComparisonStats:
        """Detailed statistics about the comparison."""
        return ComparisonStats(
            total_module_files=len(self.mod_files),
            total_other_files=len(self.other_files),
            common_files=len(self.common_files),
            different_files=len(self.diff_files),
            module_only=len(self.left_only),
            other_only=len(self.right_only),
            comparison_time=self._comparison_time,
            memory_used_mb=self._get_memory_usage(),
        )

    def get_detailed_differences(self) -> Dict[str, List[Dict[str, Any]]]:
        """Get detailed information about differences."""
        details = {
            "left_only": [],
            "right_only": [],
            "diff_files": []
        }
        
        for file_path in self._left_only_keys:
            abs_path = self.mod_files[file_path]
            details["left_only"].append({
                "path": str(file_path),
                "size": self._get_file_size(abs_path),
                "mtime": self._get_file_mtime(abs_path)
            })
        
        for file_path in self._right_only_keys:
            abs_path = self.other_files[file_path]
            details["right_only"].append({
                "path": str(file_path),
                "size": self._get_file_size(abs_path),
                "mtime": self._get_file_mtime(abs_path)
            })
        
        for file_path in self._diff_keys:
            mod_path = self.mod_files[file_path]
            other_path = self.other_files[file_path]
            details["diff_files"].append({
                "path": str(file_path),
                "module_size": self._get_file_size(mod_path),
                "other_size": self._get_file_size(other_path),
                "module_mtime": self._get_file_mtime(mod_path),
                "other_mtime": self._get_file_mtime(other_path),
                "module_hash": self._get_file_hash(mod_path),
                "other_hash": self._get_file_hash(other_path)
            })
        
        return details

    def export_report(self, output_path: Path, format: str = "json") -> None:
        """Export comparison report to file."""
        report = {
            "module_name": self.module_name,
            "other_path": str(self.other_path),
            "comparison_method": self.comparison_method.value,
            "timestamp": time.time(),
            "stats": self.stats.__dict__,
            "results": self.all_files,
            "detailed_differences": self.get_detailed_differences()
        }
        
        output_path = Path(output_path)
        
        if format.lower() == "json":
            import json
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(report, f, indent=2, ensure_ascii=False)
        else:
            raise ValueError(f"Unsupported format: {format}")

    def is_identical(self) -> bool:
        """Check if the modules are identical."""
        return (len(self._left_only_keys) == 0 and 
                len(self._right_only_keys) == 0 and 
                len(self._diff_keys) == 0)

    def get_similarity_score(self) -> float:
        """Calculate similarity score between modules (0.0 to 1.0)."""
        total_files = len(self.mod_files) + len(self.other_files)
        if total_files == 0:
            return 1.0
        
        common_count = len(self._common_keys)
        return (2 * common_count) / total_files

    def __repr__(self) -> str:
        identical = "IDENTICAL" if self.is_identical() else "DIFFERENT"
        return (
            f"ModuleComparator(module='{self.module_name}', "
            f"other='{self.other_path}', "
            f"status='{identical}', "
            f"similarity={self.get_similarity_score():.2f})"
        )

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Cleanup resources when used as context manager."""
        self._cleanup_temp_dir()


# Example usage and helper functions
def compare_modules(
    module_name: str,
    other_path: Union[str, Path],
    **kwargs
) -> ModuleComparator:
    """Convenience function for quick module comparison."""
    return ModuleComparator(module_name, other_path, **kwargs)


if __name__ == "__main__":
    # Example usage
    with ModuleComparator(
        module_name="json",
        other_path="importer-core",
        comparison_method=ComparisonMethod.HASH,
        parallel_processing=True,
        log_level=logging.INFO
    ) as comparator:
        
        print(f"Comparison results: {comparator.stats}")
        print(f"Similarity score: {comparator.get_similarity_score():.2f}")
        print(f"Identical: {comparator.is_identical()}")
        
        # Export report
        comparator.export_report("comparison_report.json")