"""
AI Service - OpenAI and Anthropic Integration with Tool Engine
"""
from openai import AsyncOpenAI
from anthropic import AsyncAnthropic
from typing import List, Dict, Optional, Any
import logging
import json
from google import genai

from app.config import settings
from app.tools.executor import tool_executor
from app.tools.registry import tool_registry

logger = logging.getLogger(__name__)


class AIService:
    """Unified AI service for chatbot responses with tool execution"""

    def __init__(self):
        self.openai_client = AsyncOpenAI(api_key=settings.OPENAI_API_KEY) if settings.OPENAI_API_KEY else None
        self.anthropic_client = AsyncAnthropic(api_key=settings.ANTHROPIC_API_KEY) if settings.ANTHROPIC_API_KEY else None
        self.gemini_client = genai.Client(api_key=settings.GEMINI_API_KEY) if settings.GEMINI_API_KEY else None
        self.tool_executor = tool_executor

    async def get_chat_response(
            self,
            messages: List[Dict[str, str]],
            system_prompt: str = None,
            client_config: Dict = None,
            model: str = None,
            agent_type: str = "general",
            context: Dict = None
    ) -> Dict:
        """
        Get AI response for chat messages with tool support.

        Args:
            messages: List of message dicts with 'role' and 'content'
            system_prompt: Custom system prompt for the client
            client_config: Client-specific configuration
            model: AI model to use (defaults to settings.DEFAULT_AI_MODEL)
            agent_type: Type of agent (ecommerce, pos, general)
            context: Additional context (client_id, conversation_id, etc.)

        Returns:
            Dict with response content and metadata
        """
        model = model or settings.DEFAULT_AI_MODEL
        context = context or {}

        # Default system prompt
        if not system_prompt:
            system_prompt = self._build_system_prompt(agent_type, client_config)

        # Get available tools for this agent type
        tools = tool_registry.get_function_schemas(agent_type)

        try:
            if model.startswith("gemini"):
                return await self._gemini_response(messages, system_prompt, model, context)
            elif model.startswith("gpt") or model.startswith("o1"):
                return await self._openai_response_with_tools(
                    messages, system_prompt, model, tools, context
                )
            elif model.startswith("claude"):
                return await self._anthropic_response_with_tools(
                    messages, system_prompt, model, tools, context
                )
            else:
                # Default to OpenAI
                return await self._openai_response_with_tools(
                    messages, system_prompt, model, tools, context
                )
        except Exception as e:
            logger.error(f"AI Service Error: {str(e)}")
            return {
                "content": "I apologize, but I'm experiencing technical difficulties. Please try again later.",
                "error": str(e),
                "tokens_used": 0
            }

    def _build_system_prompt(self, agent_type: str, client_config: Dict = None) -> str:
        """Build system prompt based on agent type and client config"""
        base_prompts = {
            "general": """You are a helpful AI assistant for a business.
Be professional, friendly, and concise.
If you don't know something, say so politely.
Always try to help the customer or direct them to the right resource.""",

            "ecommerce": """You are an AI assistant for an e-commerce business.
You can help with:
- Checking sales and revenue data
- Looking up orders and their status
- Finding product information
- Managing inventory
- Customer inquiries

Use the available tools to fetch real data when answering questions about sales, orders, products, or customers.
Be accurate with numbers and always cite the source of your data.""",

            "pos": """You are an AI assistant for a retail point-of-sale system.
You can help with:
- Today's sales and transactions
- Looking up products and prices
- Checking inventory levels
- Processing refunds
- Customer lookups

Use the available tools to get real-time data from the POS system.
Be quick and efficient in your responses.""",

            "crm": """You are an AI assistant for customer relationship management.
You can help with:
- Customer information and history
- Order tracking
- Customer analytics
- Support ticket management

Use the available tools to access customer data and provide personalized assistance."""
        }

        system_prompt = base_prompts.get(agent_type, base_prompts["general"])

        # Apply client-specific customization
        if client_config:
            if client_config.get("company_name"):
                system_prompt = f"You are an AI assistant for {client_config['company_name']}. " + system_prompt
            if client_config.get("custom_instructions"):
                system_prompt += f"\n\nAdditional instructions: {client_config['custom_instructions']}"

        return system_prompt

    async def _openai_response_with_tools(
            self,
            messages: List[Dict[str, str]],
            system_prompt: str,
            model: str,
            tools: List[Dict],
            context: Dict
    ) -> Dict:
        """Get response from OpenAI with tool calling support"""
        if not self.openai_client:
            raise ValueError("OpenAI API key not configured")

        # Prepare messages with system prompt
        full_messages = [{"role": "system", "content": system_prompt}]
        full_messages.extend(messages[-settings.MAX_CONVERSATION_HISTORY:])

        # Convert tools to OpenAI format
        openai_tools = [
            {
                "type": "function",
                "function": tool
            }
            for tool in tools
        ] if tools else None

        # Initial request
        response = await self.openai_client.chat.completions.create(
            model=model,
            messages=full_messages,
            max_tokens=settings.MAX_TOKENS_PER_MESSAGE,
            temperature=0.7,
            tools=openai_tools,
            tool_choice="auto" if openai_tools else None
        )

        message = response.choices[0].message
        total_tokens = response.usage.total_tokens if response.usage else 0

        # Handle tool calls
        if message.tool_calls:
            # Execute each tool call
            tool_results = []
            for tool_call in message.tool_calls:
                tool_name = tool_call.function.name
                try:
                    tool_args = json.loads(tool_call.function.arguments)
                except json.JSONDecodeError:
                    tool_args = {}

                logger.info(f"OpenAI requested tool: {tool_name} with args: {tool_args}")

                result = await self.tool_executor.execute(tool_name, tool_args, context)
                tool_results.append({
                    "tool_call_id": tool_call.id,
                    "role": "tool",
                    "content": json.dumps(result.model_dump())
                })

            # Add assistant message and tool results to conversation
            full_messages.append({
                "role": "assistant",
                "content": message.content,
                "tool_calls": [
                    {
                        "id": tc.id,
                        "type": "function",
                        "function": {
                            "name": tc.function.name,
                            "arguments": tc.function.arguments
                        }
                    }
                    for tc in message.tool_calls
                ]
            })
            full_messages.extend(tool_results)

            # Get final response with tool results
            final_response = await self.openai_client.chat.completions.create(
                model=model,
                messages=full_messages,
                max_tokens=settings.MAX_TOKENS_PER_MESSAGE,
                temperature=0.7
            )

            total_tokens += final_response.usage.total_tokens if final_response.usage else 0

            return {
                "content": final_response.choices[0].message.content,
                "model": model,
                "tokens_used": total_tokens,
                "finish_reason": final_response.choices[0].finish_reason,
                "tools_used": [tc.function.name for tc in message.tool_calls]
            }

        return {
            "content": message.content,
            "model": model,
            "tokens_used": total_tokens,
            "finish_reason": response.choices[0].finish_reason,
            "tools_used": []
        }

    async def _gemini_response(
            self,
            messages: List[Dict[str, str]],
            system_prompt: str,
            model: str,
            context: Dict
    ) -> Dict:
        """Get response from Google Gemini."""

        if not self.gemini_client:
            raise ValueError("Gemini API key not configured")

        contents = []

        for msg in messages[-settings.MAX_CONVERSATION_HISTORY:]:
            role = msg.get("role", "user")
            content = msg.get("content", "")

            if role == "assistant":
                role = "model"

            if role not in ("user", "model"):
                continue

            contents.append({
                "role": role,
                "parts": [
                    {
                        "text": content
                    }
                ]
            })

        response = await self.gemini_client.aio.models.generate_content(
            model=model,
            contents=contents,
            config={
                "system_instruction": system_prompt,
                "max_output_tokens": settings.MAX_TOKENS_PER_MESSAGE,
                "temperature": 0.7,
            }
        )

        return {
            "content": response.text,
            "model": model,
            "tokens_used": 0,
            "finish_reason": "stop",
            "tools_used": []
        }

    async def _anthropic_response_with_tools(
            self,
            messages: List[Dict[str, str]],
            system_prompt: str,
            model: str,
            tools: List[Dict],
            context: Dict
    ) -> Dict:
        """Get response from Anthropic Claude with tool calling support"""
        if not self.anthropic_client:
            raise ValueError("Anthropic API key not configured")

        # Convert messages format for Anthropic
        anthropic_messages = []
        for msg in messages[-settings.MAX_CONVERSATION_HISTORY:]:
            anthropic_messages.append({
                "role": msg["role"],
                "content": msg["content"]
            })

        # Convert tools to Anthropic format
        anthropic_tools = [
            {
                "name": tool["name"],
                "description": tool["description"],
                "input_schema": tool["parameters"]
            }
            for tool in tools
        ] if tools else None

        # Initial request
        response = await self.anthropic_client.messages.create(
            model=model,
            max_tokens=settings.MAX_TOKENS_PER_MESSAGE,
            system=system_prompt,
            messages=anthropic_messages,
            tools=anthropic_tools
        )

        total_tokens = response.usage.input_tokens + response.usage.output_tokens

        # Check for tool use
        tool_use_blocks = [block for block in response.content if block.type == "tool_use"]

        if tool_use_blocks:
            # Execute tool calls
            tool_results = []
            tools_used = []

            for tool_block in tool_use_blocks:
                tool_name = tool_block.name
                tool_args = tool_block.input
                tools_used.append(tool_name)

                logger.info(f"Claude requested tool: {tool_name} with args: {tool_args}")

                result = await self.tool_executor.execute(tool_name, tool_args, context)
                tool_results.append({
                    "type": "tool_result",
                    "tool_use_id": tool_block.id,
                    "content": json.dumps(result.model_dump())
                })

            # Add assistant response and tool results
            anthropic_messages.append({
                "role": "assistant",
                "content": response.content
            })
            anthropic_messages.append({
                "role": "user",
                "content": tool_results
            })

            # Get final response
            final_response = await self.anthropic_client.messages.create(
                model=model,
                max_tokens=settings.MAX_TOKENS_PER_MESSAGE,
                system=system_prompt,
                messages=anthropic_messages
            )

            total_tokens += final_response.usage.input_tokens + final_response.usage.output_tokens

            # Extract text content
            text_content = ""
            for block in final_response.content:
                if hasattr(block, "text"):
                    text_content += block.text

            return {
                "content": text_content,
                "model": model,
                "tokens_used": total_tokens,
                "finish_reason": final_response.stop_reason,
                "tools_used": tools_used
            }

        # No tool use - extract text content
        text_content = ""
        for block in response.content:
            if hasattr(block, "text"):
                text_content += block.text

        return {
            "content": text_content,
            "model": model,
            "tokens_used": total_tokens,
            "finish_reason": response.stop_reason,
            "tools_used": []
        }

    async def get_simple_response(
            self,
            messages: List[Dict[str, str]],
            system_prompt: str = None,
            client_config: Dict = None,
            model: str = None
    ) -> Dict:
        """
        Get AI response without tools (original behavior).
        Use this for simple Q&A without business data access.
        """
        model = model or settings.DEFAULT_AI_MODEL

        if not system_prompt:
            system_prompt = """You are a helpful AI assistant for a business.
            Be professional, friendly, and concise.
            If you don't know something, say so politely.
            Always try to help the customer or direct them to the right resource."""

        if client_config:
            if client_config.get("company_name"):
                system_prompt = f"You are an AI assistant for {client_config['company_name']}. " + system_prompt
            if client_config.get("custom_instructions"):
                system_prompt += f"\n\nAdditional instructions: {client_config['custom_instructions']}"

        try:
            if model.startswith("gpt") or model.startswith("o1"):
                return await self._openai_simple_response(messages, system_prompt, model)
            elif model.startswith("claude"):
                return await self._anthropic_simple_response(messages, system_prompt, model)
            else:
                return await self._openai_simple_response(messages, system_prompt, model)
        except Exception as e:
            logger.error(f"AI Service Error: {str(e)}")
            return {
                "content": "I apologize, but I'm experiencing technical difficulties. Please try again later.",
                "error": str(e),
                "tokens_used": 0
            }

    async def _openai_simple_response(
            self,
            messages: List[Dict[str, str]],
            system_prompt: str,
            model: str
    ) -> Dict:
        """Get simple response from OpenAI without tools"""
        if not self.openai_client:
            raise ValueError("OpenAI API key not configured")

        full_messages = [{"role": "system", "content": system_prompt}]
        full_messages.extend(messages[-settings.MAX_CONVERSATION_HISTORY:])

        response = await self.openai_client.chat.completions.create(
            model=model,
            messages=full_messages,
            max_tokens=settings.MAX_TOKENS_PER_MESSAGE,
            temperature=0.7
        )

        return {
            "content": response.choices[0].message.content,
            "model": model,
            "tokens_used": response.usage.total_tokens if response.usage else 0,
            "finish_reason": response.choices[0].finish_reason
        }

    async def _anthropic_simple_response(
            self,
            messages: List[Dict[str, str]],
            system_prompt: str,
            model: str
    ) -> Dict:
        """Get simple response from Anthropic without tools"""
        if not self.anthropic_client:
            raise ValueError("Anthropic API key not configured")

        anthropic_messages = []
        for msg in messages[-settings.MAX_CONVERSATION_HISTORY:]:
            anthropic_messages.append({
                "role": msg["role"],
                "content": msg["content"]
            })

        response = await self.anthropic_client.messages.create(
            model=model,
            max_tokens=settings.MAX_TOKENS_PER_MESSAGE,
            system=system_prompt,
            messages=anthropic_messages
        )

        return {
            "content": response.content[0].text,
            "model": model,
            "tokens_used": response.usage.input_tokens + response.usage.output_tokens,
            "finish_reason": response.stop_reason
        }


# Singleton instance
ai_service = AIService()