from __future__ import annotations

from typing import Optional, Dict, Any, TYPE_CHECKING, Generic
from uuid import UUID, uuid4

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

from cuid2 import cuid_wrapper

from opentelemetry import trace

from ..redact.redaction_options import (
    RedactionOptionsDict,
    RedactionOptions,
    normalize_redaction_options,
)
from .llm_reasoning import LLM_REASONING_ARG_NAME
from .tool_call_decision import ToolCallDecision, DecisionOutcome

from frisk_sdk.framework_adapter.framework_adapter import (
    NormalizedToolCall,
    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,
    ATTRIBUTE_NAME_REMOTE_SESSION_ID,
    ATTRIBUTE_NAME_LLM_REASONING,
)

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,
        *,
        tracer: Tracer,
        redact: RedactionOptions | RedactionOptionsDict | None = None,
    ):
        logger_name = f"{frisk.logger_name}.session"
        self.logger = get_logger(logger_name)

        self._cuid_generator = cuid_wrapper()
        self.id = self.generate_session_id()
        self.root_span: Optional[Span] = None
        self._root_run_id: UUID | str | None = None
        self.frisk = frisk
        self.tracer = tracer
        self.adapter = frisk._adapter
        self.options: RedactionOptionsDict = normalize_redaction_options(redact)

    def decide_tool_call(
        self, *, tool_call: NormalizedToolCall, agent_state: dict[str, Any]
    ) -> ToolCallDecision:
        """Evaluate a tool call within tracing context and capture span attributes.
        :param *:
        """
        with ToolCallSpan(
            session_id=str(self.id),
            framework_adapter=self.adapter,
            tool_call=tool_call,
            agent_state=agent_state,
            parent_span=self.get_root_span(),
            tracer=self.tracer,
            options=self.options,
        ) as tool_span:
            wrapped_tool_args = tool_call.args
            if (
                isinstance(wrapped_tool_args, dict)
                and 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(
                        ATTRIBUTE_NAME_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
            )

            try:
                tool_call_decision = self.frisk.decide_tool_call(
                    tool_call=tool_call,
                    agent_state=agent_state,
                    trace_context_carrier=trace_context_carrier,
                )
                tool_span.save_result(tool_call_decision)
                return tool_call_decision

            except Exception as e:
                wrapped = ToolCallEvaluationError(e)
                self.logger.error(wrapped)
                tool_span.save_error(wrapped)
                return ToolCallDecision(
                    outcome=DecisionOutcome.ERROR,
                    reason=str(wrapped),
                    rules_matched_count=0,
                )

    def init_tracing(
        self, *, remote_session_id: UUID | str, inputs: Dict[str, Any]
    ) -> Span:
        """Start the root session span and attach user-query context when available.
        :param *:
        """
        self._root_run_id = remote_session_id
        span = self.tracer.start_span(SPAN_NAME_FRISK_SESSION)
        span.set_attribute(ATTRIBUTE_NAME_SESSION_ID, str(self.id))
        span.set_attribute(ATTRIBUTE_NAME_REMOTE_SESSION_ID, str(remote_session_id))

        prompt = self.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.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_session_id(self):
        return str(uuid4())

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