from typing import (
    Any,
    Iterable,
    List,
    NamedTuple,
    Optional,
    Tuple,
    TypeAlias,
    Union,
)

import numpy as np
from PIL import Image
from pydantic import BaseModel, Field, PrivateAttr, ValidationInfo, field_validator
from shapely import geometry

from highlighter.client.base_models.annotation import Annotation
from highlighter.client.base_models.datum_source import DatumSource
from highlighter.client.base_models.entity import Entity
from highlighter.core.data_models.data_sample import DataSample
from highlighter.core.labeled_uuid import LabeledUUID
from highlighter.datasets.cropping import CropArgs
from highlighter.predictors.prediction import Prediction

Zone: TypeAlias = Union[Tuple[int, int, int, int] | Tuple[float, float, float, float]]


class CropProcessor(BaseModel):
    crop_args: CropArgs = Field(..., default_factory=lambda: CropArgs())
    zone_entity_ids: Optional[List[LabeledUUID]] = None
    zones: Optional[List[Zone]] = None

    class ZoneSpec(NamedTuple):
        entity_id: LabeledUUID
        location: geometry.Polygon

    _zones: Optional[List["CropProcessor.ZoneSpec"]] = PrivateAttr(default=None)

    @field_validator("zones", mode="after")
    @classmethod
    def set_zone_annotations(cls, zones: Optional[List[Zone]], info: ValidationInfo) -> Optional[List[Zone]]:
        """Instantiate zone specs for each zone"""
        if zones is None:
            zone_specs = None
        else:
            zone_entity_ids = info.data.get("zone_entity_ids")
            if zone_entity_ids:
                assert len(zone_entity_ids) == len(zones), "Must provide an entity_id for each zone, or None"
            else:
                zone_entity_ids = [LabeledUUID.uuid4(label=f"zone_{i}") for i, _ in enumerate(zones)]

            def make_zone_spec(zone, zone_entity_id):
                x0, y0, x1, y1 = zone
                location = geometry.Polygon([(x0, y0), (x1, y0), (x1, y1), (x0, y1), (x0, y0)])
                return cls.ZoneSpec(entity_id=zone_entity_id, location=location)

            zone_specs = [make_zone_spec(z, zid) for z, zid in zip(zones, zone_entity_ids)]

        info.data["_zones"] = zone_specs
        return zones

    def crop_annotations(
        self,
        annotations: Iterable[Annotation] = None,
    ) -> List[Union[np.ndarray | Annotation]]:
        """Crop each the Annotation.location for each Annotation.data_samples"""
        return [a.crop_image(crop_args=self.crop_args) for a in annotations]

    def crop_data_samples(
        self,
        data_samples: List[DataSample],
    ) -> Tuple[List[np.ndarray | Image.Image], List[ZoneSpec]]:
        """Crop each DataSample in data_samples with the locations in _zones

        Return a tuple of 2 flat lists:

            result[0]: a list of the cropped images
            result[1]: a list of Annotation objects corrisponding to each crop
        """

        if not self._zones:
            raise ValueError("No zones provided")

        # Create a copy of each zone spec for each DataSample. This is the
        # same as if a model produced N annotations per image.
        zone_specs = [[item for item in self._zones] for _ in data_samples]

        crop_imgs = []
        crop_src_zones = []
        for data_sample, specs in zip(data_samples, zone_specs):
            crop_locations = [s.location for s in specs]
            crops = data_sample.crop_content(crop_locations, crop_args=self.crop_args)
            crop_imgs.extend(crops)
            crop_src_zones.extend(specs)

        return crop_imgs, crop_src_zones
