from typing import Any, Optional, Dict, Union
from agents import tracing

import logging

from opik import context_storage, tracing_runtime_config
from opik.api_objects.span import span_data
from opik.api_objects.trace import trace_data
from opik.api_objects import opik_client
from opik.decorator import span_creation_handler, arguments_helpers

from . import span_data_parsers

LOGGER = logging.Logger(__name__)

OPENAI_SPAN_TYPES_WITH_MEANINGFUL_INPUT_OUTPUT = ["response", "generation", "custom"]


class OpikTracingProcessor(tracing.TracingProcessor):
    def __init__(
        self,
        project_name: Optional[str] = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        # NOTE: The OpikTracingProcessor maps openai traces/spans to
        # opik traces/spans. Openai traces may be logged as opik spans.
        self._opik_spans_data_map: Dict[str, span_data.SpanData] = {}
        """Map from openai span id to the opik span data created for it."""
        self._created_opik_traces_data_map: Dict[str, trace_data.TraceData] = {}
        """Map from openai trace id to the opik trace data created for it."""

        self._project_name = project_name

        self._opik_client = opik_client.get_client_cached()

        self._openai_trace_id_to_first_meaningful_input: Dict[str, Dict[str, Any]] = {}
        self._openai_trace_id_to_last_meaningful_output: Dict[str, Dict[str, Any]] = {}
        """
        Used to populate opik span/trace corresponding to the openai trace with meaningful inputs and outputs.
        By default inputs and outputs in the openai trace are always empty and it is harder to review in Opik UI.
        """

        self._openai_trace_id_to_external_opik_parent_span_id: Dict[
            str, Optional[str]
        ] = {}

        self._opik_context_storage = context_storage.get_current_context_instance()

    def force_flush(self) -> bool:
        return self._opik_client.flush()

    def shutdown(self) -> None:
        self.force_flush()

    def on_trace_start(self, trace: tracing.Trace) -> None:
        try:
            trace_metadata = trace.metadata or {}
            trace_metadata["created_from"] = "openai-agents"
            if trace.trace_id:
                trace_metadata["agents-trace-id"] = trace.trace_id

            current_trace = self._opik_context_storage.get_trace_data()
            if current_trace is None:
                current_trace = trace_data.TraceData(
                    name=trace.name,
                    project_name=self._project_name,
                    metadata=trace_metadata,
                    thread_id=trace.group_id,
                )
                self._opik_context_storage.set_trace_data(current_trace)
                self._created_opik_traces_data_map[trace.trace_id] = current_trace
                if (
                    self._opik_client.config.log_start_trace_span
                    and tracing_runtime_config.is_tracing_active()
                ):
                    self._opik_client.trace(**current_trace.as_start_parameters)
            else:
                start_span_arguments = arguments_helpers.StartSpanParameters(
                    name=trace.name,
                    project_name=self._project_name,
                    metadata=trace_metadata,
                    type="general",
                )
                result = span_creation_handler.create_span_respecting_context(
                    start_span_arguments=start_span_arguments,
                    distributed_trace_headers=None,
                    opik_context_storage=self._opik_context_storage,
                )
                self._opik_spans_data_map[trace.trace_id] = result.span_data
                self._opik_context_storage.add_span_data(result.span_data)
                self._openai_trace_id_to_external_opik_parent_span_id[
                    trace.trace_id
                ] = result.span_data.parent_span_id

                if (
                    self._opik_client.config.log_start_trace_span
                    and tracing_runtime_config.is_tracing_active()
                ):
                    self._opik_client.span(**result.span_data.as_start_parameters)

        except Exception:
            LOGGER.error("on_trace_start failed", exc_info=True)

    def on_trace_end(self, trace: tracing.Trace) -> None:
        try:
            opik_trace_or_span_data = self._try_get_span_or_trace(trace.trace_id)
            if opik_trace_or_span_data is None:
                LOGGER.error(
                    f"on_trace_end failed: no opik span/trace found for openai trace {trace.trace_id}. Probably due to an error in the previous callback execution"
                )
                return

            opik_trace_or_span_data.init_end_time()

            self._populate_root_span_or_trace_with_meaningful_input_and_output(
                trace.trace_id, opik_trace_or_span_data
            )
            if isinstance(opik_trace_or_span_data, trace_data.TraceData):
                if tracing_runtime_config.is_tracing_active():
                    self._opik_client.trace(**opik_trace_or_span_data.as_parameters)
            else:
                if tracing_runtime_config.is_tracing_active():
                    self._opik_client.span(**opik_trace_or_span_data.as_parameters)
        except Exception:
            LOGGER.error("on_trace_end failed", exc_info=True)
        finally:
            self._try_finalize_openai_trace(trace.trace_id)

    def on_span_start(self, span: tracing.Span[Any]) -> None:
        try:
            parsed_span_data = span_data_parsers.parse_spandata(span.span_data)

            opik_span_or_trace_data = (
                self._try_get_span_or_trace(span.parent_id)
                if span.parent_id is not None
                else self._try_get_span_or_trace(span.trace_id)
            )
            assert opik_span_or_trace_data is not None
            opik_span_data = opik_span_or_trace_data.create_child_span_data(
                name=parsed_span_data.name,
                type=parsed_span_data.type,
                input=parsed_span_data.input,
                output=parsed_span_data.output,
                metadata=parsed_span_data.metadata,
                usage=parsed_span_data.usage,
                model=parsed_span_data.model,
                provider=parsed_span_data.provider,
            )

            self._opik_context_storage.add_span_data(opik_span_data)
            self._opik_spans_data_map[span.span_id] = opik_span_data

            if (
                self._opik_client.config.log_start_trace_span
                and tracing_runtime_config.is_tracing_active()
            ):
                self._opik_client.span(**opik_span_data.as_start_parameters)

        except Exception:
            LOGGER.error("on_span_start failed", exc_info=True)

    def on_span_end(self, span: tracing.Span[Any]) -> None:
        try:
            parsed_span_data = span_data_parsers.parse_spandata(span.span_data)

            opik_span_data = self._opik_spans_data_map[span.span_id]
            opik_span_data.init_end_time().update(**parsed_span_data.__dict__)

            if tracing_runtime_config.is_tracing_active():
                self._opik_client.span(**opik_span_data.as_parameters)

            if (
                span.span_data.type
                not in OPENAI_SPAN_TYPES_WITH_MEANINGFUL_INPUT_OUTPUT
            ):
                return

            if (
                opik_span_data.input is not None
                and span.trace_id not in self._openai_trace_id_to_first_meaningful_input
            ):
                self._openai_trace_id_to_first_meaningful_input[span.trace_id] = (
                    opik_span_data.input
                )

            if opik_span_data.output is not None:
                self._openai_trace_id_to_last_meaningful_output[span.trace_id] = (
                    opik_span_data.output
                )

        except Exception:
            LOGGER.error("on_span_end failed", exc_info=True)
        finally:
            self._try_finalize_openai_span(span.span_id)

    def _try_get_span_or_trace(
        self, id: str
    ) -> Union[span_data.SpanData, trace_data.TraceData, None]:
        return self._opik_spans_data_map.get(
            id
        ) or self._created_opik_traces_data_map.get(id)

    def _populate_root_span_or_trace_with_meaningful_input_and_output(
        self,
        openai_trace_id: str,
        opik_trace_or_span_data: Union[trace_data.TraceData, span_data.SpanData],
    ) -> None:
        first_meaningful_input = self._openai_trace_id_to_first_meaningful_input.get(
            openai_trace_id
        )
        last_meaningful_output = self._openai_trace_id_to_last_meaningful_output.get(
            openai_trace_id
        )

        if first_meaningful_input is not None:
            opik_trace_or_span_data.update(input=first_meaningful_input)

        if last_meaningful_output is not None:
            opik_trace_or_span_data.update(output=last_meaningful_output)

    def _try_finalize_openai_trace(self, openai_trace_id: str) -> None:
        if openai_trace_id in self._opik_spans_data_map:
            opik_span_data = self._opik_spans_data_map[openai_trace_id]
            self._opik_context_storage.pop_span_data(ensure_id=opik_span_data.id)
            self._opik_spans_data_map.pop(openai_trace_id)
        elif openai_trace_id in self._created_opik_traces_data_map:
            opik_trace_data = self._created_opik_traces_data_map[openai_trace_id]
            self._opik_context_storage.pop_trace_data(ensure_id=opik_trace_data.id)
            self._created_opik_traces_data_map.pop(openai_trace_id)

        root_external_parent_span_id = (
            self._openai_trace_id_to_external_opik_parent_span_id.pop(
                openai_trace_id, None
            )
        )
        if root_external_parent_span_id is not None:
            # ensure no hanging spans left from the OpikTracingProcessor
            self._opik_context_storage.trim_span_data_stack_to_certain_span(
                root_external_parent_span_id
            )

        self._openai_trace_id_to_first_meaningful_input.pop(openai_trace_id, None)
        self._openai_trace_id_to_last_meaningful_output.pop(openai_trace_id, None)

    def _try_finalize_openai_span(self, openai_span_id: str) -> None:
        if openai_span_id in self._opik_spans_data_map:
            opik_span_data = self._opik_spans_data_map[openai_span_id]
            self._opik_context_storage.pop_span_data(ensure_id=opik_span_data.id)
            self._opik_spans_data_map.pop(openai_span_id)
