import logging
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from uuid import UUID

from highlighter.agent.capabilities._crop_processor import CropProcessor
from highlighter.agent.capabilities.base_capability import Capability, StreamEvent
from highlighter.client.gql_client import HLClient
from highlighter.client.training_runs import (
    TrainingRunArtefactType,
)
from highlighter.core.data_models import DataSample
from highlighter.predictors.onnx_yolov8 import OnnxYoloV8 as Predictor
from highlighter.predictors.output_category_descriptor import OutputCategoryDescriptor

__all__ = ["OnnxYoloV8"]

logger = logging.getLogger(__name__)


def flatten(lst):
    for x in lst:
        if isinstance(x, list):
            yield from flatten(x)
        else:
            yield x


class OnnxYoloV8(Capability):

    class InitParameters(Capability.InitParameters):
        onnx_file: Optional[str] = None
        training_run_artefact_id: Optional[UUID] = None
        num_classes: int = 80
        class_lookup: Optional[Dict[int | str, OutputCategoryDescriptor]] = None
        conf_thresh: float = 0.1
        nms_iou_thresh: float = 0.5
        nms_class_agnostic: bool = False
        is_absolute: bool = True
        crop_processor: Optional[CropProcessor] = None

    def __init__(self, context):
        super().__init__(context)
        self._predictor: Optional[Predictor] = None

        onnx_file = self.init_parameters.onnx_file
        if onnx_file is None:
            training_run_artefact_id = self.init_parameters.training_run_artefact_id
            if training_run_artefact_id is None:
                raise ValueError("Must provide `onnx_file` or `training_run_artefact_id`")
            training_run_artefact = HLClient.get_client().trainingRunArtefact(
                return_type=TrainingRunArtefactType, id=training_run_artefact_id
            )
            onnx_file = training_run_artefact.file_url

        artefact_cache_dir = Path.home() / ".cache" / "highlighter" / "agents" / "artefacts"
        # FIXME: kwargs {"device_id", "artefact_cache_dir", "onnx_file_download_timeout"}
        # should be optionally configured by the runtime context.

        self._predictor = Predictor(
            onnx_file,
            num_classes=self.init_parameters.num_classes,
            class_lookup=self.init_parameters.class_lookup,
            conf_thresh=self.init_parameters.conf_thresh,
            nms_iou_thresh=self.init_parameters.nms_iou_thresh,
            is_absolute=self.init_parameters.is_absolute,
            nms_class_agnostic=self.init_parameters.nms_class_agnostic,
            artefact_cache_dir=artefact_cache_dir,
        )

        self._crop_processor = None
        if self.init_parameters.crop_processor is not None:
            self._crop_processor = self.init_parameters.crop_processor

    def process_frame(self, stream, data_samples: List[DataSample]) -> Tuple[StreamEvent, Dict]:

        annotations = []
        if self._crop_processor is None:
            image_batch = [d.content for d in data_samples]
            batch_preds = self._predictor.predict(image_batch)

            assert len(batch_preds) == len(data_samples)
            for preds, data_sample in zip(batch_preds, data_samples):
                for pred in preds:
                    new_ann = pred.to_annotation(
                        data_sample,
                        pipeline_element_name=self.name,
                    )
                    annotations.append(new_ann)
        else:

            crops, crop_src_zones = self._crop_processor.crop_data_samples(data_samples)
            preds = self._predictor.predict(crops)

            # Translate each prediction according to the corresponding source zone
            for predictions_from_zones, zone_spec, data_sample in zip(preds, crop_src_zones, data_samples):
                for prediction in predictions_from_zones:
                    prediction.translate(*zone_spec.location.bounds[:2])

                    annotations.append(
                        prediction.to_annotation(
                            data_sample,
                            pipeline_element_name=self.name,
                        )
                    )

        entities = [a.get_or_create_entity() for a in annotations]

        return StreamEvent.OKAY, {"entities": {e.id: e for e in entities}}
