"""Pipeline stage that consumes SHA256 bins and outputs groups of identical files and a list of unique exemplar files."""

from __future__ import annotations

import random

from .config import CONFIG
from .models import IdenticalGroup, ReviewType
from .photo_file import (
    PhotoFile,
    pick_exemplar_from_class,
)
from .pipeline_stage import PipelineStage, PrepareResult, WorkerResult
from .ports import InputPort, OutputPort
from .review_utils import build_identical_group
from .template_parsing import INDEX_T


class ComputeIdentical(
    PipelineStage[
        list[PhotoFile],  # WorkItem: SHA bins
        list[PhotoFile],  # Accumulator: flat list of exemplars during processing
        dict[str, list[tuple[INDEX_T, PhotoFile]]],  # Result: template bins after finalise()
    ]
):
    """Pipeline stage that detects byte-identical files and outputs template bins.

    This stage consumes SHA256 bins, identifies byte-identical files within each bin,
    picks exemplars from non-identical files, and bins the exemplars by filename template.

    Input:
        SHA256 bins (from ComputeSHABins)

    Output:
        Template bins: dict mapping template patterns to (index, photo) tuples
        Example: {"IMG_{P0}.jpg": [("1234", photo1), ("5678", photo2)]}

    Review data:
        Identical photo groups for user review
    """

    # Typed result field - populated after run() completes
    result: dict[str, list[tuple[INDEX_T, PhotoFile]]]

    def __init__(self) -> None:
        """Initialize identical files detection stage."""
        super().__init__(
            path=CONFIG.paths.identical_pkl,
            stage_name="Byte-identical detection",
        )

        # Worker args
        self.args = None

        # Create input port for SHA bins (from ComputeShaBins)
        self.sha_bins_i: InputPort[dict[str, list[PhotoFile]]] = InputPort("sha_bins")

        # Create output ports
        # - nonidentical_o: template bins for next stage (versions)
        self.nonidentical_o: OutputPort[dict[str, list[tuple[INDEX_T, PhotoFile]]]] = OutputPort(
            self, getter=lambda: self.result
        )

    def prepare(
        self,
    ) -> PrepareResult[list[PhotoFile], dict[str, list[tuple[INDEX_T, PhotoFile]]]]:
        """Prepare identical file detection by splitting bins into work items.

        Reads SHA bins from input port and prepares work items for parallel processing.

        Returns:
            Tuple of (work_items, accumulator) where:
            - work_items: List of bins with multiple photos (potential duplicates)
            - accumulator: Empty template bins dict (populated during accumulation)
        """
        # Read SHA bins from input port
        sha_bins: dict[str, list[PhotoFile]] = self.sha_bins_i.read()
        # Get reference counts from upstream (for photos without sequences, ref == total)
        self.ref_photos_init = self.sha_bins_i.get_ref_photo_count()
        self.ref_seqs_init = self.sha_bins_i.get_ref_sequence_count()
        # Count total photos for internal invariant checking (ref == total at this stage)
        self.total_photos = sum(len(b) for b in sha_bins.values())

        multiple_bins: list[list[PhotoFile]] = [b for b in sha_bins.values() if len(b) > 1]
        singleton_bins: list[list[PhotoFile]] = [b for b in sha_bins.values() if len(b) == 1]

        # Pre-bin singleton exemplars by template (optimization)
        template_bins: dict[str, list[tuple[INDEX_T, PhotoFile]]] = {}
        for bin in singleton_bins:
            photo = bin[0]
            assert "TEMPLATE" in photo.cache, "Template must be cached by SHA stage"
            template_bins.setdefault(photo.template, []).append((photo.template_index, photo))

        return multiple_bins, template_bins

    def finalise(self) -> None:
        """Update reference counts after template binning.

        Template binning already happened during accumulation, so this just
        computes the final reference counts.
        """
        # self.result is already a dict[str, list[tuple[INDEX_T, PhotoFile]]] from accumulator
        template_bins = self.result

        # Update reference counting
        self.ref_photos_final = sum(len(photos) for photos in template_bins.values())
        self.ref_seqs_final = len(template_bins)  # Number of unique templates

        # Existing invariant check (photos count unchanged)
        photos_final: int = (
            sum(len(cl.photos) for cl in self.identical_review_result)
            - len(self.identical_review_result)
            + self.ref_photos_final
        )

        assert photos_final == self.total_photos, (
            f"ComputeIdentical started with {self.total_photos} photos and ended up with {photos_final}"
        )

        # Shuffle review groups for variety in review UI (byte-identical groups are 100% confident)
        random.shuffle(self.identical_review_result)

    @classmethod
    def stage_worker(cls, photo_list: list[PhotoFile], _args: str) -> WorkerResult[list[PhotoFile]]:
        """Process one SHA bin to find byte-identical files.

        Creates both review data (IdenticalGroups) and working data (classes/exemplars)
        from a single SHA bin.

        Args:
            photo_list: Photos from one SHA256 bin
            _args: Stage name (unused, required to match PipelineStage interface)

        Returns:
            Tuple of (review_data, work_data) where:
            - review_data: (list[IdenticalGroup], dict[int, PhotoFile]) for review UI
            - work_data: (list[list[PhotoFile]], list[PhotoFile]) for pipeline flow
        """
        # Singleton bins are filtered out by prepare()
        assert len(photo_list) >= 2

        # If skipping byte-identical check, trust SHA256 uniqueness
        # Treat all photos in the bin as identical (same SHA = identical files)
        if CONFIG.processing.SKIP_BYTE_IDENTICAL:
            # Pick best exemplar using same logic as non-skip path
            # (prefer higher resolution, larger file size, stable path/ID tiebreaker)
            photos_dict: dict[int, PhotoFile] = {pf.id: pf for pf in photo_list}
            photo_ids: set[int] = set(photos_dict.keys())
            exemplar_id: int = pick_exemplar_from_class(photos_dict, photo_ids)
            exemplar: PhotoFile = photos_dict[exemplar_id]

            # Mark all other photos as identical to the exemplar
            for photo in photo_list:
                if photo.id != exemplar_id:
                    photo.cache["IDENTICAL"] = exemplar

            # Skip building review data - no need to review when trusting SHA256
            # (building review data would open every image file just to get dimensions)
            # Return empty review list and single exemplar
            return [build_identical_group(photo_list, exemplar_id)], [], [exemplar]

        # Otherwise, perform byte-by-byte comparison (original behavior)
        groups: list[IdenticalGroup] = []
        exemplars: list[PhotoFile] = []

        # Build dict and set for pick_exemplar_from_class
        photos_dict = {pf.id: pf for pf in photo_list}
        remaining_ids: set[int] = set(photos_dict.keys())

        # Process bin until empty
        while remaining_ids:
            # Pick exemplar from remaining files
            exemplar_id = pick_exemplar_from_class(photos_dict, remaining_ids)
            exemplar = photos_dict[exemplar_id]

            # Create new equivalence class starting with exemplar
            eq_class: list[PhotoFile] = [exemplar]
            remaining_ids.remove(exemplar_id)

            # Read exemplar once into memory for comparison against all candidates
            assert exemplar.path is not None, "Exemplar path cannot be None in production code"
            with exemplar.path.open("rb") as f:
                exemplar_data: bytes = f.read()

            # Compare exemplar against all remaining files
            to_remove: set[int] = set()
            photo_id: int
            for photo_id in remaining_ids:
                candidate: PhotoFile = photos_dict[photo_id]
                assert candidate.path is not None, "Candidate path cannot be None in production code"

                # Read candidate and compare (exemplar already in memory)
                with candidate.path.open("rb") as f:
                    candidate_data: bytes = f.read()

                if exemplar_data == candidate_data:
                    # Mark as identical to exemplar
                    candidate.cache["IDENTICAL"] = exemplar
                    eq_class.append(candidate)
                    to_remove.add(photo_id)

            # Remove matched files from remaining
            remaining_ids -= to_remove

            if len(eq_class) > 1:
                groups.append(build_identical_group(eq_class, exemplar_id))
            exemplars.append(exemplar)

        return groups, [], exemplars

    def accumulate_results(
        self,
        accum: dict[str, list[tuple[INDEX_T, PhotoFile]]],
        job: list[PhotoFile],
    ) -> None:
        """Accumulate exemplars from workers into template bins.

        Args:
            accum: Template bins dictionary being built incrementally
            job: List of exemplars from one worker (flat list)
        """
        # Bin each exemplar by its template
        for exemplar in job:
            assert "TEMPLATE" in exemplar.cache, "Template must be cached by SHA stage"
            accum.setdefault(exemplar.template, []).append((exemplar.template_index, exemplar))

    def needs_review(self) -> ReviewType:
        """This stage produces photo groups (byte-identical duplicates).

        Returns:
            "photos" to indicate this stage produces reviewable photo groups
        """
        return "photos"

    def has_review_data(self) -> bool:
        """Check if there are any identical photo groups to review.

        Returns:
            True if there are identical photo groups with 2+ photos, False otherwise
        """
        # Check if there are any classes with duplicates (2+ photos)
        return len(self.identical_review_result) > 0
