from __future__ import annotations
from pydantic import BaseModel, field_validator, field_serializer
from PIL import Image
from io import BytesIO
from pathlib import Path

from src.model_experiential import ExperientialModelOutput
from src.model_interpersonal import (
    ContactModelOutput,
    AngleModelOutput,
    PointOfViewModelOutput,
    DistanceModelOutput,
    ModalityLightingModelOutput,
    ModalityColorModelOutput,
    ModalityDepthModelOutput
)
from src.model_textual import (
    InformationValueModelOutput,
    FramingModelOutput,
    SalienceModelOutput
)

class ModelOutputs(BaseModel):
    experiential: ExperientialModelOutput
    contact: ContactModelOutput
    angle: AngleModelOutput
    point_of_view: PointOfViewModelOutput
    distance: DistanceModelOutput
    modality_lighting: ModalityLightingModelOutput
    modality_color: ModalityColorModelOutput
    modality_depth: ModalityDepthModelOutput
    information_value: InformationValueModelOutput
    framing: FramingModelOutput
    salience: SalienceModelOutput


class VisualCommunication(BaseModel):
    name: str
    image: Image.Image | BytesIO | bytes
    annotation: ModelOutputs | None = None
    prediction: ModelOutputs | None = None

    class Config:
        arbitrary_types_allowed = True

    @classmethod
    def classname(cls) -> str:
        """Return classname."""
        return cls.__name__

    @classmethod
    def from_file(cls, path: Path) -> VisualCommunication:
        """Instantiate from file."""
        name = path.stem
        image = Image.open(path)
        image.load()
        return VisualCommunication(name=name, image=image)

    @field_serializer("image")
    def serialize_image(image: Image.Image) -> bytes:
        buffer = BytesIO()
        image.save(buffer, format="JPEG")
        return buffer.getvalue()

    @field_validator("image", mode="before")
    @classmethod
    def convert_to_image(cls, image: Image.Image | BytesIO | bytes) -> Image.Image:
        if isinstance(image, bytes):
            image = BytesIO(image)
        if isinstance(image, BytesIO):
            image = Image.open(image)
        return image

    def __repr__(self) -> str:
        return f"{self.classname()}(name='{self.name}')"