"""
Claude Agent SDK callback handler for Aigie SDK.

Provides automatic tracing for Claude Agent SDK query execution,
tool usage, and conversation sessions.
"""

import logging
import uuid
from datetime import datetime
from typing import Any, Dict, List, Optional, TYPE_CHECKING

from ...buffer import EventType
from .cost_tracking import calculate_claude_cost

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    from .session_context import ClaudeSessionContext


class ClaudeAgentSDKHandler:
    """
    Claude Agent SDK handler for Aigie tracing.

    Automatically traces Claude Agent SDK executions including:
    - Query execution (stateless)
    - Client sessions (stateful)
    - Tool use with PreToolUse/PostToolUse hooks
    - Message streaming and costs
    - Token usage tracking

    Example:
        >>> from claude_agent_sdk import query
        >>> from aigie.integrations.claude_agent_sdk import ClaudeAgentSDKHandler
        >>>
        >>> handler = ClaudeAgentSDKHandler(
        ...     trace_name='research-agent',
        ...     metadata={'project': 'analysis'}
        ... )
        >>>
        >>> # Manual usage
        >>> await handler.handle_query_start("What is the capital of France?", {})
        >>> # ... execute query
        >>> await handler.handle_query_end(query_id, messages, result_message)
    """

    def __init__(
        self,
        trace_name: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
        tags: Optional[List[str]] = None,
        user_id: Optional[str] = None,
        session_id: Optional[str] = None,
        capture_tool_results: bool = True,
        capture_messages: bool = True,
        session_context: Optional["ClaudeSessionContext"] = None,
    ):
        """
        Initialize Claude Agent SDK handler.

        Args:
            trace_name: Name for the trace
            metadata: Additional metadata to attach
            tags: Tags to apply to trace and spans
            user_id: User ID for the trace
            session_id: Session ID for the trace
            capture_tool_results: Whether to capture tool results
            capture_messages: Whether to capture message content
            session_context: Optional session context for trace sharing
        """
        self.trace_name = trace_name
        self.metadata = metadata or {}
        self.tags = tags or []
        self.user_id = user_id
        self.session_id = session_id
        self.capture_tool_results = capture_tool_results
        self.capture_messages = capture_messages

        # Session context for trace sharing
        self._session_context = session_context

        # State tracking - inherit trace_id from session context if available
        self.trace_id: Optional[str] = session_context.trace_id if session_context else None
        self.query_span_id: Optional[str] = None
        self.session_span_id: Optional[str] = None
        self.tool_map: Dict[str, Dict[str, Any]] = {}  # tool_use_id -> {spanId, startTime}
        self.turn_map: Dict[str, Dict[str, Any]] = {}  # turn_id -> {spanId, startTime}

        # Current context for parent relationships
        self._current_query_span_id: Optional[str] = None
        self._current_turn_span_id: Optional[str] = None
        self._aigie = None
        self._trace_context: Optional[Any] = None

        # Statistics (delegate to session context if available)
        self._local_total_turns = 0
        self._local_total_tool_calls = 0
        self._local_total_input_tokens = 0
        self._local_total_output_tokens = 0
        self._local_total_cache_read_tokens = 0
        self._local_total_cache_creation_tokens = 0
        self._local_total_cost = 0.0

    @property
    def total_turns(self) -> int:
        if self._session_context:
            return self._session_context.total_turns
        return self._local_total_turns

    @total_turns.setter
    def total_turns(self, value: int) -> None:
        if self._session_context:
            self._session_context.total_turns = value
        else:
            self._local_total_turns = value

    @property
    def total_tool_calls(self) -> int:
        if self._session_context:
            return self._session_context.total_tool_calls
        return self._local_total_tool_calls

    @total_tool_calls.setter
    def total_tool_calls(self, value: int) -> None:
        if self._session_context:
            self._session_context.total_tool_calls = value
        else:
            self._local_total_tool_calls = value

    @property
    def total_input_tokens(self) -> int:
        if self._session_context:
            return self._session_context.total_input_tokens
        return self._local_total_input_tokens

    @total_input_tokens.setter
    def total_input_tokens(self, value: int) -> None:
        if self._session_context:
            self._session_context.total_input_tokens = value
        else:
            self._local_total_input_tokens = value

    @property
    def total_output_tokens(self) -> int:
        if self._session_context:
            return self._session_context.total_output_tokens
        return self._local_total_output_tokens

    @total_output_tokens.setter
    def total_output_tokens(self, value: int) -> None:
        if self._session_context:
            self._session_context.total_output_tokens = value
        else:
            self._local_total_output_tokens = value

    @property
    def total_cache_read_tokens(self) -> int:
        if self._session_context:
            return self._session_context.total_cache_read_tokens
        return self._local_total_cache_read_tokens

    @total_cache_read_tokens.setter
    def total_cache_read_tokens(self, value: int) -> None:
        if self._session_context:
            self._session_context.total_cache_read_tokens = value
        else:
            self._local_total_cache_read_tokens = value

    @property
    def total_cache_creation_tokens(self) -> int:
        if self._session_context:
            return self._session_context.total_cache_creation_tokens
        return self._local_total_cache_creation_tokens

    @total_cache_creation_tokens.setter
    def total_cache_creation_tokens(self, value: int) -> None:
        if self._session_context:
            self._session_context.total_cache_creation_tokens = value
        else:
            self._local_total_cache_creation_tokens = value

    @property
    def total_cost(self) -> float:
        if self._session_context:
            return self._session_context.total_cost
        return self._local_total_cost

    @total_cost.setter
    def total_cost(self, value: float) -> None:
        if self._session_context:
            self._session_context.total_cost = value
        else:
            self._local_total_cost = value

    def _get_aigie(self):
        """Lazy load Aigie client."""
        if self._aigie is None:
            from ...client import get_aigie
            self._aigie = get_aigie()
        return self._aigie

    def set_trace_context(self, trace_context: Any) -> None:
        """Set an existing trace context to use."""
        self._trace_context = trace_context
        if hasattr(trace_context, 'id'):
            self.trace_id = str(trace_context.id)

    async def handle_query_start(
        self,
        prompt: str,
        options: Dict[str, Any],
        model: Optional[str] = None,
    ) -> str:
        """
        Called when a query() call starts.

        Args:
            prompt: The prompt being sent
            options: Query options (tools, system_prompt, etc.)
            model: The model being used

        Returns:
            The query ID for tracking
        """
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return ""

        # Generate trace ID if not set - use session context if available
        if not self.trace_id:
            if self._session_context:
                self.trace_id = self._session_context.trace_id
            elif self._trace_context and hasattr(self._trace_context, 'id'):
                self.trace_id = str(self._trace_context.id)
            else:
                self.trace_id = str(uuid.uuid4())

        # Build trace name - use session context name if available
        trace_name = self.trace_name
        if not trace_name and self._session_context:
            trace_name = self._session_context.trace_name
        if not trace_name:
            # Generate descriptive name with model and prompt preview
            model_short = (model or "claude").split("-")[0].capitalize()
            if prompt:
                # Create prompt preview (30 chars, remove newlines)
                preview = prompt[:30].replace('\n', ' ').strip()
                if len(prompt) > 30:
                    preview += "..."
                trace_name = f"{model_short}: {preview}"
            else:
                trace_name = f"{model_short} Agent"

        # Extract tool names if present
        tools = options.get('tools', [])
        tool_names = []
        if tools:
            for t in tools[:10]:
                if hasattr(t, 'name'):
                    tool_names.append(t.name)
                elif isinstance(t, dict) and 'name' in t:
                    tool_names.append(t['name'])

        # Build metadata
        trace_metadata = {
            **self.metadata,
            'model': model or options.get('model', 'claude-sonnet-4-20250514'),
            'tool_count': len(tools),
            'tool_names': tool_names[:10],
            'framework': 'claude_agent_sdk',
            'max_tokens': options.get('max_tokens'),
            'max_turns': options.get('max_turns'),
        }

        # Only create trace if we don't have a trace context AND haven't created one yet
        should_create_trace = not self._trace_context
        if self._session_context and self._session_context.trace_created:
            should_create_trace = False

        if should_create_trace:
            trace_data = {
                'id': self.trace_id,
                'name': trace_name,
                'type': 'agent',
                'input': {
                    'prompt': prompt[:2000] if self.capture_messages else '[redacted]',
                    'model': trace_metadata['model'],
                    'tool_count': len(tools),
                },
                'status': 'pending',
                'tags': [*self.tags, 'claude_agent_sdk'],
                'metadata': trace_metadata,
                'start_time': datetime.now().isoformat(),
                'created_at': datetime.now().isoformat(),
            }

            if self.user_id:
                trace_data['user_id'] = self.user_id
            if self.session_id:
                trace_data['session_id'] = self.session_id

            # Send trace via buffer
            if aigie._buffer:
                logger.debug(f"[AIGIE] TRACE_CREATE: id={self.trace_id}, name={trace_name}")
                await aigie._buffer.add(EventType.TRACE_CREATE, trace_data)

            # Mark trace as created in session context
            if self._session_context:
                self._session_context.mark_trace_created()

        # Create query span
        self.query_span_id = str(uuid.uuid4())
        query_span_data = {
            'id': self.query_span_id,
            'trace_id': self.trace_id,
            'name': f'query:{trace_metadata["model"]}',
            'type': 'llm',
            'input': {
                'prompt': prompt[:2000] if self.capture_messages else '[redacted]',
                'model': trace_metadata['model'],
                'tools': tool_names,
                'system_prompt': options.get('system_prompt', '')[:500] if self.capture_messages else None,
            },
            'status': 'pending',
            'tags': self.tags or [],
            'metadata': trace_metadata,
            'model': trace_metadata['model'],
            'start_time': datetime.now().isoformat(),
            'created_at': datetime.now().isoformat(),
        }

        self._current_query_span_id = self.query_span_id

        # Set query span as current parent in session context for child span nesting
        if self._session_context:
            self._session_context.current_query_span_id = self.query_span_id
            self._session_context.set_current_parent(self.query_span_id)

        if aigie._buffer:
            logger.debug(f"[AIGIE] SPAN_CREATE: id={self.query_span_id}, name=query:{trace_metadata['model']}, parent=None (trace root)")
            await aigie._buffer.add(EventType.SPAN_CREATE, query_span_data)

        return self.query_span_id

    async def handle_query_end(
        self,
        query_id: str,
        messages: List[Any],
        result_message: Any,
        error: Optional[str] = None,
    ) -> None:
        """
        Called when a query() call completes.

        Args:
            query_id: Query ID from handle_query_start
            messages: List of messages from the conversation
            result_message: The ResultMessage with final output and costs
            error: Error message if query failed
        """
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized or not self.trace_id:
            return

        end_time = datetime.now()
        success = error is None

        # Extract usage and cost from ResultMessage
        usage = {}
        cost = 0.0
        model = None

        if result_message:
            if hasattr(result_message, 'usage'):
                usage_obj = result_message.usage
                # Handle both dict and object formats
                if isinstance(usage_obj, dict):
                    usage = {
                        'input_tokens': usage_obj.get('input_tokens', 0),
                        'output_tokens': usage_obj.get('output_tokens', 0),
                        'cache_read_input_tokens': usage_obj.get('cache_read_input_tokens', 0),
                        'cache_creation_input_tokens': usage_obj.get('cache_creation_input_tokens', 0),
                    }
                else:
                    usage = {
                        'input_tokens': getattr(usage_obj, 'input_tokens', 0),
                        'output_tokens': getattr(usage_obj, 'output_tokens', 0),
                        'cache_read_input_tokens': getattr(usage_obj, 'cache_read_input_tokens', 0),
                        'cache_creation_input_tokens': getattr(usage_obj, 'cache_creation_input_tokens', 0),
                    }
                # Update totals
                self.total_input_tokens += usage.get('input_tokens', 0)
                self.total_output_tokens += usage.get('output_tokens', 0)
                self.total_cache_read_tokens += usage.get('cache_read_input_tokens', 0)
                self.total_cache_creation_tokens += usage.get('cache_creation_input_tokens', 0)

            if hasattr(result_message, 'total_cost_usd'):
                cost = result_message.total_cost_usd or 0.0
            elif hasattr(result_message, 'model'):
                # Calculate cost from usage
                model = result_message.model
                cost = calculate_claude_cost(model, usage)

            self.total_cost += cost

            if hasattr(result_message, 'model'):
                model = result_message.model

        # Extract final output
        output = None
        if messages and self.capture_messages:
            last_message = messages[-1] if messages else None
            if last_message:
                if hasattr(last_message, 'content'):
                    content = last_message.content
                    if isinstance(content, str):
                        output = content[:2000]
                    elif isinstance(content, list):
                        # Extract text blocks
                        text_parts = []
                        for block in content:
                            if hasattr(block, 'text'):
                                text_parts.append(block.text)
                            elif isinstance(block, dict) and 'text' in block:
                                text_parts.append(block['text'])
                        output = '\n'.join(text_parts)[:2000]

        # Update query span
        if self.query_span_id:
            # Get model name for span name
            model_name = model or (self.metadata.get('model') if self.metadata else None) or 'claude-sonnet-4-20250514'

            # Get token values - prefer local usage, fallback to accumulated totals
            span_input_tokens = usage.get('input_tokens') or self.total_input_tokens
            span_output_tokens = usage.get('output_tokens') or self.total_output_tokens
            span_total_tokens = span_input_tokens + span_output_tokens
            span_cost = cost if cost > 0 else self.total_cost

            query_update = {
                'id': self.query_span_id,
                'trace_id': self.trace_id,  # Required for backend merge
                'name': f'query:{model_name}',  # Include name to handle race conditions
                'type': 'llm',  # Include type to handle race conditions
                'status': 'success' if success else 'failed',
                'output': {
                    'response': output,
                    'message_count': len(messages),
                    'usage': {
                        'input_tokens': span_input_tokens,
                        'output_tokens': span_output_tokens,
                        'total_tokens': span_total_tokens,
                    },
                    'cost': span_cost,
                },
                'end_time': end_time.isoformat(),
                'model': model_name,
                'prompt_tokens': span_input_tokens,
                'completion_tokens': span_output_tokens,
                'total_tokens': span_total_tokens,
                'total_cost': span_cost,
            }

            if error:
                query_update['error'] = error
                query_update['error_message'] = error

            if aigie._buffer:
                logger.debug(f"[AIGIE] SPAN_UPDATE: id={query_update['id']}, tokens={query_update.get('total_tokens')}, status={query_update.get('status')}")
                await aigie._buffer.add(EventType.SPAN_UPDATE, query_update)

        # Update trace with top-level token fields for backend aggregation
        update_data = {
            'id': self.trace_id,
            'status': 'success' if success else 'failed',
            'output': {
                'response': output,
                'message_count': len(messages),
                'total_tokens': self.total_input_tokens + self.total_output_tokens,
                'total_cost': self.total_cost,
                'tool_calls': self.total_tool_calls,
            },
            'end_time': end_time.isoformat(),
            # Top-level token/cost fields for backend aggregation display
            'total_tokens': self.total_input_tokens + self.total_output_tokens,
            'prompt_tokens': self.total_input_tokens,
            'completion_tokens': self.total_output_tokens,
            'total_cost': self.total_cost,
        }

        if error:
            update_data['error'] = error
            update_data['error_message'] = error

        if aigie._buffer:
            # Debug: Log trace update data
            logger.debug(f"[AIGIE] TRACE_UPDATE: id={self.trace_id}, total_tokens={update_data['total_tokens']}, cost={update_data['total_cost']}")
            await aigie._buffer.add(EventType.TRACE_UPDATE, update_data)

        # Clear query span from session context since query is done
        self._current_query_span_id = None
        if self._session_context:
            self._session_context.current_query_span_id = None
            self._session_context.set_current_parent(None)

    async def handle_tool_use_start(
        self,
        tool_name: str,
        tool_input: Dict[str, Any],
        tool_use_id: str,
    ) -> str:
        """
        Called when a tool use starts (PreToolUse hook).

        Args:
            tool_name: Name of the tool being called
            tool_input: Input arguments to the tool
            tool_use_id: Unique ID for this tool use

        Returns:
            The span ID for this tool call
        """
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized or not self.trace_id:
            return ""

        self.total_tool_calls += 1
        span_id = str(uuid.uuid4())
        start_time = datetime.now()

        # Get parent span ID - check session context first for shared parent tracking
        parent_id = self._current_turn_span_id or self._current_query_span_id
        if self._session_context:
            ctx_parent = self._session_context.get_current_parent()
            if ctx_parent:
                parent_id = ctx_parent

        self.tool_map[tool_use_id] = {
            'spanId': span_id,
            'startTime': start_time,
            'toolName': tool_name,
        }

        # Serialize tool input
        input_data = {}
        if self.capture_tool_results:
            input_data = {k: str(v)[:500] for k, v in tool_input.items()} if tool_input else {}

        span_data = {
            'id': span_id,
            'trace_id': self.trace_id,
            'parent_id': parent_id,
            'name': f'tool:{tool_name}',
            'type': 'tool',
            'input': input_data,
            'status': 'pending',
            'tags': self.tags or [],
            'metadata': {
                'toolName': tool_name,
                'toolUseId': tool_use_id,
                'framework': 'claude_agent_sdk',
            },
            'start_time': start_time.isoformat(),
            'created_at': start_time.isoformat(),
        }

        if aigie._buffer:
            logger.debug(f"[AIGIE] SPAN_CREATE: id={span_id}, name=tool:{tool_name}, parent={parent_id}")
            await aigie._buffer.add(EventType.SPAN_CREATE, span_data)

        return span_id

    async def handle_tool_use_end(
        self,
        tool_use_id: str,
        result: Any,
        is_error: bool = False,
    ) -> None:
        """
        Called when a tool use completes (PostToolUse hook).

        Args:
            tool_use_id: Unique ID for this tool use
            result: Result from the tool execution
            is_error: Whether the tool execution failed
        """
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        tool_data = self.tool_map.get(tool_use_id)
        if not tool_data:
            return

        end_time = datetime.now()
        duration = (end_time - tool_data['startTime']).total_seconds()

        output_data = {}
        if self.capture_tool_results:
            if isinstance(result, str):
                output_data['result'] = result[:1000]
            elif hasattr(result, 'model_dump'):
                output_data['result'] = str(result.model_dump())[:1000]
            else:
                output_data['result'] = str(result)[:1000]
        output_data['is_error'] = is_error

        update_data = {
            'id': tool_data['spanId'],
            'trace_id': self.trace_id,  # Required for backend merge
            'name': f"tool:{tool_data['toolName']}",  # Include name for race conditions
            'type': 'tool',  # Include type for race conditions
            'output': output_data,
            'status': 'failed' if is_error else 'success',
            'end_time': end_time.isoformat(),
            'duration_ns': int(duration * 1_000_000_000),
        }

        if is_error:
            update_data['error'] = str(result)[:500]
            update_data['error_message'] = str(result)[:500]

        if aigie._buffer:
            await aigie._buffer.add(EventType.SPAN_UPDATE, update_data)

        del self.tool_map[tool_use_id]

    async def complete_pending_tool_spans(self) -> None:
        """
        Complete any pending tool spans that weren't explicitly closed.

        This ensures all tool spans have end_time populated even if
        the corresponding ToolResultBlock was missed.
        """
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        end_time = datetime.now()
        pending_ids = list(self.tool_map.keys())

        for tool_use_id in pending_ids:
            tool_data = self.tool_map.get(tool_use_id)
            if not tool_data:
                continue

            duration = (end_time - tool_data['startTime']).total_seconds()

            update_data = {
                'id': tool_data['spanId'],
                'trace_id': self.trace_id,  # Required for backend merge
                'name': f"tool:{tool_data['toolName']}",  # Include name for race conditions
                'type': 'tool',  # Include type for race conditions
                'status': 'success',  # Assume success if not explicitly failed
                'end_time': end_time.isoformat(),
                'duration_ns': int(duration * 1_000_000_000),
            }

            if aigie._buffer:
                await aigie._buffer.add(EventType.SPAN_UPDATE, update_data)

            del self.tool_map[tool_use_id]

    async def handle_session_start(
        self,
        client: Any,
        options: Dict[str, Any],
    ) -> str:
        """
        Called when a ClaudeSDKClient session starts.

        Args:
            client: The ClaudeSDKClient instance
            options: Session options

        Returns:
            The session ID for tracking
        """
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return ""

        # Generate trace ID if not set
        if not self.trace_id:
            self.trace_id = str(uuid.uuid4())

        # Build descriptive trace name
        model = options.get('model', 'claude-sonnet-4-20250514')
        model_short = model.split("-")[0].capitalize() if model else "Claude"
        trace_name = self.trace_name or f"{model_short} Session"

        # Build metadata
        trace_metadata = {
            **self.metadata,
            'framework': 'claude_agent_sdk',
            'session_type': 'stateful',
            'model': options.get('model', 'claude-sonnet-4-20250514'),
        }

        # Create trace
        trace_data = {
            'id': self.trace_id,
            'name': trace_name,
            'type': 'agent',
            'input': {
                'session_type': 'stateful',
                'model': trace_metadata['model'],
            },
            'status': 'pending',
            'tags': [*self.tags, 'claude_agent_sdk', 'session'],
            'metadata': trace_metadata,
            'start_time': datetime.now().isoformat(),
            'created_at': datetime.now().isoformat(),
        }

        if self.user_id:
            trace_data['user_id'] = self.user_id
        if self.session_id:
            trace_data['session_id'] = self.session_id

        if aigie._buffer:
            logger.debug(f"[AIGIE] TRACE_CREATE (session): id={self.trace_id}, name={trace_name}")
            await aigie._buffer.add(EventType.TRACE_CREATE, trace_data)

        # Create session span
        self.session_span_id = str(uuid.uuid4())
        session_span_data = {
            'id': self.session_span_id,
            'trace_id': self.trace_id,
            'name': 'session',
            'type': 'chain',
            'input': {'session_type': 'stateful'},
            'status': 'pending',
            'tags': self.tags or [],
            'metadata': trace_metadata,
            'start_time': datetime.now().isoformat(),
            'created_at': datetime.now().isoformat(),
        }

        if aigie._buffer:
            logger.debug(f"[AIGIE] SPAN_CREATE: id={self.session_span_id}, name=session, parent=None (trace root)")
            await aigie._buffer.add(EventType.SPAN_CREATE, session_span_data)

        return self.trace_id

    async def handle_session_end(
        self,
        turn_count: int,
        total_cost: float,
        error: Optional[str] = None,
    ) -> None:
        """
        Called when a ClaudeSDKClient session ends.

        Args:
            turn_count: Number of conversation turns
            total_cost: Total cost in USD
            error: Error message if session failed
        """
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized or not self.trace_id:
            return

        end_time = datetime.now()
        success = error is None

        # Update session span with token data
        if self.session_span_id:
            session_tokens = self.total_input_tokens + self.total_output_tokens
            session_update = {
                'id': self.session_span_id,
                'trace_id': self.trace_id,  # Required for backend merge
                'name': 'session',  # Include name for race conditions
                'type': 'chain',  # Include type for race conditions
                'status': 'success' if success else 'failed',
                'output': {
                    'turn_count': turn_count,
                    'total_cost': total_cost,
                    'total_tokens': session_tokens,
                    'total_tool_calls': self.total_tool_calls,
                    'usage': {
                        'input_tokens': self.total_input_tokens,
                        'output_tokens': self.total_output_tokens,
                        'total_tokens': session_tokens,
                    },
                },
                'end_time': end_time.isoformat(),
                'prompt_tokens': self.total_input_tokens,
                'completion_tokens': self.total_output_tokens,
                'total_tokens': session_tokens,
                'total_cost': total_cost,
            }

            if error:
                session_update['error'] = error
                session_update['error_message'] = error

            if aigie._buffer:
                await aigie._buffer.add(EventType.SPAN_UPDATE, session_update)

        # Update trace with top-level token fields for backend aggregation
        update_data = {
            'id': self.trace_id,
            'status': 'success' if success else 'failed',
            'output': {
                'turn_count': turn_count,
                'total_cost': total_cost,
                'total_tokens': self.total_input_tokens + self.total_output_tokens,
                'total_tool_calls': self.total_tool_calls,
            },
            'end_time': end_time.isoformat(),
            # Top-level token/cost fields for backend aggregation display
            'total_tokens': self.total_input_tokens + self.total_output_tokens,
            'prompt_tokens': self.total_input_tokens,
            'completion_tokens': self.total_output_tokens,
            'total_cost': total_cost,
        }

        if error:
            update_data['error'] = error
            update_data['error_message'] = error

        if aigie._buffer:
            logger.debug(f"[AIGIE] TRACE_UPDATE (session): id={self.trace_id}, total_tokens={update_data['total_tokens']}, cost={update_data['total_cost']}")
            await aigie._buffer.add(EventType.TRACE_UPDATE, update_data)

    async def handle_turn_start(
        self,
        turn_id: str,
        user_message: str,
        turn_number: Optional[int] = None,
    ) -> str:
        """
        Called when a conversation turn starts.

        Args:
            turn_id: Unique turn identifier
            user_message: The user's message
            turn_number: Turn number in the conversation (optional, uses session count if not provided)

        Returns:
            The span ID for this turn
        """
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized or not self.trace_id:
            return ""

        # Use session context turn number if not provided
        if turn_number is None:
            if self._session_context:
                turn_number = self._session_context.increment_turn()
            else:
                self._local_total_turns += 1
                turn_number = self._local_total_turns
        else:
            # Update total_turns to match provided turn_number
            self.total_turns = turn_number

        span_id = str(uuid.uuid4())
        start_time = datetime.now()

        # Get parent span ID
        parent_id = self.session_span_id or self.query_span_id

        self.turn_map[turn_id] = {
            'spanId': span_id,
            'startTime': start_time,
            'turnNumber': turn_number,
        }

        span_data = {
            'id': span_id,
            'trace_id': self.trace_id,
            'parent_id': parent_id,
            'name': f'turn:{turn_number}',
            'type': 'chain',
            'input': {
                'user_message': user_message[:1000] if self.capture_messages else '[redacted]',
                'turn_number': turn_number,
            },
            'status': 'pending',
            'tags': self.tags or [],
            'metadata': {
                'turnId': turn_id,
                'turnNumber': turn_number,
                'framework': 'claude_agent_sdk',
            },
            'start_time': start_time.isoformat(),
            'created_at': start_time.isoformat(),
        }

        self._current_turn_span_id = span_id

        # Update session context current turn span ID and set as current parent
        if self._session_context:
            self._session_context.current_turn_span_id = span_id
            self._session_context.set_current_parent(span_id)

        if aigie._buffer:
            logger.debug(f"[AIGIE] SPAN_CREATE: id={span_id}, name=turn:{turn_number}, parent={parent_id}")
            await aigie._buffer.add(EventType.SPAN_CREATE, span_data)

        return span_id

    async def handle_turn_end(
        self,
        turn_id: str,
        output: Optional[str] = None,
        usage: Optional[Dict[str, int]] = None,
        cost: float = 0.0,
        error: Optional[str] = None,
    ) -> None:
        """
        Called when a conversation turn completes.

        Args:
            turn_id: Unique turn identifier
            output: The assistant's response (optional)
            usage: Token usage for this turn
            cost: Cost for this turn
            error: Error message if turn failed
        """
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        turn_data = self.turn_map.get(turn_id)
        if not turn_data:
            return

        end_time = datetime.now()
        duration = (end_time - turn_data['startTime']).total_seconds()

        # Update totals
        if usage:
            self.total_input_tokens += usage.get('input_tokens', 0)
            self.total_output_tokens += usage.get('output_tokens', 0)
        self.total_cost += cost

        # Determine status
        success = error is None

        update_data = {
            'id': turn_data['spanId'],
            'trace_id': self.trace_id,  # Required for backend merge
            'name': f"turn:{turn_data['turnNumber']}",  # Include name for race conditions
            'type': 'chain',  # Include type for race conditions
            'status': 'success' if success else 'error',
            'end_time': end_time.isoformat(),
            'duration_ns': int(duration * 1_000_000_000),
        }

        # Add output
        if output:
            update_data['output'] = {
                'assistant_message': output[:1000] if self.capture_messages else '[redacted]',
            }

        # Add error info
        if error:
            update_data['error'] = error
            update_data['error_message'] = error

        if usage:
            update_data['prompt_tokens'] = usage.get('input_tokens', 0)
            update_data['completion_tokens'] = usage.get('output_tokens', 0)
            update_data['total_tokens'] = usage.get('input_tokens', 0) + usage.get('output_tokens', 0)

        if cost > 0:
            update_data['total_cost'] = cost

        if aigie._buffer:
            await aigie._buffer.add(EventType.SPAN_UPDATE, update_data)

        del self.turn_map[turn_id]
        self._current_turn_span_id = None

        # Update session context - clear turn span and parent
        if self._session_context:
            self._session_context.current_turn_span_id = None
            # Revert parent to query span (if exists) since turn is done
            self._session_context.set_current_parent(self._current_query_span_id)

    async def handle_turn_error(self, turn_id: str, error: str) -> None:
        """Called when a turn encounters an error."""
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        turn_data = self.turn_map.get(turn_id)
        if not turn_data:
            return

        end_time = datetime.now()
        duration = (end_time - turn_data['startTime']).total_seconds()

        update_data = {
            'id': turn_data['spanId'],
            'trace_id': self.trace_id,  # Required for backend merge
            'name': f"turn:{turn_data['turnNumber']}",  # Include name for race conditions
            'type': 'chain',  # Include type for race conditions
            'status': 'failed',
            'error': error,
            'error_message': error,
            'end_time': end_time.isoformat(),
            'duration_ns': int(duration * 1_000_000_000),
        }

        if aigie._buffer:
            await aigie._buffer.add(EventType.SPAN_UPDATE, update_data)

        del self.turn_map[turn_id]
        self._current_turn_span_id = None

        # Update session context - clear turn span and parent
        if self._session_context:
            self._session_context.current_turn_span_id = None
            # Revert parent to query span (if exists) since turn errored
            self._session_context.set_current_parent(self._current_query_span_id)

    async def handle_llm_response(
        self,
        message: Any,
        model: Optional[str] = None,
        response_index: int = 0,
    ) -> str:
        """
        Called when an LLM response (AssistantMessage with text) is received.

        Creates a span for the reasoning/response step.

        Args:
            message: The AssistantMessage or similar message with text content
            model: The model name
            response_index: Index of this response in the conversation

        Returns:
            The span ID for this LLM response
        """
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized or not self.trace_id:
            return ""

        span_id = str(uuid.uuid4())
        start_time = datetime.now()

        # Get parent span ID
        parent_id = self._current_turn_span_id or self._current_query_span_id

        # Extract text content from message
        text_content = ""
        if hasattr(message, 'content'):
            content = message.content
            if isinstance(content, str):
                text_content = content[:500]
            elif isinstance(content, list):
                for block in content:
                    if hasattr(block, 'text'):
                        text_content = block.text[:500]
                        break
                    elif hasattr(block, 'type') and block.type == 'text':
                        text_content = getattr(block, 'text', '')[:500]
                        break

        # Get model name
        model_name = model or getattr(message, 'model', None) or 'claude'

        span_data = {
            'id': span_id,
            'trace_id': self.trace_id,
            'parent_id': parent_id,
            'name': f'llm:{model_name}',
            'type': 'llm',
            'input': {
                'response_index': response_index,
            },
            'output': {
                'text': text_content if self.capture_messages else '[redacted]',
            },
            'status': 'success',
            'tags': self.tags or [],
            'metadata': {
                'model': model_name,
                'framework': 'claude_agent_sdk',
                'response_index': response_index,
            },
            'model': model_name,
            'start_time': start_time.isoformat(),
            'end_time': start_time.isoformat(),  # Instant span
            'created_at': start_time.isoformat(),
        }

        if aigie._buffer:
            await aigie._buffer.add(EventType.SPAN_CREATE, span_data)

        return span_id

    async def handle_message(
        self,
        message_type: str,
        content: Any,
        role: str = "assistant",
    ) -> None:
        """
        Called for each message in the stream.

        Args:
            message_type: Type of message (AssistantMessage, ToolUseBlock, etc.)
            content: Message content
            role: Message role (user, assistant, tool)
        """
        # This can be used for detailed message tracking if needed
        pass

    def __repr__(self) -> str:
        return (
            f"ClaudeAgentSDKHandler("
            f"trace_id={self.trace_id}, "
            f"turns={self.total_turns}, "
            f"tool_calls={self.total_tool_calls}, "
            f"tokens={self.total_input_tokens + self.total_output_tokens}, "
            f"cost=${self.total_cost:.4f})"
        )
