from types import ModuleType as module
from typing import Union, List, Tuple
from pathlib import Path
import os
import importlib
from concurrent.futures import ThreadPoolExecutor


def _files_(path: Path) -> List[Path]:
    """
    Get all files recursively in the given directory path.

    Args:
            path (Path): A pathinit.Path object representing the directory
                         where files are to be listed.

    Returns:
            List[Path]: A list of pathinit.Path objects representing all files
                        in the directory and its subdirectories.

    Notes:
            - Uses Path.rglob('*') to recursively traverse directories.
            - Returns all entries (files and folders), filtering for files
              can be done later if needed.
            - Useful for modules/packages where we want to analyze file sizes
              or perform file-based searches.
    """
    return list(path.rglob("*"))


def _file_size(file: Path) -> int:
    """
    Get the size of a single file safely.

    Args:
            file (Path): A pathinit.Path object representing the file.

    Returns:
            int: Size of the file in bytes.
                 Returns 0 if the file is not accessible or does not exist.

    Notes:
            - Uses file.stat().st_size to get file size.
            - Wraps in try-except to handle:
                    * FileNotFoundError: in case the file was deleted or moved.
                    * PermissionError: in case access is denied.
            - This function ensures that large batch operations don't break
              due to inaccessible files.
    """
    try:
        return file.stat().st_size
    except (FileNotFoundError, PermissionError):
        return 0


class size:
    """
    Class to represent the size and files of a Python module or package.

    This class provides:
    - The total size in bytes of all files in the module/package.
    - The ability to search for files based on size criteria.
    - Options to return file names or full paths.
    - Multi-threaded processing for faster performance on large directories.

    Attributes:
            init (module): The imported Python module object.
            path (Path): The root directory path of the module/package.

    Methods:
            - size: Property that returns total size in bytes.
            - find: Search for files based on size and comparison operators.
    """

    def __init__(
        self, ModuleName: Union[str, module], worker: int = os.cpu_count() * 5
    ) -> None:
        """
        Initialize the size object with a module or module name.

        Args:
                init (Union[str, module]): Either a Python module object (imported)
                                           or a string representing the module name.

        Behavior:
                1. Determines module name if a module object is passed.
                2. Imports the module dynamically if only a string is passed.
                3. Extracts the path of the module by removing the module's filename
                   from the full path.
                4. Stores both the imported module and its path for further processing.

        Example:
                sz = size("os")
                print(sz.path)  # /usr/init/python3.11/os
        """
        name = init.__name__ if isinstance(ModuleName, module) else ModuleName
        self.init = importlib.import_module(ModuleName)
        self.file = self.init.__file__
        self.path = Path(self.file).resolve().parent  # Directory path
        self.worker = worker

    def __str__(self) -> str:
        """
        Return a human-readable representation of the object.

        Returns:
                str: A string containing the module name, total size, and path.

        Example:
                sz = size("os")
                print(sz)
                # Output: Size(size='123456',initname='os', path='/usr/init/python3.11/os')

        Notes:
                - Accesses self.size property dynamically.
                - Shows current total size of all files in the module/package.
        """
        return f"size(size='{self.size}',initname='{self.init.__name__}', path='{self.path}')"

    @property
    def size(self) -> int:
        """
        Calculate total size of all files in the module/package directory.

        Returns:
                int: Total size in bytes of all files.

        Behavior:
                1. Recursively collects all files using Path.rglob('*').
                2. Warns if the total number of files exceeds 2000 to alert
                   user about potential time consumption.
                3. Filters only files (ignores directories) using a generator.
                4. Calculates file sizes concurrently using ThreadPoolExecutor
                   for faster computation on large directories.

        Notes:
                - ThreadPoolExecutor automatically scales to the number of available threads.
                - _file_size ensures safe handling of inaccessible files.
                - Total size is computed as the sum of individual file sizes.

        Example:
                sz = size("os")
                total = sz.size
                print(f"Total size: {total} bytes")
        """
        files = list(self.path.rglob("*"))
        if len(files) >= 2000:
            print(
                "<Note:: The sum of the files is more than 2000 files, so the process may take time>\n"
            )
        with ThreadPoolExecutor(max_workers=self.worker) as executor:
            total_size = sum(
                executor.map(_file_size, (f for f in files if f.is_file()))
            )
        return total_size

    def find(
        self,
        size: int = None,
        maxfiles: int = None,
        cmp: str = "<=",
        withsize: bool = False,
        fullname: bool = False,
        ignore: List[str] = [],
    ) -> List[Union[str, Tuple[str, int]]]:
        """
        Find files in the module/package directory based on size comparison.

        Args:
                size (int, optional): Size in bytes to compare each file against.
                maxfiles (int, optional): Maximum number of files to process.
                cmp (str, optional): Comparison operator, one of:
                                      '==', '>=', '<=', '>', '<'.
                                      Default is '<='.
                withsize (bool, optional): If True, returns tuples (filename, size),
                                           otherwise returns just filenames. Default False.
                fullname (bool, optional): If True, returns full file paths,
                                           otherwise just the filename. Default False.
                ignore (List[str], optional): List of file suffixes to ignore. Default [].

        Returns:
                List[Union[str, Tuple[str, int]]]: A list of matching files or (filename, size) tuples.

        Behavior:
                1. Collects all files in the module/package directory (optionally limited by maxfiles).
                2. Iterates through each file, ignoring specified suffixes.
                3. Calculates file size safely using _file_size.
                4. Uses Python 3.10+ match/case statements to apply comparison operator.
                5. Returns list of files meeting the comparison criteria.
                6. Uses ThreadPoolExecutor to accelerate computation on large directories.

        Notes:
                - File paths are converted to Path objects internally for consistency.
                - match/case ensures clean and readable comparison logic.
                - Safe handling of inaccessible files prevents exceptions from breaking iteration.
                - maxfiles allows partial scanning for extremely large modules.

        Examples:
                sz = size("os")
                # Find all files >= 1KB
                files_large = sz.find(size=1024, cmp=">=", withsize=True)

                # Find first 50 files < 100 bytes
                files_small = sz.find(size=100, cmp="<", maxfiles=50)
        """
        files = _files_(self.path)[:maxfiles]
        results = []

        def process_file(file: Path):
            if ignore and file.suffix in ignore:
                return None

            name = str(file) if fullname else file.name
            fsize = _file_size(file)
            item = (name, fsize) if withsize else name

            match cmp:
                case "==":
                    if fsize == size:
                        return item
                case ">=":
                    if fsize >= size:
                        return item
                case "<=":
                    if fsize <= size:
                        return item
                case ">":
                    if fsize > size:
                        return item
                case "<":
                    if fsize < size:
                        return item
                case _:
                    raise ValueError(f"Invalid comparison operator '{cmp}'")
            return None

        with ThreadPoolExecutor(max_workers=self.worker) as executor:
            for res in executor.map(process_file, (Path(f) for f in files)):
                if res is not None:
                    results.append(res)

        return results
