from __future__ import annotations

import io
from functools import cached_property

from PIL import Image
from pydantic import BaseModel, ConfigDict, model_validator

from .bbox import BBox
from .coordinate import Coordinate
from .exceptions import (
    InvalidResizeModeError,
    MissingResizeMetadataError,
    SpaceMismatchError,
)
from .resize import (
    ResizeMetadata,
    compute_letterbox_metadata,
    compute_stretch_metadata,
)
from .space import Space
from .types import ResizeMode


class Screenshot(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)

    image_bytes: bytes
    space: Space
    resize_metadata: ResizeMetadata | None = None

    @model_validator(mode="after")
    def _validate_image_bytes_and_space(self) -> Screenshot:
        try:
            with Image.open(io.BytesIO(self.image_bytes)) as image:
                image_width, image_height = image.size
        except Exception as exc:  # pragma: no cover
            raise ValueError("image_bytes must contain a valid image") from exc
        if image_width != self.space.width or image_height != self.space.height:
            raise ValueError(
                f"space {self.space.width}x{self.space.height} does not match "
                f"image dimensions {image_width}x{image_height}"
            )
        return self

    @cached_property
    def image(self) -> Image.Image:
        return Image.open(io.BytesIO(self.image_bytes))

    @staticmethod
    def from_image(img: Image.Image) -> Screenshot:
        buf = io.BytesIO()
        img.save(buf, format="PNG")
        return Screenshot(
            image_bytes=buf.getvalue(),
            space=Space(width=img.width, height=img.height),
        )

    def _require_resize_metadata(self) -> ResizeMetadata:
        if self.resize_metadata is None:
            raise MissingResizeMetadataError(
                "This screenshot has no resize_metadata. Resize first or pass metadata explicitly."
            )
        if self.space != self.resize_metadata.target_space:
            raise SpaceMismatchError(
                f"Screenshot space {self.space} must match resize metadata target_space "
                f"{self.resize_metadata.target_space}"
            )
        return self.resize_metadata

    def map_coord_to_original(self, coord: Coordinate) -> Coordinate:
        """Map a coordinate in this resized screenshot back into original source space."""
        metadata = self._require_resize_metadata()
        if coord.space != self.space:
            raise SpaceMismatchError(
                f"Coordinate space {coord.space} must match screenshot space {self.space}"
            )
        return coord.to_source_space(metadata)

    def map_coord_from_original(self, coord: Coordinate) -> Coordinate:
        """Map a coordinate in original source space into this resized screenshot."""
        metadata = self._require_resize_metadata()
        return coord.to_target_space(metadata)

    def map_bbox_to_original(self, bbox: BBox, offset: BBox | None = None) -> BBox:
        """Map a bbox in this resized screenshot back into original source space."""
        metadata = self._require_resize_metadata()
        if bbox.space != self.space:
            raise SpaceMismatchError(
                f"BBox space {bbox.space} must match screenshot space {self.space}"
            )
        return bbox.to_source_space(metadata, offset=offset)

    def map_bbox_from_original(self, bbox: BBox, offset: BBox | None = None) -> BBox:
        """Map a bbox in original source space into this resized screenshot."""
        metadata = self._require_resize_metadata()
        return bbox.to_target_space(metadata, offset=offset)

    def resize(self, target: Space, mode: ResizeMode) -> Screenshot:
        if not isinstance(mode, ResizeMode):
            raise InvalidResizeModeError(f"mode must be a ResizeMode, got {mode!r}")
        img = self.image

        if mode == ResizeMode.LETTERBOX:
            metadata = compute_letterbox_metadata(self.space, target)
            scaled = img.resize(
                (metadata.scaled_width, metadata.scaled_height), Image.LANCZOS
            )
            canvas = Image.new("RGB", (target.width, target.height), (0, 0, 0))
            canvas.paste(scaled, (metadata.offset_x, metadata.offset_y))
            result_img = canvas
        else:
            metadata = compute_stretch_metadata(self.space, target)
            result_img = img.resize((target.width, target.height), Image.LANCZOS)

        buf = io.BytesIO()
        result_img.save(buf, format="PNG")
        return Screenshot(
            image_bytes=buf.getvalue(),
            space=target,
            resize_metadata=metadata,
        )
