"""Definition of VisualCommunication model."""
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 shared.dto import ModelData


class VisualCommunication(BaseModel):
    """Visual communication model."""

    name: str
    image: Image.Image
    annotation: ModelData | None = None
    prediction: ModelData | None = None

    class Config:
        """BaseModel configuration."""
        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:
        """Extract image from webencoded content."""
        _, content_data = content.split(',')
        return Image.open(BytesIO(b64decode(content_data)))

    @field_serializer('image')
    @classmethod
    def serialize_image(cls, image: Image.Image) -> bytes:   # type: ignore
        """Convert image to bytes for storage in database."""
        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:
        """Convert bytes input from database into 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()