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

from opentelemetry.trace import Tracer, Span
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator

from cuid2 import cuid_wrapper

from opentelemetry import trace

from .frisk_options import FriskOptionsDict, 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
from ..telemetry.constants import (
    SPAN_NAME_FRISK_SESSION,
    ATTRIBUTE_NAME_SESSION_PROMPT,
    ATTRIBUTE_NAME_SESSION_ID,
)

if TYPE_CHECKING:
    from .frisk import Frisk


class FriskSession(Generic[ToolCallT]):
    """Per-session policy and tracing context used during agent execution."""

    def __init__(
        self,
        frisk: Frisk,
        *,
        session_id: UUID,
        tracer: Tracer,
        framework_adapter: FrameworkAdapter[ToolCallT, Any],
        options: object | None = 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._root_run_id: UUID | str | None = 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: UUID | str, inputs: Dict[str, Any]) -> Span:
        """Start the root session span and attach user-query context when available."""
        self._root_run_id = run_id
        span = self.tracer.start_span(SPAN_NAME_FRISK_SESSION)
        span.set_attribute(ATTRIBUTE_NAME_SESSION_ID, str(self.session_id))

        prompt = self.framework_adapter.extract_prompt(agent_state=inputs)
        if prompt:
            span.set_attribute(ATTRIBUTE_NAME_SESSION_PROMPT, prompt)
        self.set_root_span(span)
        return span

    def end_tracing(self):
        """End the root span and unregister the session from the global registry."""
        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:
        """Set the root span used to parent per-tool evaluation spans."""
        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) -> Optional[Span]:
        """Return the current root span for this session."""
        return self.root_span

    def generate_tool_call_id(self) -> str:
        """Generate a unique tool-call identifier for this session."""
        return self._cuid_generator()

    def evaluate_tool_call(
        self, tool_call_info: ToolCallInfo, agent_state: dict[str, Any]
    ) -> ProcessToolCallResult:
        """Evaluate a tool call within tracing context and capture span attributes."""
        with ToolCallSpan(
            str(self.session_id),
            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 and tool_span.span is not None:
                    tool_span.span.set_attribute("llm_reasoning", llm_reasoning)

            if tool_span.span is None:
                raise RuntimeError("Tool call span was not initialized")
            span_context = trace.set_span_in_context(tool_span.span)
            trace_context_carrier: dict[str, str] = {}
            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,
                )
