"""
LockStock OpenAI Agents SDK Tracing
------------------------------------
Tracing integration for audit trail logging.

Uses the thin client pattern - all verification happens server-side via /v1/stamp.

The OpenAI Agents SDK supports extensible tracing. This module
integrates LockStock's hash chain audit trail with the SDK's
tracing infrastructure.
"""

from typing import Dict, Any, Optional
from dataclasses import dataclass, field
import time

from lockstock_core import LockStockClient


@dataclass
class TraceSpan:
    """A span in the trace."""
    name: str
    start_time: float
    end_time: Optional[float] = None
    metadata: Dict[str, Any] = field(default_factory=dict)
    parent_span_id: Optional[str] = None
    span_id: Optional[str] = None
    status: str = "running"


class LockStockTracer:
    """
    Tracer that logs agent activity to LockStock's audit trail.

    Integrates with OpenAI Agents SDK's tracing infrastructure
    while maintaining cryptographic audit guarantees via /v1/stamp.
    """

    def __init__(
        self,
        agent_id: str,
        api_key: str,
        endpoint: str = "https://lockstock-api-i9kp.onrender.com"
    ):
        """
        Initialize LockStock tracer.

        Args:
            agent_id: The agent's unique identifier
            api_key: API key for authentication (from dashboard)
            endpoint: LockStock API endpoint
        """
        self.client = LockStockClient(
            agent_id=agent_id,
            api_key=api_key,
            endpoint=endpoint
        )
        self.agent_id = agent_id
        self._spans: Dict[str, TraceSpan] = {}
        self._span_counter = 0

    async def start_span(
        self,
        name: str,
        parent_span_id: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None
    ) -> str:
        """
        Start a new trace span.

        Args:
            name: Name of the span
            parent_span_id: Optional parent span ID
            metadata: Optional metadata

        Returns:
            The span ID
        """
        self._span_counter += 1
        span_id = f"{self.agent_id}:span:{self._span_counter}"

        span = TraceSpan(
            name=name,
            start_time=time.time(),
            metadata=metadata or {},
            parent_span_id=parent_span_id,
            span_id=span_id
        )

        self._spans[span_id] = span

        # Stamp span start (server handles audit logging)
        await self.client.stamp(task=f"span_start")

        return span_id

    async def end_span(
        self,
        span_id: str,
        status: str = "success",
        metadata: Optional[Dict[str, Any]] = None
    ):
        """
        End a trace span.

        Args:
            span_id: The span ID to end
            status: Final status of the span
            metadata: Optional additional metadata
        """
        if span_id not in self._spans:
            return

        span = self._spans[span_id]
        span.end_time = time.time()
        span.status = status

        if metadata:
            span.metadata.update(metadata)

        # Stamp span end (server handles audit logging)
        await self.client.stamp(task=f"span_end")

    async def log_event(
        self,
        name: str,
        span_id: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None
    ):
        """
        Log an event within a span.

        Args:
            name: Event name
            span_id: Optional span ID
            metadata: Optional metadata
        """
        # Stamp the event (server handles audit logging)
        await self.client.stamp(task=f"event")

    async def log_tool_call(
        self,
        tool_name: str,
        tool_input: Dict[str, Any],
        tool_output: Optional[Any] = None,
        span_id: Optional[str] = None,
        authorized: bool = True
    ):
        """
        Log a tool call with its input and output.

        Args:
            tool_name: Name of the tool
            tool_input: Tool input parameters
            tool_output: Tool output (optional)
            span_id: Optional span ID
            authorized: Whether the call was authorized
        """
        # Stamp the tool call via stamp_tool
        await self.client.stamp_tool(tool_name=tool_name)

    def get_trace_summary(self) -> Dict[str, Any]:
        """
        Get a summary of all traced spans.

        Returns:
            Dictionary with trace summary
        """
        completed = [s for s in self._spans.values() if s.end_time is not None]
        running = [s for s in self._spans.values() if s.end_time is None]

        return {
            "agent_id": self.agent_id,
            "total_spans": len(self._spans),
            "completed_spans": len(completed),
            "running_spans": len(running),
            "spans": [
                {
                    "span_id": s.span_id,
                    "name": s.name,
                    "status": s.status,
                    "duration_ms": (s.end_time - s.start_time) * 1000 if s.end_time else None
                }
                for s in self._spans.values()
            ]
        }

    async def close(self):
        """Close the underlying client."""
        await self.client.close()
