import sys
import importlib
from pathlib import Path
import zipfile
import tarfile
import hashlib
import filecmp
import mmap
from typing import Union, List, Dict, Set
from concurrent.futures import ThreadPoolExecutor
import os


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

    Features:
        - Compares files by hash, size, or content for accuracy
        - Supports modules from sys.modules or importable names
        - Supports folder paths or archives (ZIP/TAR variants)
        - Parallel processing for better performance
        - Property-based access to comparison results
    """

    def __init__(
        self,
        module_name: str,
        other_path: Union[str, Path],
        temp_dir: Path = Path(".tmp_compare"),
        use_hash: bool = True,
        parallel_processing: bool = True,
        max_workers: int = None,
    ):
        self.module_name = module_name
        self.other_path = Path(other_path).resolve()
        self.temp_dir = Path(temp_dir).resolve()
        self.use_hash = use_hash
        self.parallel_processing = parallel_processing
        self.max_workers = max_workers or os.cpu_count()

        # Ensure module is loaded
        if module_name not in sys.modules:
            importlib.import_module(module_name)
        mod = sys.modules[module_name]

        if not hasattr(mod, "__file__"):
            raise TypeError(
                f"Cannot locate source for built-in module '{module_name}'."
            )

        self.module_path = Path(mod.__file__).resolve().parent

        # Extract other if it's an archive
        if self.other_path.is_file() and self.other_path.suffix in (
            ".zip",
            ".tar",
            ".gz",
            ".bz2",
            ".xz",
        ):
            self.other_path = self._extract_archive(
                self.other_path, self.temp_dir / self.other_path.stem
            )

        # Collect file dictionaries
        self.mod_files = self._collect_files(self.module_path)
        self.other_files = self._collect_files(self.other_path)

        # Perform comparison
        self._perform_comparison()

    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
        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

    def _compare_files_sequential(self, common_keys: Set[Path]) -> Set[Path]:
        """Compare files sequentially."""
        diff_keys = set()
        for key in common_keys:
            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."""
        diff_keys = set()
        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            futures = {
                executor.submit(
                    self._files_equal, self.mod_files[key], self.other_files[key]
                ): key
                for key in common_keys
            }
            for future in futures:
                if not future.result():
                    diff_keys.add(futures[future])
        return diff_keys

    def _files_equal(self, file1: Path, file2: Path) -> bool:
        """Check if two files are equal using the configured strategy."""
        if self.use_hash:
            return self._hash_file(file1) == self._hash_file(file2)
        else:
            return self._compare_files_fast(file1, file2)

    @staticmethod
    def _hash_file(path: Path, algo="md5") -> str:
        """Calculate file hash efficiently."""
        h = hashlib.new(algo)
        try:
            with path.open("rb") as f:
                for chunk in iter(lambda: f.read(8192), b""):
                    h.update(chunk)
        except (IOError, OSError):
            return ""
        return h.hexdigest()

    @staticmethod
    def _compare_files_fast(file1: Path, file2: Path) -> bool:
        """Fast file comparison using multiple strategies."""
        # First check file size
        try:
            stat1 = file1.stat()
            stat2 = file2.stat()
            if stat1.st_size != stat2.st_size:
                return False

            # If sizes match, compare content
            return filecmp.cmp(file1, file2, shallow=False)
        except (OSError, IOError):
            return False

    def _extract_archive(self, archive_path: Path, extract_to: Path) -> Path:
        """Extract archive to temporary directory."""
        extract_to.mkdir(parents=True, exist_ok=True)
        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"):
            with tarfile.open(archive_path, "r:*") as tf:
                tf.extractall(extract_to)
        else:
            raise ValueError("Unsupported archive format.")
        return extract_to

    @staticmethod
    def _collect_files(base: Path) -> Dict[Path, Path]:
        """Return dict: relative_path -> absolute_path for all .py files."""
        files = {}
        for f in base.rglob("*.py"):
            try:
                files[f.relative_to(base)] = f
            except ValueError:
                continue
        return files

    # 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 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 self._right_only_keys]

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

    @property
    def common_files(self) -> List[str]:
        """Files present in both with identical content."""
        return [str(f) for f in 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) -> Dict[str, int]:
        """Statistics about the comparison."""
        return {
            "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),
        }

    def __repr__(self) -> str:
        return (
            f"ModuleComparator(module='{self.module_name}', "
            f"other='{self.other_path}', "
            f"files_comparison={self.stats})"
        )
