"""
Anthropic provider integration for Amplitude AI SDK.

Wraps the Anthropic Python SDK to automatically track LLM usage in Amplitude.
"""

import time
import uuid
from typing import Any, Dict, List, Optional, Union, Iterator
try:
    import anthropic
    from anthropic import Anthropic as AnthropicClient
    from anthropic.types import Message, MessageStreamEvent
    ANTHROPIC_AVAILABLE = True
except ImportError:
    ANTHROPIC_AVAILABLE = False
    AnthropicClient = object  # Fallback for type hints
    Message = object  # Fallback for type hints
    MessageStreamEvent = object  # Fallback for type hints

from amplitude import Amplitude
from .base import BaseAIProvider
from ..core.privacy import PrivacyConfig
from ..exceptions import ProviderError
from ..utils.logger import get_logger
from ..utils.tokens import count_tokens
from ..utils.costs import calculate_cost


class Anthropic(BaseAIProvider, AnthropicClient):
    """
    Anthropic client wrapper that automatically tracks usage in Amplitude.
    
    This class inherits from both BaseAIProvider (for tracking) and Anthropic's client
    (for API functionality), providing seamless drop-in replacement.
    """
    
    def __init__(
        self,
        amplitude: Amplitude,
        api_key: Optional[str] = None,
        privacy_mode: bool = False,
        **anthropic_kwargs
    ):
        """
        Initialize Anthropic client with Amplitude tracking.
        
        Args:
            amplitude: Amplitude client instance
            api_key: Anthropic API key
            privacy_mode: Enable privacy mode (content hashing)
            **anthropic_kwargs: Additional arguments passed to Anthropic client
        """
        if not ANTHROPIC_AVAILABLE:
            raise ProviderError(
                "Anthropic SDK not installed. Install with: pip install anthropic"
            )
        
        # Initialize base provider
        privacy_config = PrivacyConfig(privacy_mode=privacy_mode)
        BaseAIProvider.__init__(
            self,
            amplitude=amplitude,
            privacy_config=privacy_config
        )
        
        # Initialize Anthropic client
        self._propagate_context = anthropic_kwargs.pop("propagate_context", False)
        AnthropicClient.__init__(self, api_key=api_key, **anthropic_kwargs)
        
        # Wrap the messages interface
        self._original_messages = self.messages
        self.messages = WrappedMessages(self)
    
    def __getattr__(self, name):
        """Forward unimplemented attributes to original Anthropic client.
        
        This allows access to APIs we haven't wrapped yet (completions, count_tokens, etc.)
        They will work but won't be tracked in Amplitude.
        """
        # Avoid recursion on internal attributes
        if name.startswith('_') or name in ('amplitude', 'privacy_config', 'provider_name'):
            raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
        # Forward to parent Anthropic client
        return getattr(super(AnthropicClient, self), name)
    
    def _get_provider_name(self) -> str:
        return "anthropic"
    
    def _extract_model_name(self, **kwargs) -> str:
        return kwargs.get('model', 'unknown')
    
    def _extract_messages(self, **kwargs) -> List[Dict[str, Any]]:
        return kwargs.get('messages', [])

    def _extract_system_prompt(self, **kwargs) -> Optional[str]:
        """Extract system prompt from Anthropic's top-level ``system`` kwarg.

        Anthropic passes the system prompt as ``system=...`` (str or list of
        content blocks), *not* inside ``messages``.
        """
        system = kwargs.get("system")
        if system is None:
            return None
        if isinstance(system, str):
            return system
        # List of content blocks: [{"type": "text", "text": "..."}]
        if isinstance(system, list):
            parts = []
            for block in system:
                if isinstance(block, dict) and block.get("type") == "text":
                    parts.append(block.get("text", ""))
                elif isinstance(block, str):
                    parts.append(block)
            return "".join(parts) if parts else None
        return str(system)

    def _format_messages_for_tracking(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Format Anthropic messages for tracking (apply privacy settings)."""
        from ..core.privacy import extract_text_from_anthropic_message, get_text_from_llm_message, sanitize_any_content
        
        formatted = []
        for msg in messages:
            formatted_msg = msg.copy()
            
            # Extract text and apply universal sanitization
            text_content = extract_text_from_anthropic_message(msg)
            if text_content:
                sanitized_result = sanitize_any_content(
                    text_content,
                    privacy_mode=self.privacy_config.privacy_mode,
                    redact_pii=self.privacy_config.redact_pii
                )
                
                if self.privacy_config.privacy_mode:
                    formatted_msg['content_hash'] = sanitized_result.get('content_hash')
                    formatted_msg.pop('content', None)
                else:
                    # Keep structure simple for tracking
                    formatted_msg['content'] = get_text_from_llm_message(sanitized_result['$llm_message'])
            
            formatted.append(formatted_msg)
        return formatted
    
    def _extract_response_content(self, response: Any) -> str:
        """Extract text content from Anthropic response."""
        if hasattr(response, 'content') and response.content:
            content_parts = []
            for content_block in response.content:
                if hasattr(content_block, 'text'):
                    content_parts.append(content_block.text)
                elif hasattr(content_block, 'type') and content_block.type == 'text':
                    content_parts.append(getattr(content_block, 'text', ''))
            return ''.join(content_parts)
        return ""
    
    def _extract_reasoning_content(self, response: Any) -> Optional[str]:
        """Extract thinking/reasoning content from Anthropic extended thinking.

        Anthropic models with extended thinking produce content blocks with
        ``type == "thinking"`` containing the model's chain-of-thought.
        """
        if not hasattr(response, 'content') or not response.content:
            return None

        thinking_parts = []
        for block in response.content:
            block_type = getattr(block, 'type', None)
            if block_type == 'thinking':
                # Prefer 'thinking' attribute; fall back to 'text' only if 'thinking' is absent.
                # Use 'is not None' to preserve empty-string values.
                text = getattr(block, 'thinking', None)
                if text is None:
                    text = getattr(block, 'text', None)
                if text:
                    thinking_parts.append(text)

        return ''.join(thinking_parts) if thinking_parts else None

    def _extract_token_usage(self, response: Any) -> Dict[str, Optional[int]]:
        """Extract token usage from Anthropic response, normalizing input_tokens to total.

        Anthropic's raw API returns input_tokens as non-cached only. This override
        normalizes it to TOTAL (non-cached + cache_read + cache_creation) so that
        emitted Agent events are consistent with cost calculations.
        """
        usage_dict = super()._extract_token_usage(response)

        # Normalize: Anthropic raw input_tokens excludes cached tokens.
        # Add cache tokens to get the true total.
        cache_read = usage_dict.get('cache_read_input_tokens') or 0
        cache_creation = usage_dict.get('cache_creation_input_tokens') or 0
        if usage_dict.get('input_tokens') is not None and (cache_read or cache_creation):
            usage_dict['input_tokens'] = usage_dict['input_tokens'] + cache_read + cache_creation
            # Recalculate total_tokens with the normalized input
            if usage_dict.get('output_tokens') is not None:
                usage_dict['total_tokens'] = usage_dict['input_tokens'] + usage_dict['output_tokens']

        return usage_dict

    def _extract_finish_reason(self, response: Any) -> Optional[str]:
        """Extract finish reason from Anthropic response."""
        if hasattr(response, 'stop_reason'):
            return response.stop_reason
        return None
    
    def _extract_tool_calls(self, response: Any) -> Optional[List[Dict[str, Any]]]:
        """Extract tool calls from Anthropic response."""
        if hasattr(response, 'content') and response.content:
            tool_calls = []
            for content_block in response.content:
                if hasattr(content_block, 'type') and content_block.type == 'tool_use':
                    tool_calls.append({
                        'type': 'function',
                        'id': getattr(content_block, 'id', ''),
                        'function': {
                            'name': getattr(content_block, 'name', ''),
                            'arguments': str(getattr(content_block, 'input', {}))
                        }
                    })
            return tool_calls if tool_calls else None
        return None


class WrappedMessages:
    """Wrapper for Anthropic messages that adds tracking."""
    
    def __init__(self, client: Anthropic):
        self.client = client
        self._original = client._original_messages
    
    def create(
        self,
        *,
        amplitude_user_id: Optional[str] = None,
        amplitude_conversation_id: Optional[str] = None,
        amplitude_trace_id: Optional[str] = None,
        amplitude_turn_id: Optional[int] = None,
        amplitude_event_properties: Optional[Dict[str, Any]] = None,
        amplitude_user_properties: Optional[Dict[str, Any]] = None,
        amplitude_groups: Optional[Dict[str, Any]] = None,
        **kwargs
    ):
        """
        Create a message with automatic Amplitude tracking.
        
        Args:
            amplitude_user_id: User ID for tracking
            amplitude_conversation_id: Conversation ID for linking events
            amplitude_trace_id: Trace ID for linking related events
            amplitude_turn_id: Turn number in conversation
            amplitude_event_properties: Additional event properties
            amplitude_user_properties: User properties to set
            amplitude_groups: Group analytics properties
            **kwargs: Arguments passed to Anthropic's create method
        """
        amplitude_user_id = amplitude_user_id or getattr(self.client, '_default_amplitude_user_id', None)

        # Generate missing IDs
        if not amplitude_trace_id:
            amplitude_trace_id = str(uuid.uuid4())
        if not amplitude_conversation_id:
            amplitude_conversation_id = str(uuid.uuid4())
        
        # Check if streaming
        is_streaming = kwargs.get('stream', False)
        
        if is_streaming:
            return self._create_streaming(
                amplitude_user_id=amplitude_user_id,
                amplitude_conversation_id=amplitude_conversation_id,
                amplitude_trace_id=amplitude_trace_id,
                amplitude_turn_id=amplitude_turn_id,
                amplitude_event_properties=amplitude_event_properties,
                amplitude_user_properties=amplitude_user_properties,
                amplitude_groups=amplitude_groups,
                **kwargs
            )
        else:
            return self._create_non_streaming(
                amplitude_user_id=amplitude_user_id,
                amplitude_conversation_id=amplitude_conversation_id,
                amplitude_trace_id=amplitude_trace_id,
                amplitude_turn_id=amplitude_turn_id,
                amplitude_event_properties=amplitude_event_properties,
                amplitude_user_properties=amplitude_user_properties,
                amplitude_groups=amplitude_groups,
                **kwargs
            )
    
    def _create_non_streaming(
        self,
        amplitude_user_id: Optional[str] = None,
        amplitude_conversation_id: Optional[str] = None,
        amplitude_trace_id: Optional[str] = None,
        amplitude_turn_id: Optional[int] = None,
        amplitude_event_properties: Optional[Dict[str, Any]] = None,
        amplitude_user_properties: Optional[Dict[str, Any]] = None,
        amplitude_groups: Optional[Dict[str, Any]] = None,
        **kwargs
    ) -> Message:
        """Handle non-streaming message creation with tracking."""
        if getattr(self.client, "_propagate_context", False):
            from ..propagation import inject_context
            kwargs["extra_headers"] = inject_context(kwargs.get("extra_headers"))

        start_time = time.time()
        
        try:
            # Make the actual Anthropic API call
            response = self._original.create(**kwargs)
            
            # Calculate latency
            end_time = time.time()
            latency_ms = (end_time - start_time) * 1000
            
            # Track the completion if user ID is provided
            if amplitude_user_id:
                # Calculate cost
                token_usage = self.client._extract_token_usage(response)
                model_name = kwargs.get('model', 'unknown')
                
                total_cost_usd = None
                if token_usage['input_tokens'] is not None and token_usage['output_tokens'] is not None:
                    # _extract_token_usage already normalizes input_tokens to total
                    # (including cached tokens), so pass directly to calculate_cost.
                    total_cost_usd = calculate_cost(
                        model_name=model_name,
                        input_tokens=token_usage['input_tokens'],
                        output_tokens=token_usage['output_tokens'] or 0,
                        cache_read_input_tokens=token_usage['cache_read_input_tokens'] or 0,
                        cache_creation_input_tokens=token_usage['cache_creation_input_tokens'] or 0,
                    )
                
                # Track the completion
                self.client.track_completion(
                    user_id=amplitude_user_id,
                    response=response,
                    latency_ms=latency_ms,
                    conversation_id=amplitude_conversation_id,
                    trace_id=amplitude_trace_id,
                    turn_id=amplitude_turn_id,
                    event_properties=amplitude_event_properties,
                    user_properties=amplitude_user_properties,
                    groups=amplitude_groups,
                    total_cost_usd=total_cost_usd,
                    **kwargs
                )
            
            return response
            
        except Exception as e:
            # Calculate latency even for errors
            end_time = time.time()
            latency_ms = (end_time - start_time) * 1000
            
            # Track error if user ID is provided
            if amplitude_user_id:
                self.client.track_completion(
                    user_id=amplitude_user_id,
                    response=None,
                    latency_ms=latency_ms,
                    conversation_id=amplitude_conversation_id,
                    trace_id=amplitude_trace_id,
                    turn_id=amplitude_turn_id,
                    is_error=True,
                    error_message=str(e),
                    event_properties=amplitude_event_properties,
                    user_properties=amplitude_user_properties,
                    groups=amplitude_groups,
                    **kwargs
                )
            
            # Re-raise the original exception
            raise
    
    def _create_streaming(
        self,
        amplitude_user_id: Optional[str] = None,
        amplitude_conversation_id: Optional[str] = None,
        amplitude_trace_id: Optional[str] = None,
        amplitude_turn_id: Optional[int] = None,
        amplitude_event_properties: Optional[Dict[str, Any]] = None,
        amplitude_user_properties: Optional[Dict[str, Any]] = None,
        amplitude_groups: Optional[Dict[str, Any]] = None,
        **kwargs
    ) -> Iterator[MessageStreamEvent]:
        """Handle streaming message creation with tracking."""
        if getattr(self.client, "_propagate_context", False):
            from ..propagation import inject_context
            kwargs["extra_headers"] = inject_context(kwargs.get("extra_headers"))

        if not amplitude_user_id:
            # If no user ID, just return the original stream without tracking
            return self._original.create(**kwargs)
        
        # Create streaming tracker
        streaming_tracker = self.client.create_streaming_tracker(
            user_id=amplitude_user_id,
            conversation_id=amplitude_conversation_id,
            trace_id=amplitude_trace_id,
            turn_id=amplitude_turn_id,
            event_properties=amplitude_event_properties,
            user_properties=amplitude_user_properties,
            groups=amplitude_groups,
            **kwargs
        )
        
        # Separate API call from processing to avoid masking Anthropic errors
        try:
            # Make the streaming API call - this is the only part that should fail "legitimately"
            stream = self._original.create(**kwargs)
        except Exception as e:
            # This is definitely an Anthropic API error
            streaming_tracker.set_error(str(e))
            streaming_tracker.finalize(kwargs.get('model', 'unknown'))
            raise
        
        # Process stream with defensive error handling for wrapper logic
        try:
            for event in stream:
                # Process event with defensive wrapper logic
                try:
                    # Handle different event types
                    if hasattr(event, 'type'):
                        if event.type == 'content_block_delta':
                            delta = getattr(event, 'delta', None)
                            if delta is not None:
                                delta_type = getattr(delta, 'type', None)
                                if delta_type == 'thinking_delta':
                                    # Extended thinking content
                                    thinking_text = getattr(delta, 'thinking', None)
                                    if thinking_text:
                                        streaming_tracker.add_reasoning_chunk(thinking_text)
                                elif hasattr(delta, 'text') and delta.text:
                                    streaming_tracker.add_content_chunk(delta.text)
                        
                        elif event.type == 'content_block_start':
                            if hasattr(event, 'content_block'):
                                if hasattr(event.content_block, 'type') and event.content_block.type == 'tool_use':
                                    streaming_tracker.add_tool_call({
                                        'type': 'function',
                                        'id': getattr(event.content_block, 'id', ''),
                                        'function': {
                                            'name': getattr(event.content_block, 'name', ''),
                                            'arguments': str(getattr(event.content_block, 'input', {}))
                                        }
                                    })
                        
                        elif event.type == 'message_stop':
                            # Extract final usage if available
                            if hasattr(event, 'message') and hasattr(event.message, 'usage'):
                                usage = event.message.usage
                                token_update = {}
                                # Anthropic raw API returns input_tokens as non-cached only.
                                # Normalize to TOTAL (including cached) for consistency with
                                # calculate_cost which expects total input tokens.
                                raw_input = getattr(usage, 'input_tokens', 0) or 0
                                cache_read = getattr(usage, 'cache_read_input_tokens', 0) or 0
                                cache_creation = getattr(usage, 'cache_creation_input_tokens', 0) or 0
                                token_update['input_tokens'] = raw_input + cache_read + cache_creation
                                if hasattr(usage, 'output_tokens'):
                                    token_update['output_tokens'] = usage.output_tokens
                                token_update['cache_read_input_tokens'] = cache_read
                                token_update['cache_creation_input_tokens'] = cache_creation
                                
                                streaming_tracker.update_token_usage(token_update)
                            
                            # Extract stop reason
                            if hasattr(event, 'message') and hasattr(event.message, 'stop_reason'):
                                streaming_tracker.set_finish_reason(event.message.stop_reason)
                
                except Exception as tracking_error:
                    # Log wrapper errors but don't break the stream
                    get_logger(self.client.amplitude).warning(f"Tracking failed for event: {tracking_error}")
                
                # Always yield the original event regardless of tracking success/failure
                yield event
            
            # Calculate cost if we have token usage
            model_name = kwargs.get('model', 'unknown')
            if streaming_tracker.token_usage:
                try:
                    total_cost_usd = calculate_cost(
                        model_name=model_name,
                        input_tokens=streaming_tracker.token_usage.get('input_tokens', 0),
                        output_tokens=streaming_tracker.token_usage.get('output_tokens', 0),
                        cache_read_input_tokens=streaming_tracker.token_usage.get('cache_read_input_tokens', 0),
                        cache_creation_input_tokens=streaming_tracker.token_usage.get('cache_creation_input_tokens', 0)
                    )
                    streaming_tracker.token_usage['total_cost_usd'] = total_cost_usd
                except Exception as cost_error:
                    get_logger(self.client.amplitude).warning(f"Cost calculation failed: {cost_error}")
            
            # Finalize tracking
            try:
                streaming_tracker.finalize(model_name)
            except Exception as finalize_error:
                get_logger(self.client.amplitude).warning(f"Streaming tracker finalization failed: {finalize_error}")
                
        except Exception as e:
            # This catches errors in the stream iteration itself (network issues, etc.)
            # These could be from Anthropic or from our processing, but we should still track them
            try:
                streaming_tracker.set_error(str(e))
                streaming_tracker.finalize(kwargs.get('model', 'unknown'))
            except Exception as tracking_error:
                get_logger(self.client.amplitude).warning(f"Error tracking failed: {tracking_error}")
            raise
    
    def __getattr__(self, name):
        """Forward other attributes to original messages object."""
        return getattr(self._original, name)