import json
import os
from pathlib import Path
from tempfile import mkdtemp
from typing import Callable, Dict, Optional, Tuple, Union
from uuid import UUID

import aiko_services as aiko
import yaml
from aiko_services.main import aiko as aiko_main

__all__ = [
    "HLAgent",
    "set_mock_aiko_messager",
]


def load_pipeline_definition(path) -> Dict:
    path = Path(path)
    suffix = path.suffix

    if suffix in (".json",):
        with path.open("r") as f:
            pipeline_def = json.load(f)
    elif suffix in (".yml", ".yaml"):
        with path.open("r") as f:
            pipeline_def = yaml.safe_load(f)
    else:
        raise NotImplementedError(
            f"Unsupported pipeline_definition file, '{path}'." " Expected .json|.yml|.yaml"
        )

    def remove_dict_keys_starting_with_a_hash(data):
        if isinstance(data, dict):
            # Create a new dictionary excluding keys starting with "#"
            return {
                key: remove_dict_keys_starting_with_a_hash(value)
                for key, value in data.items()
                if not key.startswith("#")
            }
        elif isinstance(data, list):
            # If the item is a list, recursively apply the function to each element
            return [remove_dict_keys_starting_with_a_hash(item) for item in data]
        else:
            # If the item is neither a dict nor a list, return it as-is
            return data

    pipeline_def = remove_dict_keys_starting_with_a_hash(pipeline_def)
    return pipeline_def


class HLPipelineImpl(aiko.PipelineImpl):

    def __init__(self, context):
        self.hl_task_status_info = {
            "state": aiko.StreamState.RUN,
            "message": "",
        }
        self.on_agent_error: Optional[Callable] = None
        self.on_agent_stop: Optional[Callable] = None
        self.on_before_process_frame: Optional[Callable] = None
        self.on_after_process_frame: Optional[Callable] = None
        super().__init__(context)

    def _process_stream_event(self, element_name, stream, stream_event, diagnostic, in_destroy_stream=False):
        # class StreamState:
        #    ERROR =   -2  # Don't generate new frames and ignore queued frames
        #    STOP  =   -1  # Don't generate new frames and process queued frames
        #    RUN   =    0  # Generate new frames and process queued frames

        most_recent_state = self.hl_task_status_info["state"]

        if stream_event == aiko.StreamEvent.ERROR:
            if aiko.StreamState.ERROR < most_recent_state:
                self.hl_task_status_info["state"] = aiko.StreamState.ERROR
                self.hl_task_status_info["message"] = diagnostic

        elif stream_event == aiko.StreamEvent.STOP:
            if aiko.StreamState.STOP < most_recent_state:
                self.hl_task_status_info["state"] = aiko.StreamState.STOP
                self.hl_task_status_info["message"] = diagnostic

        return super()._process_stream_event(
            element_name, stream, stream_event, diagnostic, in_destroy_stream=in_destroy_stream
        )

    def run_pipeline_with_callabcks(self, *args, **kwargs):
        self.run(*args, **kwargs)

        hl_task_status = self.hl_task_status_info

        if (hl_task_status["state"] == aiko.StreamState.ERROR) and self.on_agent_error is not None:
            self.on_agent_error(self)
        elif (hl_task_status["state"] == aiko.StreamState.STOP) and self.on_agent_stop is not None:
            self.on_agent_stop(self)

    def process_frame(self, stream_dict, frame_data) -> Tuple[aiko.StreamEvent, dict]:
        try:
            if self.on_before_process_frame is not None:
                self.on_before_process_frame(self)
        except Exception as e:
            return aiko.StreamEvent.ERROR, {"diagnostic": e}

        result = self._process_frame_common(stream_dict, frame_data, True)

        try:
            if self.on_after_process_frame is not None:
                self.on_after_process_frame(self)
        except Exception as e:
            return aiko.StreamEvent.ERROR, {"diagnostic": e}

        return result


aiko.PipelineImpl = HLPipelineImpl


def _validate_uuid(s):
    try:
        u = UUID(s)
        return u
    except Exception as e:
        return None


def _validate_path(s) -> Optional[Path]:
    try:
        p = Path(s)
        if p.exists():
            return p
    except Exception as e:
        return None


class HLAgent:

    def __init__(
        self,
        pipeline_definition: Union[str, dict, os.PathLike],
        name: Optional[str] = None,
        dump_definition: Optional[os.PathLike] = None,
    ):
        if dump_definition is not None:
            pipeline_path = Path(dump_definition)
        else:
            pipeline_path = Path(mkdtemp()) / "pipeline_def.json"

        if pipeline_definition_path := _validate_path(pipeline_definition):
            definition_dict = load_pipeline_definition(pipeline_definition_path)
            if name is None:
                name = pipeline_definition_path.name

        elif def_uuid := _validate_uuid(pipeline_definition):
            raise NotImplementedError("Need to implement pulling the pipeline_definition from Highlighter")
        elif isinstance(pipeline_definition, dict):
            if name is None:
                raise ValueError(
                    "If pipeline_definition is a dict you must provide the 'name' arg to HLAgent.__init__"
                )
            definition_dict = pipeline_definition

        else:
            if Path(pipeline_definition).suffix not in (".json", ".yml", ".yaml"):
                raise SystemExit(f"pipeline_definition '{pipeline_definition}' path does not exist")
            else:
                raise SystemExit(f"pipeline_definition '{pipeline_definition}' id does not exist")

        self._dump_definition(definition_dict, pipeline_path)

        parsed_definition = aiko.PipelineImpl.parse_pipeline_definition(pipeline_path)

        init_args = aiko.pipeline_args(
            name,
            protocol=aiko.PROTOCOL_PIPELINE,
            definition=parsed_definition,
            definition_pathname=pipeline_path,
        )
        pipeline = aiko.compose_instance(aiko.PipelineImpl, init_args)

        self.pipeline = pipeline
        self.pipeline_definition = parsed_definition

    def _dump_definition(self, pipeline_def: Dict, path: Path):
        with path.open("w") as f:
            json.dump(pipeline_def, f, sort_keys=True, indent=2)

    def run(self, stream_id, stream_parameters, mqtt_connection_required=False, queue_response=None):
        self.pipeline.create_stream(stream_id, parameters=stream_parameters, queue_response=queue_response)
        self.pipeline.run_pipeline_with_callabcks(mqtt_connection_required=mqtt_connection_required)

    def get_head_data_source_capability(self):
        head_capability = [x for x in self.pipeline.pipeline_graph][0]
        if hasattr(head_capability.element, "_is_data_source"):
            return head_capability
        return None

    def set_callback(self, callback_name: str, callback: Callable):
        setattr(self.pipeline, callback_name, callback)

    def process_frame(self, frame_data, stream_id=0, frame_id=0):
        stream = {
            "stream_id": stream_id,
            "frame_id": frame_id,
        }
        return self.pipeline.process_frame(stream, frame_data)


def set_mock_aiko_messager():
    # ToDo: Chat with Andy about if this is a requirement. The issue is
    # in pipeline.py +999 causes an error because if I use `process_frame`
    # directly, without setting the aiko.message object to something I
    # get an attribute error when .publish is called
    class MockMessage:
        def publish(self, *args, **kwargs):
            pass

        def subscribe(self, *args, **kwargs):
            pass

        def unsubscribe(self, *args, **kwargs):
            pass

    aiko_main.message = MockMessage()
