from __future__ import annotations

import hashlib
import io
import mimetypes
import os
from collections import defaultdict
from collections.abc import Iterator
from pathlib import Path

from PIL import Image

from .config import CONFIG
from .photo_file import PhotoFile
from .pipeline_stage import PipelineStage, PrepareResult, WorkerResult
from .ports import OutputPort
from .template_parsing import extract_template


def _get_oriented_dimensions(img: Image.Image) -> tuple[int, int]:
    """Extract dimensions from PIL Image with EXIF orientation applied.

    Args:
        img: Opened PIL Image

    Returns:
        Tuple of (width, height) with EXIF orientation applied
    """
    # Get EXIF orientation if present
    orientation: int = 0
    if hasattr(img, "_getexif") and img._getexif() is not None:
        exif = img._getexif()
        orientation = exif.get(274, 0)  # 274 = Orientation EXIF tag

    # Get raw dimensions
    raw_width: int = img.width
    raw_height: int = img.height

    # Apply EXIF orientation (swap dimensions for rotations 5,6,7,8)
    # https://www.impulseadventure.com/photo/exif-orientation.html
    if orientation in {5, 6, 7, 8}:
        return raw_height, raw_width
    return raw_width, raw_height


class ComputeShaBins(PipelineStage[tuple[int, tuple[Path, str]], tuple[PhotoFile, str], dict[str, list[PhotoFile]]]):
    """Pipeline stage that walks source directory and bins photos by SHA256 hash.

    Creates PhotoFile objects with full metadata extracted from a single file read:
    - SHA256 hash (for binning, then discarded)
    - File size, MIME type, path
    - Image dimensions with EXIF orientation applied

    PhotoFile is a pure data container - this stage performs ALL file I/O.
    """

    def __init__(self, source_path: Path) -> None:
        """Initialize SHA256 binning stage.

        Args:
            source_path: Root directory to scan for images
        """
        super().__init__(
            path=CONFIG.paths.sha_bins_pkl,
            stage_name="Directory Walk",
        )
        # Store input
        self.source_path = source_path
        self.args = source_path

        # Create output ports
        self.sha_bins_o: OutputPort[dict[str, list[PhotoFile]]] = OutputPort(self, getter=lambda: self.result)
        self.photofiles_o: OutputPort[dict[int, PhotoFile]] = OutputPort(self, getter=lambda: self.photofiles)

    def prepare(
        self,
    ) -> PrepareResult[tuple[int, tuple[Path, str]], dict[str, list[PhotoFile]]]:
        """Prepare source directory for processing.

        Walks directory tree lazily to find image files and creates work items.

        Returns:
            Tuple of (enumerated file paths, empty SHA256 bins accumulator)
        """
        self.ref_photos_init = None
        self.ref_seqs_init = None

        # Use stored source_path
        path = self.source_path

        def walk_image_files(root: Path) -> Iterator[tuple[Path, str]]:
            """Walk directory tree and yield image file paths with mime types.

            Args:
                    root: Root directory to walk

            Yields:
                    Tuples of (file_path, mime_subtype) for each image file
                    Example: (Path("/photos/img.jpg"), "jpeg")
            """
            dp: str
            fns: list[str]
            for dp, _, fns in os.walk(root):
                fn: str
                for fn in fns:
                    mime: str | None
                    mime, _ = mimetypes.guess_type(fn)
                    if mime is not None and mime.startswith("image/"):
                        p: Path = Path(dp) / fn
                        yield p, mime.removeprefix("image/")

        return enumerate(walk_image_files(path)), defaultdict(list)

    @classmethod
    def stage_worker(cls, param: tuple[int, tuple[Path, str]], _args: str) -> WorkerResult[tuple[PhotoFile, str]]:
        """Create PhotoFile with dimensions and compute SHA256 in single file read.

        Work function for parallel processing that reads file once, computes SHA256,
        extracts dimensions with EXIF orientation, and creates PhotoFile with all
        metadata. PhotoFile.__init__ never opens files - this is the only file I/O.

        There is no exception handling in here. All exceptions should be surfaced
        to be dealt with by the user.

        Args:
                param: (photo_id, (path, mime)) tuple
                _args: Placeholder to match pattern

        Returns:
                (PhotoFile with all metadata, SHA256 hash) tuple
        """
        photo_id: int
        path: Path
        mime: str
        photo_id, (path, mime) = param

        # Read file once into memory for both SHA256 and dimensions (optimization)
        with path.open("rb") as f:
            file_data: bytes = f.read()

        # Compute SHA256 from in-memory data
        sha256_hash: str = hashlib.sha256(file_data).hexdigest()

        # Extract dimensions with EXIF orientation from in-memory data
        with Image.open(io.BytesIO(file_data)) as img:
            width, height = _get_oriented_dimensions(img)

        # Get file size from data length (avoids separate stat call)
        size_bytes: int = len(file_data)

        # Create PhotoFile with all metadata (no file I/O in PhotoFile.__init__)
        photo = PhotoFile(
            path=path,
            mime=mime,
            size_bytes=size_bytes,
            file_id=photo_id,
            width=width,
            height=height,
        )

        # Extract template from filename (with extension) and include full directory path
        # This ensures files with same name in different directories have different templates
        template_pattern, index = extract_template(path.name)
        template_with_path = str(path.with_name(template_pattern))
        photo.cache["TEMPLATE"] = (template_with_path, index)

        # Return PhotoFile and SHA256 separately (SHA256 used for binning only)
        return (
            [],
            [],
            (photo, sha256_hash),
        )

    def accumulate_results(self, result: dict[str, list[PhotoFile]], job: tuple[PhotoFile, str]) -> None:
        """Add PhotoFile to appropriate SHA256 bin.

        Args:
            result: Dictionary mapping SHA256 hash to list of PhotoFile objects
            job: Tuple of (PhotoFile, sha256) from stage_worker
        """
        photo, sha256 = job
        result[sha256].append(photo)

    def finalise(self) -> None:
        self.ref_photos_final = sum(len(b) for b in self.result.values())
        self.ref_seqs_final = None

    @property
    def photofiles(self) -> dict[int, PhotoFile]:
        """Property that returns the photofile dictionary mapping id to structure."""
        return {p.id: p for b in self.result.values() for p in b}
