import hashlib
import inspect
import json
import threading
import time
import traceback
from functools import lru_cache
from pathlib import Path
from queue import Empty, Queue
from typing import Any, Dict, Iterable, Optional, Tuple, Union
from uuid import UUID

import aiko_services as aiko
from aiko_services import (
    PROTOCOL_PIPELINE,
    ActorTopic,
)
from aiko_services import DataSource as AikoDataSource
from aiko_services import (
    PipelineImpl,
    StreamEvent,
    StreamState,
    compose_instance,
    pipeline_args,
    pipeline_element_args,
)
from pydantic import BaseModel, ConfigDict, Field

from highlighter.agent.capabilities.recorder import (
    Recorder,
    RecordingState,
    RecordMode,
)
from highlighter.client.json_tools import HLJSONEncoder
from highlighter.core.database.database import Database

# ToDo: Move CropArgs out of .dataset
from highlighter.datasets.cropping import CropArgs

__all__ = [
    "ActorTopic",
    "Capability",
    "DataSourceCapability",
    "ContextPipelineElement",
    "EntityUUID",
    "PROTOCOL_PIPELINE",
    "PipelineElement",
    "PipelineImpl",
    "StreamEvent",
    "StreamState",
    "compose_instance",
    "compose_instance",
    "pipeline_args",
    "pipeline_element_args",
]

VIDEO = "VIDEO"
TEXT = "TEXT"
IMAGE = "IMAGE"

EntityUUID = UUID

"""Decouple the rest of the code from aiko.PipelineElement"""
ContextPipelineElement = aiko.ContextPipelineElement
PipelineElement = aiko.PipelineElement

# SEPARATOR = b"\x1c"  # ASCII 28 (File Separator)
SEPARATOR = 28  # ASCII 28 (File Separator)


class _BaseCapability:
    class InitParameters(BaseModel):
        """Populate with init parameter fields"""

        model_config = ConfigDict(extra="ignore", arbitrary_types_allowed=True)

    class StreamParameters(InitParameters):
        """Populate with stream parameter fields"""

        pass

    @classmethod
    def definition(cls, outputs: Iterable[Tuple[str, str]], **parameters):
        """Generate a dictionary containing information about the Capability

        Args:

        """
        # Get class name
        name = cls.__name__

        # Process InitParameters
        if hasattr(cls, "InitParameters"):
            _ = cls.InitParameters(**parameters)

        # Get input and output parameters from process_frame method signature
        ins = []
        if hasattr(cls, "process_frame"):
            sig = inspect.signature(cls.process_frame)

            # Get input parameters
            for param_name, param in sig.parameters.items():
                if param_name not in ["self", "stream"]:
                    param_type = (
                        param.annotation.__name__ if param.annotation != inspect.Parameter.empty else "Any"
                    )
                    ins.append({"name": param_name, "type": param_type})
        else:
            raise ValueError(f"{cls.__name__} has no `process_frame` function.")

        outs = []
        for n, t in outputs:
            outs.append({"name": n, "type": t})

        # Get module information for deploy
        module_name = cls.__module__

        return {
            "name": name,
            "parameters": parameters,
            "input": ins,
            "output": outs,
            "deploy": {
                "local": {
                    "class_name": name,
                    "module": module_name,
                }
            },
        }

    @property
    @lru_cache(maxsize=1)
    def init_parameters(self):
        # Take init parameters from the agent definition at the
        # agent level and the capability level, with capability parameters
        # overriding agent-level parameters
        parameters = {
            **self.pipeline.definition.parameters,
            **self.definition.parameters,
        }
        return self.InitParameters(**parameters)

    def stream_parameters(self, stream_id):
        stream_parameters = self._get_stream_parameters()
        qualified_stream_parameters = {
            param_name.removeprefix(f"{self.definition.name}."): param_value
            for param_name, param_value in stream_parameters.items()
            if param_name.startswith(self.definition.name)
        }
        # Take stream parameters from the definition and from the stream,
        # with qualified stream parameters (prefixed with the capability name)
        # taking highest precendence.
        # Unqualified stream parameters will not override parameters
        # set on the capability definition, for compatibility with Aiko's parameter precedence.
        parameters = {
            **self.pipeline.definition.parameters,
            **stream_parameters,
            **self.definition.parameters,
            **qualified_stream_parameters,
        }

        # Cache construction of the StreamParameters pydantic model
        # which may be expensive.
        # TODO: Check if serializing to JSON is just as expensive
        # as building pydantic StreamParameters
        if stream_id not in self._cached_stream_parameters:
            self._cached_stream_parameters[stream_id] = (None, None)
        parameters_string = json.dumps(parameters, sort_keys=True, separators=(",", ":"), cls=HLJSONEncoder)
        parameters_hash = hashlib.sha256(parameters_string.encode()).hexdigest()
        cached_hash, cached_params = self._cached_stream_parameters[stream_id]
        if parameters_hash == cached_hash:
            return cached_params
        else:
            validated_params = self.StreamParameters(**parameters)
            self._cached_stream_parameters[stream_id] = (parameters_hash, validated_params)
            return validated_params

    def stop_stream(self, stream, stream_id):
        if stream_id in self._cached_stream_parameters:
            del self._cached_stream_parameters[stream_id]


