from __future__ import annotations
from typing import Optional, Dict, Any, TYPE_CHECKING, Generic
from uuid import UUID

from opentelemetry.propagators.textmap import CarrierT
from opentelemetry.sdk.trace import Span
from opentelemetry.trace import Tracer
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator

from cuid2 import cuid_wrapper

from opentelemetry import trace

from .frisk_options import FriskOptionsDict, FriskOptionsInput, normalize_frisk_options
from .llm_reasoning import LLM_REASONING_ARG_NAME, remove_llm_reasoning_arg
from .process_tool_call_result import ProcessToolCallResult, ProcessToolCallDecision

from frisk_sdk.framework_adapter.framework_adapter import (
    FrameworkAdapter,
    ToolCallInfo,
    ToolCallT,
)
from .tool_call_span import ToolCallSpan
from ..errors.invocation_errors import ToolCallEvaluationError
from ..logging import get_logger

if TYPE_CHECKING:
    from .frisk import Frisk


class FriskSession(Generic[ToolCallT]):
    def __init__(
        self,
        frisk: Frisk,
        *,
        session_id: UUID,
        tracer: Tracer,
        framework_adapter: FrameworkAdapter[ToolCallT],
        options: Optional[FriskOptionsInput] = None,
    ):
        logger_name = f"{frisk.logger_name}.session"
        self.logger = get_logger(logger_name)

        self.session_id = session_id
        self.root_span: Optional[Span] = None
        self.frisk = frisk
        self.tracer = tracer
        self._cuid_generator = cuid_wrapper()
        self.framework_adapter = framework_adapter
        self.options: FriskOptionsDict = normalize_frisk_options(options)

    def init_tracing(self, run_id: str, inputs: Dict[str, Any]) -> Span:
        self._root_run_id = run_id
        span = self.tracer.start_span("frisk_session")
        user_query = self.framework_adapter.get_user_query(agent_state=inputs)

        if user_query:
            span.set_attribute("user_query", user_query)
        self.set_root_span(span)
        return span

    def end_tracing(self):
        if self.root_span:
            self.logger.debug("Closing root span")
            self.root_span.end()

        # Unregister session from registry when tracing ends
        from .session_registry import SessionRegistry

        registry = SessionRegistry()
        registry.unregister(self.session_id)

    def set_root_span(self, span: Span) -> None:
        self.root_span = span  # todo: prevent override. https://linear.app/friskai/issue/POL-95/encapsulate-all-otel-span-creation-into-frisk-session

    def get_root_span(self) -> Span:
        return self.root_span

    def generate_tool_call_id(self) -> str:
        return self._cuid_generator()

    def evaluate_tool_call(
        self, tool_call_info: ToolCallInfo, agent_state: dict[str, Any]
    ) -> ProcessToolCallResult:
        with ToolCallSpan(
            self.framework_adapter,
            tool_call_info,
            agent_state,
            self.get_root_span(),
            self.tracer,
            options=self.options,
        ) as tool_span:
            wrapped_tool_args = tool_call_info.args
            if LLM_REASONING_ARG_NAME in wrapped_tool_args:
                llm_reasoning = wrapped_tool_args.get(LLM_REASONING_ARG_NAME)
                if llm_reasoning is not None:
                    tool_span.span.set_attribute("llm_reasoning", llm_reasoning)

            span_context = trace.set_span_in_context(tool_span.span)
            trace_context_carrier: CarrierT = {}
            TraceContextTextMapPropagator().inject(
                trace_context_carrier, context=span_context
            )

            base_tool_args = remove_llm_reasoning_arg(wrapped_tool_args)
            try:
                process_tool_call_result = self.frisk.evaluate_tool_call(
                    tool_call_info.name,
                    base_tool_args,
                    agent_state,
                    tool_call_info.id,
                    trace_context_carrier,
                )
                tool_span.save_result(process_tool_call_result)
                return process_tool_call_result

            except Exception as e:
                wrapped = ToolCallEvaluationError(e)
                self.logger.error(wrapped)
                tool_span.save_error(wrapped)
                return ProcessToolCallResult(
                    decision=ProcessToolCallDecision.ERROR,
                    reason=str(wrapped),
                    rules_matched_count=0,
                )
