import json

import time
from typing import Dict, Any, Optional

from frisk_sdk.core.llm_reasoning import remove_llm_reasoning_arg
from frisk_sdk.core.tool_call_decision import ToolCallDecision
from frisk_sdk.telemetry.constants import (
    SPAN_NAME_DECIDE_TOOL_CALL,
    ATTRIBUTE_NAME_SESSION_ID,
    ATTRIBUTE_NAME_TOOL_NAME,
    ATTRIBUTE_NAME_TOOL_CALL_ID,
    ATTRIBUTE_NAME_TOOL_ARGS_JSON,
    ATTRIBUTE_NAME_TOOL_ARGS_REDACTED_PATHS_JSON,
    ATTRIBUTE_NAME_AGENT_STATE_JSON,
    ATTRIBUTE_NAME_AGENT_STATE_REDACTED_PATHS_JSON,
    ATTRIBUTE_NAME_DECISION_OUTCOME,
    ATTRIBUTE_NAME_DECISION_MATCHED_RULE_COUNT,
    ATTRIBUTE_NAME_DECISION_REASON,
    ATTRIBUTE_NAME_ERROR_MESSAGE,
    ATTRIBUTE_NAME_ERROR_TYPE,
    ATTRIBUTE_NAME_LATENCY_NS,
)
from opentelemetry import trace
from opentelemetry.trace import Tracer, use_span, Span

from frisk_sdk.redact.redaction_options import (
    RedactionOptionsDict,
    normalize_redaction_options,
)
from frisk_sdk.framework_adapter.framework_adapter import (
    FrameworkAdapter,
    NormalizedToolCall,
)


class ToolCallSpan:
    """
    Context manager wrapping an OpenTelemetry span for a single tool call.
    - __enter__: starts the span with initial attributes
    - close(...): sets ToolCallDecision-related attributes and ends the span
    - __exit__: ensures span ends even on exception
    """

    def __init__(
        self,
        *,
        session_id: str,
        framework_adapter: FrameworkAdapter[Any, Any],
        tool_call: NormalizedToolCall,
        agent_state: Dict[str, Any],
        parent_span: Optional[Span],
        tracer: Tracer,
        options: object | None = None,
    ):
        self.session_id = session_id
        self.tool_call = tool_call
        self.agent_state = agent_state
        self.parent_span = parent_span
        self.tracer = tracer
        self.framework_adapter = framework_adapter

        self.span: Optional[Span] = None
        self.start_time_ns: Optional[int] = None
        self.options: RedactionOptionsDict = normalize_redaction_options(options)

    def __enter__(self) -> "ToolCallSpan":
        # Prepare parent context from root span (if any)
        root_span = self.parent_span
        parent_span_context = (
            trace.set_span_in_context(root_span) if root_span else None
        )
        self.start_time_ns = time.perf_counter_ns()

        args_without_llm_reasoning = remove_llm_reasoning_arg(self.tool_call.args)
        redacted_tool_args_result = self.framework_adapter.serialize_tool_args(
            args_without_llm_reasoning,
            redact=self.options.get("redact_tool_args", False),
        )

        redacted_agent_state_result = self.framework_adapter.serialize_agent_state(
            self.agent_state, redact=self.options.get("redact_agent_state", False)
        )

        attributes: dict[str, Any] = {
            ATTRIBUTE_NAME_SESSION_ID: str(self.session_id),
            ATTRIBUTE_NAME_TOOL_NAME: self.tool_call.name,
            ATTRIBUTE_NAME_TOOL_CALL_ID: self.tool_call.id,
            ATTRIBUTE_NAME_TOOL_ARGS_JSON: redacted_tool_args_result.value,
            ATTRIBUTE_NAME_TOOL_ARGS_REDACTED_PATHS_JSON: json.dumps(
                redacted_tool_args_result.redacted_paths
            ),
            ATTRIBUTE_NAME_AGENT_STATE_JSON: redacted_agent_state_result.value,
            ATTRIBUTE_NAME_AGENT_STATE_REDACTED_PATHS_JSON: json.dumps(
                redacted_agent_state_result.redacted_paths
            ),
        }

        self.span = self.tracer.start_span(
            SPAN_NAME_DECIDE_TOOL_CALL,
            context=parent_span_context,
            attributes=attributes,
        )

        use_span(self.span)
        return self

    def save_result(self, tool_call_decision: ToolCallDecision) -> None:
        if not self.span:
            return  # todo: Gracefully handle close() being called twice. https://linear.app/friskai/issue/POL-98/gracefully-handle-spans-being-opened-or-closed-twice
        # Set final attributes from the policy engine result
        self.span.set_attributes(
            {
                ATTRIBUTE_NAME_DECISION_OUTCOME: tool_call_decision.outcome,
                ATTRIBUTE_NAME_DECISION_MATCHED_RULE_COUNT: tool_call_decision.rules_matched_count,
            }
        )

        if tool_call_decision.reason is not None:
            self.span.set_attribute(
                ATTRIBUTE_NAME_DECISION_REASON, tool_call_decision.reason
            )

    def save_error(self, error: Exception):
        if not self.span:
            return
        self.span.set_attribute(ATTRIBUTE_NAME_ERROR_MESSAGE, str(error))

    def __exit__(self, exc_type, exc_val, exc_tb):
        # If an exception occurs inside the context, end the span gracefully
        if self.span:
            # Record latency even on exception
            end_time = time.perf_counter_ns()
            latency_ns = (
                None if self.start_time_ns is None else end_time - self.start_time_ns
            )
            try:
                if latency_ns is not None:
                    self.span.set_attribute(ATTRIBUTE_NAME_LATENCY_NS, latency_ns)
                if exc_type:
                    # Minimal error tagging
                    self.span.set_attribute(ATTRIBUTE_NAME_ERROR_TYPE, str(exc_type))
                    if exc_val:
                        self.span.set_attribute(
                            ATTRIBUTE_NAME_ERROR_MESSAGE, str(exc_val)
                        )
                    else:
                        self.span.set_attribute(
                            ATTRIBUTE_NAME_ERROR_MESSAGE, "Unknown error"
                        )
            finally:
                self.span.end()
        # Propagate exception (do not suppress)
        return False
