from typing import List, Optional, Tuple
from uuid import UUID

from pydantic import BaseModel

from highlighter.client.base_models.annotation import Annotation
from highlighter.client.base_models.observation import Observation
from highlighter.core import LabeledUUID
from highlighter.core.data_models.data_sample import DataSample


class Prediction(BaseModel):
    pixel_location: Tuple[float, float, float, float] | List[Tuple[float, float]] | None
    attribute: LabeledUUID
    value: LabeledUUID
    conf: float

    @property
    def points(self):
        if self.pixel_location is None:
            raise ValueError("Prediction has no pixel location")
        if isinstance(self.pixel_location, tuple):
            x0, y0, x1, y1 = self.pixel_location
            points = [(x0, y0), (x1, y0), (x1, y1), (x0, y1), (x0, y0)]
        else:
            points = list(self.pixel_location)
            if points and points[0] != points[-1]:
                points.append(points[0])
        return points

    def to_annotation(
        self,
        data_sample: DataSample,
        pipeline_element_name: Optional[str] = None,
    ) -> Annotation:
        """
        recorded_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))  # FIXME remove default
        stream_frame_index: int = 0
        media_frame_index: int = 0
        # TODO: Replace with data_source_id once we re-target annotations and submissions from data_files to data_sources
        data_file_id: Optional[UUID] = None

        """
        if self.pixel_location is None:
            raise ValueError("Cannot convert Prediction to Annotation if Prediction has no pixel location")

        ann = Annotation.from_points(
            self.points,
            self.conf,
            data_sample=data_sample,
        )
        self.add_to_annotation(ann)
        return ann

    def add_to_annotation(self, src_annotation: Annotation) -> Observation:
        obs = Observation.make_enum_observation(
            self.attribute,
            self.attribute.label,
            self.value.label,
            self.value,
            self.conf,
            src_annotation.occurred_at,
            pipeline_element_name=src_annotation.datum_source.pipeline_element_name,
            training_run_id=src_annotation.datum_source.training_run_id,
            host_id=src_annotation.datum_source.host_id,
            frame_id=src_annotation.datum_source.frame_id,
        )
        src_annotation.observations.add(obs)
        return obs

    def translate(self, xoff=0, yoff=0):
        if isinstance(self.pixel_location, tuple):
            x0, y0, x1, y1 = self.pixel_location
            self.pixel_location = (x0 + xoff, y0 + yoff, x1 + xoff, y1 + yoff)
        else:
            self.pixel_location = [(x + xoff, y + yoff) for x, y in self.pixel_location]
