from __future__ import annotations

import logging
from base64 import b64decode
from base64 import b64encode
from io import BytesIO
from pathlib import Path

from PIL import Image
from pydantic import BaseModel
from pydantic import field_serializer
from pydantic import field_validator

from core.dto import ModelData


class NoDocumentFoundException(Exception):
    pass


class VisualCommunication(BaseModel):
    name: str
    image: Image.Image
    annotation: ModelData | None = None
    prediction: ModelData | 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)

    @classmethod
    def decode_image(cls, content: str) -> Image.Image:
        """Decode image."""
        _, content_data = content.split(',')
        return Image.open(BytesIO(b64decode(content_data)))

    @field_serializer('image')
    def serialize_image(image: Image.Image) -> bytes:   # type: ignore
        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}')"

    def webencoded_image(self) -> str:
        """Convert image to be displayed on webpage."""
        # convert images to bytes string
        buffer = BytesIO()
        self.image.save(buffer, format='png')
        img_enc = b64encode(buffer.getvalue()).decode('utf-8')
        return f"data:image/png;base64, {img_enc}"

    def generate_random_prediction(self, force: bool = False) -> None:
        """Generate random prediction values."""
        if not force and self.prediction is not None:
            logging.warning('set force=True to overwrite existing values.')
        self.prediction = ModelData.from_random()