"""
Anthropic provider integration for Amplitude AI SDK.

Wraps the Anthropic Python SDK to automatically track LLM usage in Amplitude.
"""
from __future__ import annotations

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, AmplitudeOrAI
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: AmplitudeOrAI,
        api_key: Optional[str] = None,
        privacy_mode: bool = False,
        content_mode: Optional[str] = None,
        **anthropic_kwargs
    ):
        """
        Initialize Anthropic client with Amplitude tracking.
        
        Args:
            amplitude: ``Amplitude`` client or ``AmplitudeAI`` instance
            api_key: Anthropic API key
            privacy_mode: Enable privacy mode (content hashing). Deprecated
                in favor of *content_mode*.
            content_mode: One of ``"full"``, ``"metadata_only"``,
                ``"customer_enriched"``. When provided, takes precedence
                over *privacy_mode*.
            **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, content_mode=content_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)

        # Only auto-generate IDs when no active session context can supply them
        if not amplitude_trace_id:
            amplitude_trace_id = str(uuid.uuid4())
        if not amplitude_conversation_id:
            from ..context import get_active_context
            ctx = get_active_context()
            if ctx is None or not ctx.session_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)
                # Prefer the versioned model name from the response
                model_name = getattr(response, 'model', None) or 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,
                        default_provider='anthropic',
                    )
                
                # 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
        response_model_name = None
        try:
            for event in stream:
                # Process event with defensive wrapper logic
                try:
                    # Handle different event types
                    if hasattr(event, 'type'):
                        # Capture the versioned model name from message_start.
                        if not response_model_name and getattr(event, 'type', None) == 'message_start':
                            _msg = getattr(event, 'message', None)
                            if _msg is not None:
                                _model = getattr(_msg, 'model', None)
                                if isinstance(_model, str):
                                    response_model_name = _model
                                    streaming_tracker._response_model = _model

                        if event.type == 'message_start':
                            message = getattr(event, 'message', None)
                            if message is not None:
                                usage = getattr(message, 'usage', None)
                                if usage is not None:
                                    token_update: Dict[str, int] = {}
                                    raw_input = getattr(usage, 'input_tokens', None)
                                    if isinstance(raw_input, int):
                                        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
                                        token_update['cache_read_input_tokens'] = cache_read
                                        token_update['cache_creation_input_tokens'] = cache_creation
                                    if token_update:
                                        streaming_tracker.update_token_usage(token_update)

                        elif 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':
                                    thinking_text = getattr(delta, 'thinking', None)
                                    if thinking_text:
                                        streaming_tracker.add_reasoning_chunk(thinking_text)
                                elif delta_type == 'input_json_delta':
                                    partial_json = getattr(delta, 'partial_json', '')
                                    if partial_json and streaming_tracker.tool_calls:
                                        streaming_tracker.tool_calls[-1]['function']['arguments'] += partial_json
                                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': ''
                                        }
                                    })

                        elif event.type == 'message_delta':
                            delta = getattr(event, 'delta', None)
                            stop_reason = getattr(delta, 'stop_reason', None)
                            if isinstance(stop_reason, str):
                                streaming_tracker.set_finish_reason(stop_reason)
                            else:
                                fallback_stop_reason = getattr(event, 'stop_reason', None)
                                if isinstance(fallback_stop_reason, str):
                                    streaming_tracker.set_finish_reason(fallback_stop_reason)

                            usage = getattr(event, 'usage', None)
                            if usage is not None:
                                output_tokens = getattr(usage, 'output_tokens', None)
                                if isinstance(output_tokens, int):
                                    streaming_tracker.update_token_usage(
                                        {'output_tokens': output_tokens}
                                    )
                        
                        elif event.type == 'message_stop':
                            if hasattr(event, 'message') and hasattr(event.message, 'usage'):
                                usage = event.message.usage
                                token_update = {}
                                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)
                            
                            if hasattr(event, 'message') and hasattr(event.message, 'stop_reason'):
                                streaming_tracker.set_finish_reason(event.message.stop_reason)
                
                except Exception as tracking_error:
                    get_logger(self.client.amplitude).warning(f"Tracking failed for event: {tracking_error}")
                
                yield event
            
            model_name = response_model_name or 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),
                        default_provider='anthropic',
                    )
                    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}")
            
            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:
            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
        finally:
            try:
                streaming_tracker.finalize_on_abandon(kwargs.get('model', 'unknown'))
            except Exception:
                pass
    
    def __getattr__(self, name):
        """Forward other attributes to original messages object."""
        return getattr(self._original, name)


# TODO: Add AsyncAnthropic wrapper with async streaming support.
# When implementing, use the same finalize_on_abandon guard pattern
# from the sync WrappedMessages above (and OpenAI's WrappedAsyncCompletions).