# llm_providers/claude_provider.py
import anthropic
import logging
import requests
import base64
from mcpserver.mcp_client import get_llm_function_definitions

logger = logging.getLogger(__name__)

def get_mcp_tools_anthropic():
    """Returns MCP tools definitions for Claude Anthropic"""
    tools = []

    # Add MCP functions
    mcp_functions = get_llm_function_definitions()
    for func_def in mcp_functions:
        logger.info(f"Adding Anthropic MCP function: {func_def['name']}")

        # Convert properties to Anthropic format
        properties = {}
        for prop_name, prop_info in func_def['parameters']['properties'].items():
            properties[prop_name] = {
                "type": "string",  # MCP functions typically use string parameters
                "description": prop_info['description']
            }

        anthropic_func = {
            "name": func_def['name'],
            "description": func_def['description'],
            "input_schema": {
                "type": "object",
                "properties": properties,
                "required": func_def['parameters']['required']
            }
        }
        tools.append(anthropic_func)

    logger.info(f"Created {len(tools)} MCP tools for Anthropic")
    return tools

def call_claude_stream(chat_history, model_name, api_key, system_prompt):
    """
    Sends messages to Anthropic API and returns the response as a stream.
    Handles message formatting and error handling.
    """
    client = anthropic.Anthropic(
        api_key=api_key
    )
    logger.info(f"API model (Anthropic): {model_name}")

    # Convert message format for Anthropic
    anthropic_messages = []
    for msg in chat_history:
        if msg["role"] == "system":
            continue  # Handle system message separately
        
        role = "assistant" if msg["role"] == "assistant" else "user"
        
        content = msg["content"]
        if isinstance(content, str):
            anthropic_messages.append({"role": role, "content": content})
        elif isinstance(content, list):
            content_parts = []
            for item in content:
                if item.get('type') == 'text':
                    text_content = item.get('text', '').strip()
                    if text_content:
                        content_parts.append({"type": "text", "text": text_content})
                elif item.get('type') == 'image_url':
                    image_url = item['image_url']['url']
                    try:
                        if image_url.startswith('data:image'):
                            header, encoded = image_url.split(',', 1)
                            media_type = header.split(':')[1].split(';')[0]
                            image_data = encoded
                            
                            content_parts.append({
                                "type": "image",
                                "source": {
                                    "type": "base64",
                                    "media_type": media_type,
                                    "data": image_data
                                }
                            })
                        else:
                            logger.info(f"Fetching image from URL for Claude: {image_url}")
                            response = requests.get(image_url)
                            response.raise_for_status()
                            
                            content_type = response.headers.get('content-type', 'image/jpeg')
                            if 'image' not in content_type:
                                content_type = 'image/jpeg'
                            
                            image_data = base64.b64encode(response.content).decode('utf-8')
                            
                            content_parts.append({
                                "type": "image",
                                "source": {
                                    "type": "base64",
                                    "media_type": content_type,
                                    "data": image_data
                                }
                            })
                            
                    except Exception as e:
                        logger.error(f"Error processing image URL for Claude: {e}", exc_info=True)
                        content_parts.append({"type": "text", "text": f"(не удалось загрузить изображение: {e})"})
            
            if content_parts:
                anthropic_messages.append({"role": role, "content": content_parts})
            else:
                anthropic_messages.append({"role": role, "content": str(content)})
        else:
            anthropic_messages.append({"role": role, "content": str(content)})
    
    logger.debug(f"Messages sent to Anthropic: {anthropic_messages}")
    try:
        logger.info(f"Calling Anthropic API with model: {model_name}")
        
        with client.messages.stream(
            model=model_name,
            messages=anthropic_messages,
            system=system_prompt,
            max_tokens=4000,
            thinking={
                "type": "enabled",
                "budget_tokens": 3000
            },
            timeout=300
        ) as stream:
            message_text = ""
            for text in stream.text_stream:
                message_text += text
                yield text
            
    except anthropic.APIError as ae:
        logger.error(f"API error: {ae}", exc_info=True)
        yield f"API error: {str(ae)}"
    except anthropic.RateLimitError as rle:
        logger.error(f"API rate limit exceeded: {rle}", exc_info=True)
        yield "API rate limit exceeded. Please try again later."
    except anthropic.APIConnectionError as ace:
        logger.error(f"API connection error: {ace}", exc_info=True)
        yield "Failed to connect to API. Please check your internet connection."
    except Exception as e:
        logger.error(f"Unexpected error calling API: {e}", exc_info=True)
        yield f"An error occurred while calling the API: {str(e)}"