"""LangChain-based chat service."""

import json
import asyncio
import os
from typing import AsyncGenerator, Optional, List, Dict, Any
from sqlalchemy.orm import Session

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult

from ..core.config import settings
from ..models.message import MessageRole
from utils.util_schemas import ChatResponse, StreamChunk, MessageResponse
from utils.util_exceptions import ChatServiceError, OpenAIError, AuthenticationError, RateLimitError
from loguru import logger
from .conversation import ConversationService


class StreamingCallbackHandler(BaseCallbackHandler):
    """Custom callback handler for streaming responses."""
    
    def __init__(self):
        self.tokens = []
        
    def on_llm_new_token(self, token: str, **kwargs) -> None:
        """Handle new token from LLM."""
        self.tokens.append(token)
        
    def get_response(self) -> str:
        """Get the complete response."""
        return "".join(self.tokens)
        
    def clear(self):
        """Clear the tokens."""
        self.tokens = []


class LangChainChatService:
    """LangChain-based chat service for AI model integration."""
    
    def __init__(self, session: Session):
        self.session = session
        self.conversation_service = ConversationService(session)
        
    async def initialize(self):
        from ..core.new_agent import new_agent
        
        # Initialize LangChain ChatOpenAI    
        self.llm = await new_agent(self.session, streaming=False)    
        self.session.desc = "LangChainChatService初始化 - llm 实例化完毕"
        
        # Streaming LLM for stream responses    
        self.streaming_llm = await new_agent(self.session, streaming=True)    
        self.session.desc = "LangChainChatService初始化 - streaming_llm 实例化完毕"
        
        self.streaming_handler = StreamingCallbackHandler()
        self.session.desc = "LangChainChatService初始化 - streaming_handler 实例化完毕"
    
    def _prepare_langchain_messages(self, conversation, history: List) -> List:
        """Prepare messages for LangChain format."""
        messages = []   
        
        # Add system message if conversation has system prompt
        if hasattr(conversation, 'system_prompt') and conversation.system_prompt:
            messages.append(SystemMessage(content=conversation.system_prompt))
        else:
            # Default system message
            messages.append(SystemMessage(
                content="You are a helpful AI assistant. Please provide accurate and helpful responses."
            ))
        
        # Add conversation history
        for msg in history[:-1]:  # Exclude the last message (current user message)
            if msg.role == MessageRole.USER:
                messages.append(HumanMessage(content=msg.content))
            elif msg.role == MessageRole.ASSISTANT:
                messages.append(AIMessage(content=msg.content))
        
        # Add current user message
        if history:
            last_msg = history[-1]
            if last_msg.role == MessageRole.USER:
                messages.append(HumanMessage(content=last_msg.content))
        
        return messages
    
    async def chat(
        self, 
        conversation_id: int, 
        message: str, 
        stream: bool = False,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None
    ) -> ChatResponse:
        """Send a message and get AI response using LangChain."""
        logger.info(f"Processing LangChain chat request for conversation {conversation_id}")
        
        try:
            # Get conversation details
            conversation = await self.conversation_service.get_conversation(conversation_id)
            if not conversation:
                raise ChatServiceError("Conversation not found")
            
            # Add user message to database
            user_message = await self.conversation_service.add_message(
                conversation_id=conversation_id,
                content=message,
                role=MessageRole.USER
            )
            
            # Get conversation history for context
            history = await self.conversation_service.get_conversation_history(
                conversation_id, limit=20
            )
            
            # Prepare messages for LangChain
            langchain_messages = self._prepare_langchain_messages(conversation, history)
            
            # Update LLM parameters if provided
            llm_to_use = self.llm
            if temperature is not None or max_tokens is not None:
                llm_config = await settings.llm.get_current_config()
                llm_to_use = ChatOpenAI(
                    model=llm_config["model"],
                    openai_api_key=llm_config["api_key"],
                    openai_api_base=llm_config["base_url"],
                    temperature=temperature if temperature is not None else float(conversation.temperature),
                    max_tokens=max_tokens if max_tokens is not None else conversation.max_tokens,
                    streaming=False
                )
            
            # Call LangChain LLM
            response = await llm_to_use.ainvoke(langchain_messages)
            
            # Extract response content
            assistant_content = response.content
            
            # Add assistant message to database
            assistant_message = await self.conversation_service.add_message(
                conversation_id=conversation_id,
                content=assistant_content,
                role=MessageRole.ASSISTANT,
                message_metadata={
                    "model": llm_to_use.model_name,
                    "langchain_version": "0.1.0",
                    "provider": "langchain_openai"
                }
            )
            
            # Update conversation timestamp
            await self.conversation_service.update_conversation_timestamp(conversation_id)
            
            logger.info(f"Successfully processed LangChain chat request for conversation {conversation_id}")
            
            return ChatResponse(
                user_message=MessageResponse.from_orm(user_message),
                assistant_message=MessageResponse.from_orm(assistant_message),
                total_tokens=None,  # LangChain doesn't provide token count by default
                model_used=llm_to_use.model_name
            )
            
        except Exception as e:
            logger.error(f"Failed to process LangChain chat request for conversation {conversation_id}: {str(e)}", exc_info=True)
            
            # Classify error types for better handling
            error_type = type(e).__name__
            error_message = self._format_error_message(e)
            
            # Add error message to database
            assistant_message = await self.conversation_service.add_message(
                conversation_id=conversation_id,
                content=error_message,
                role=MessageRole.ASSISTANT,
                message_metadata={
                    "error": True, 
                    "error_type": error_type,
                    "original_error": str(e),
                    "langchain_error": True
                }
            )
            
            # Re-raise specific exceptions for proper error handling
            if "rate limit" in str(e).lower():
                raise RateLimitError(str(e))
            elif "api key" in str(e).lower() or "authentication" in str(e).lower():
                raise AuthenticationError(str(e))
            elif "openai" in str(e).lower():
                raise OpenAIError(str(e))
            
            return ChatResponse(
                user_message=MessageResponse.from_orm(user_message),
                assistant_message=MessageResponse.from_orm(assistant_message),
                total_tokens=0,
                model_used=self.llm.model_name
            )
    
    async def chat_stream(
        self, 
        conversation_id: int, 
        message: str,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None
    ) -> AsyncGenerator[str, None]:
        """Send a message and get streaming AI response using LangChain."""
        logger.info(f"通过 LangChain 进行流式处理对话 请求，会话 ID: {conversation_id}")
        
        try:
            # Get conversation details
            conversation = await self.conversation_service.get_conversation(conversation_id)
            conv = conversation.to_dict()
            if not conversation:
                raise ChatServiceError("Conversation not found")
            
            # Add user message to database
            user_message = await self.conversation_service.add_message(
                conversation_id=conversation_id,
                content=message,
                role=MessageRole.USER
            )
            
            # Get conversation history for context
            history = await self.conversation_service.get_conversation_history(
                conversation_id, limit=20
            )
            
            # Prepare messages for LangChain
            langchain_messages = self._prepare_langchain_messages(conv, history)
            # Update streaming LLM parameters if provided
            streaming_llm_to_use = self.streaming_llm
            if temperature is not None or max_tokens is not None:
                llm_config = await settings.llm.get_current_config()
                streaming_llm_to_use = ChatOpenAI(
                    model=llm_config["model"],
                    openai_api_key=llm_config["api_key"],
                    openai_api_base=llm_config["base_url"],
                    temperature=temperature if temperature is not None else float(conversation.temperature),
                    max_tokens=max_tokens if max_tokens is not None else conversation.max_tokens,
                    streaming=True
                )
            # Clear previous streaming handler state
            self.streaming_handler.clear()
            
            # Stream response 
            full_response = ""
            try:
                async for chunk in streaming_llm_to_use._astream(langchain_messages):
                    # Handle different chunk types to avoid KeyError
                    chunk_content = None
                    if hasattr(chunk, 'content'):
                        # For object-like chunks with content attribute
                        chunk_content = chunk.content
                    elif isinstance(chunk, dict) and 'content' in chunk:
                        # For dict-like chunks with content key
                        chunk_content = chunk['content']
                    elif isinstance(chunk, dict) and 'error' in chunk:
                        # Handle error chunks explicitly
                        logger.error(f"Error in LLM response: {chunk['error']}")
                        yield self._format_error_message(Exception(chunk['error']))
                        continue
                    
                    if chunk_content:
                        full_response += chunk_content
                        yield chunk_content
            except Exception as e:
                logger.error(f"Error in LLM streaming: {e}")
                yield f"{self._format_error_message(e)} >>> {e}"
            # Add complete assistant message to database
            assistant_message = await self.conversation_service.add_message(
                conversation_id=conversation_id,
                content=full_response,
                role=MessageRole.ASSISTANT,
                message_metadata={
                    "model": streaming_llm_to_use.model_name,
                    "langchain_version": "0.1.0",
                    "provider": "langchain_openai",
                    "streaming": True
                }
            )
            
            # Update conversation timestamp
            await self.conversation_service.update_conversation_timestamp(conversation_id)
            logger.info(f"完成 LangChain 流式处理对话，会话 ID: {conversation_id}")
            
        except Exception as e:
            # 安全地格式化异常信息，避免再次引发KeyError
            error_info = f"Failed to process LangChain streaming chat request for conversation {conversation_id} >>> {e}"
            logger.error(error_info, exc_info=True)
            
            # Format error message for user
            error_message = self._format_error_message(e)
            yield error_message
            
            # Add error message to database
            await self.conversation_service.add_message(
                conversation_id=conversation_id,
                content=error_message,
                role=MessageRole.ASSISTANT,
                message_metadata={
                    "error": True, 
                    "error_type": type(e).__name__,
                    "original_error": str(e),
                    "langchain_error": True,
                    "streaming": True
                }
            )
    
    async def get_available_models(self) -> List[str]:
        """Get list of available models from LangChain."""
        try:
            # LangChain doesn't have a direct method to list models
            # Return commonly available OpenAI models
            return [
                "gpt-3.5-turbo",
                "gpt-3.5-turbo-16k",
                "gpt-4",
                "gpt-4-turbo-preview",
                "gpt-4o",
                "gpt-4o-mini"
            ]
        except Exception as e:
            logger.error(f"Failed to get available models: {str(e)}")
            return ["gpt-3.5-turbo"]
    
    async def update_model_config(
        self, 
        model: Optional[str] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None
    ):
        """Update LLM configuration."""
        from ..core.new_agent import new_agent
        
        # 重新创建LLM实例
        self.llm = await new_agent(
            model=model,
            temperature=temperature,
            streaming=False
        )
        
        self.streaming_llm = await new_agent(
            model=model,
            temperature=temperature,
            streaming=True
        )
        
        logger.info(f"Updated LLM configuration: model={model}, temperature={temperature}, max_tokens={max_tokens}")
    
    def _format_error_message(self, error: Exception) -> str:
        """Format error message for user display."""
        error_type = type(error).__name__
        error_str = str(error)
        
        # Provide user-friendly error messages
        if "rate limit" in error_str.lower():
            return "服务器繁忙，请稍后再试。"
        elif "api key" in error_str.lower() or "authentication" in error_str.lower():
            return f"API认证失败，请检查配置文件。"
        elif "timeout" in error_str.lower():
            return "请求超时，请重试。"
        elif "connection" in error_str.lower():
            return "网络连接错误，请检查网络连接。"
        elif "model" in error_str.lower() and "not found" in error_str.lower():
            return "指定的模型不可用，请选择其他模型。"
        else:
            return f"处理请求时发生错误：{error_str}"
    
    async def _retry_with_backoff(self, func, max_retries: int = 3, base_delay: float = 1.0):
        """Retry function with exponential backoff."""
        for attempt in range(max_retries):
            try:
                return await func()
            except Exception as e:
                if attempt == max_retries - 1:
                    raise e
                
                # Check if error is retryable
                if not self._is_retryable_error(e):
                    raise e
                
                delay = base_delay * (2 ** attempt)
                logger.warning(f"Attempt {attempt + 1} failed, retrying in {delay}s: {str(e)}")
                await asyncio.sleep(delay)
    
    def _is_retryable_error(self, error: Exception) -> bool:
        """Check if an error is retryable."""
        error_str = str(error).lower()
        retryable_errors = [
            "timeout",
            "connection",
            "server error",
            "internal error",
            "rate limit"
        ]
        return any(err in error_str for err in retryable_errors)