class _AppendableIterator:
    def __init__(self, max_size: int = 10):
        """Initialize the iterator with an optional `max_size`."""
        self._maxsize = max_size
        self._items = Queue(maxsize=self._maxsize)

    def __iter__(self):
        """Make the class iterable."""
        return self

    def __next__(self):
        """Return the next item in the sequence or raise StopIteration."""
        try:
            item = self._items.get(timeout=10)
        except Empty:
            raise StopIteration
        return item

    def append(self, item):
        """Append an item to the underlying list."""
        self._items.put(item)

    def qsize(self) -> int:
        return self._items.qsize()

    def maxsize(self) -> int:
        return self._maxsize


class Capability(PipelineElement, _BaseCapability):

    class InitParameters(_BaseCapability.InitParameters):
        account_uuid: Optional[UUID] = None
        data_source_uuid: Optional[UUID] = None
        output_filename_template: Optional[str] = None
        output_folder: Optional[str] = None
        recording_state: RecordingState = RecordingState.ON
        record: RecordMode = RecordMode.OFF
        samples_per_file: Optional[int] = None
        seconds_per_file: Optional[float] = None
        database: Optional[Database] = None
        writer_opts: Dict[str, Any] = Field(default_factory=lambda: {})

    class StreamParameters(InitParameters):
        pass

    def recording_enabled(self, stream_id) -> bool:
        parameters = self.stream_parameters(stream_id)
        if not hasattr(parameters, "record") or parameters.record == RecordMode.OFF:
            return False

        required_args = {
            "account_uuid": parameters.account_uuid,
            "data_source_uuid": parameters.data_source_uuid,
        }

        if not all(required_args.values()):
            self.logger.warning(
                f"Recorder args are missing required parameters: {[name for name, value in required_args.items() if not value]}. recording disabled",
                extra={"stream_id": stream_id, "capability_name": self.name},
            )
            return False

        return True

    def __init__(self, context: aiko.ContextPipelineElement):
        context.get_implementation("PipelineElement").__init__(self, context)
        self._cached_stream_parameters = {}

    def process_frame(self, stream, *args) -> Tuple[StreamEvent, dict]:
        raise NotImplementedError()

    def start_stream(self, stream, stream_id, use_create_frame=True):
        return StreamEvent.OKAY, {}

    def stop_stream(self, stream, stream_id):
        _BaseCapability.stop_stream(self, stream, stream_id)
        # Recording cleanup is now handled at the agent level
        return StreamEvent.OKAY, None


# ToDO: Remove
class DataSourceType(BaseModel):
    # class MediaType(str, Enum):
    #    IMAGE = "IMAGE"
    #    TEXT = "TEXT"
    #    VIDEO = "VIDEO"

    media_type: str
    url: str
    id: UUID
    content: Optional[Any] = None

    @classmethod
    def image_iter(cls, images: Iterable[Union[str, Path, bytes]]):
        pass

    @classmethod
    def video_iter(cls, videos: Iterable[Union[str, Path, bytes]]):
        pass

    @classmethod
    def text_iter(cls, tests: Iterable[Union[str, Path, bytes]]):
        pass


class DataSourceCapability(AikoDataSource, _BaseCapability):

    stream_media_type = None

    class InitParameters(Capability.InitParameters):
        rate: Optional[float] = None
        batch_size: int = 1
        data_sources: Optional[str] = None
        file_ids: Optional[Iterable] = None
        task_id: Optional[UUID] = None

    class StreamParameters(InitParameters):
        pass

    def __init__(self, context: aiko.ContextPipelineElement):
        context.get_implementation("PipelineElement").__init__(self, context)
        self._dsps: Dict[str, Recorder] = {}
        self._cached_stream_parameters = {}

    def frame_generator(self, stream, pipeline_iter_idx):
        """Produce a batch of frames.

        Args:
            stream: The Stream context
            pipeline_iter_idx: An integer counting the number of times the
                               pipeline has been executed, (ie: process_frame
                               has been called)

        """
        if (
            "disable_create_frame_event" in stream.variables
            and stream.variables["disable_create_frame_event"].is_set()
        ):
            return StreamEvent.STOP, {"diagnostic": "Frame generation disabled"}

        parameters = self.stream_parameters(stream.stream_id)
        batch_size = parameters.batch_size
        task_id = parameters.task_id

        recording_state = parameters.recording_state

        frame_gen_start = time.perf_counter()
        frame_data_batch = {"data_samples": [], "entities": {}}
        for _ in range(batch_size):
            try:
                # self.set_dsp_recording_state(stream)
                sample_start = time.perf_counter()
                data_sample, entities = self.get_next_data_sample(stream)
                sample_elapsed = time.perf_counter() - sample_start
                if sample_elapsed > 0.1:  # Only log if > 100ms
                    self.logger.debug(
                        f"Stream {stream.stream_id} [{self.name}]: get_next_data_sample took {sample_elapsed:.3f}s"
                    )

                frame_data_batch["data_samples"].append(data_sample)
                frame_data_batch["entities"].update(entities)
                self.logger.debug(f"data_sample: {data_sample}, entities: {entities}")
            except StopIteration:
                pass
            except Exception as e:
                frame_data = {"diagnostic": traceback.format_exc()}
                return StreamEvent.ERROR, frame_data

        if not frame_data_batch["data_samples"]:
            return StreamEvent.STOP, {"diagnostic": "All frames generated"}

        # For each pipeline iteration the is a batch of file_ids and frame_ids
        stream.variables["task_id"] = task_id

        frame_gen_elapsed = time.perf_counter() - frame_gen_start
        if frame_gen_elapsed > 0.1:  # Only log if > 100ms
            self.logger.debug(
                f"Stream {stream.stream_id} [{self.name}]: frame_generator took {frame_gen_elapsed:.3f}s for {batch_size} samples"
            )

        return StreamEvent.OKAY, frame_data_batch

    def start_stream(self, stream, stream_id):
        stream.variables["video_capture"] = None
        stream.variables["video_frame_generator"] = None

        stream.variables["disable_create_frame_event"] = threading.Event()

        return super().start_stream(
            stream, stream_id, frame_generator=self.frame_generator, use_create_frame=False
        )

    def stop_stream(self, stream, stream_id):
        if "disable_create_frame_event" in stream.variables:
            stream.variables["disable_create_frame_event"].set()
        return super().stop_stream(stream, stream_id)

    def get_next_data_sample(self, stream):
        raise NotImplementedError()

    def process_frame(
        self, stream, data_samples, entities: Optional[Dict] = None
    ) -> Tuple[StreamEvent, Dict]:
        return StreamEvent.OKAY, {
            "data_samples": data_samples,
            "entities": entities if entities is not None else {},
        }

    def using_hl_data_scheme(self, stream) -> bool:
        return "hl_source_data" in stream.variables
