"""
对话服务
"""
import logging
import time
from typing import AsyncGenerator, Dict, List, Optional

from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import AsyncSession

from ai_platform.models import AIApp, Conversation, Message, LLMModel
from utils.context import get_current_user_id_from_context
from .llm_service import LLMService

logger = logging.getLogger(__name__)


class ChatService:
    """
    对话服务
    
    管理对话和消息
    """
    
    def __init__(self, db: AsyncSession):
        self._db = db
        self.llm_service = LLMService(db)
    
    async def create_conversation(
        self,
        app_id: str,
        title: str = '',
        user_id: Optional[str] = None,
    ) -> Conversation:
        """
        创建对话
        
        Args:
            app_id: 应用 ID
            title: 对话标题
            
        Returns:
            Conversation
        """
        current_user_id = user_id or get_current_user_id_from_context()
        if not current_user_id:
            raise ValueError('未登录或登录已过期')

        # 查询应用
        result = await self._db.execute(
            select(AIApp).where(AIApp.id == app_id, AIApp.is_deleted == False)
        )
        app = result.scalar_one_or_none()
        if not app:
            raise ValueError(f'应用不存在: {app_id}')
        
        conversation = Conversation(
            app_id=app_id,
            user_id=current_user_id,
            title=title or '新对话',
        )
        self._db.add(conversation)
        
        # 更新应用统计
        app.conversation_count = (app.conversation_count or 0) + 1
        
        await self._db.commit()
        await self._db.refresh(conversation)
        
        return conversation
    
    async def get_conversation(self, conversation_id: str) -> Optional[Conversation]:
        """获取对话"""
        result = await self._db.execute(
            select(Conversation).where(
                Conversation.id == conversation_id,
                Conversation.is_deleted == False
            )
        )
        return result.scalar_one_or_none()
    
    async def list_conversations(
        self,
        app_id: str,
        page: int = 1,
        page_size: int = 20,
    ) -> tuple:
        """
        获取对话列表
        
        Returns:
            (conversations, total)
        """
        query = select(Conversation).where(
            Conversation.app_id == app_id,
            Conversation.is_deleted == False
        ).order_by(Conversation.is_pinned.desc(), Conversation.sys_update_datetime.desc())
        
        # 获取总数
        count_result = await self._db.execute(
            select(func.count()).select_from(query.subquery())
        )
        total = count_result.scalar() or 0
        
        # 分页
        offset = (page - 1) * page_size
        query = query.offset(offset).limit(page_size)
        
        result = await self._db.execute(query)
        conversations = result.scalars().all()
        
        return list(conversations), total
    
    async def delete_conversation(self, conversation_id: str) -> bool:
        """删除对话"""
        result = await self._db.execute(
            select(Conversation).where(
                Conversation.id == conversation_id,
                Conversation.is_deleted == False
            )
        )
        conversation = result.scalar_one_or_none()
        if not conversation:
            return False
        
        conversation.is_deleted = True
        await self._db.commit()
        return True
    
    async def get_messages(
        self,
        conversation_id: str,
        limit: int = 50,
    ) -> List[Message]:
        """获取对话消息"""
        result = await self._db.execute(
            select(Message).where(
                Message.conversation_id == conversation_id,
                Message.is_deleted == False
            ).order_by(Message.sys_create_datetime).limit(limit)
        )
        return list(result.scalars().all())
    
    async def send_message(
        self,
        conversation_id: str,
        content: str,
    ) -> tuple:
        """
        发送消息并获取 AI 回复
        
        Args:
            conversation_id: 对话 ID
            content: 消息内容
            
        Returns:
            (user_message, assistant_message)
        """
        conversation = await self.get_conversation(conversation_id)
        if not conversation:
            raise ValueError('对话不存在')
        
        # 获取应用
        app_result = await self._db.execute(
            select(AIApp).where(AIApp.id == conversation.app_id)
        )
        app = app_result.scalar_one_or_none()
        if not app:
            raise ValueError('应用不存在')
        
        # 获取模型
        model = await self._get_effective_model(conversation, app)
        if not model:
            raise ValueError('未配置模型')
        
        # 创建用户消息
        user_message = Message(
            conversation_id=conversation_id,
            role='user',
            content=content,
            status='completed',
        )
        self._db.add(user_message)
        await self._db.flush()
        
        # 构建消息列表
        messages = await self._build_messages(conversation, app, content)
        
        # 创建助手消息（pending 状态）
        assistant_message = Message(
            conversation_id=conversation_id,
            role='assistant',
            content='',
            status='pending',
            model_name=model.model_name,
        )
        self._db.add(assistant_message)
        await self._db.flush()
        
        try:
            start_time = time.time()
            
            # 调用 LLM
            response = await self.llm_service.chat_async(
                model_id=str(model.id),
                messages=messages,
                temperature=app.temperature or 0.7,
                max_tokens=app.max_tokens or 2048,
            )
            
            latency = int((time.time() - start_time) * 1000)
            
            # 更新助手消息
            assistant_message.content = response.content
            assistant_message.status = 'completed'
            assistant_message.prompt_tokens = response.prompt_tokens
            assistant_message.completion_tokens = response.completion_tokens
            assistant_message.total_tokens = response.total_tokens
            assistant_message.latency = latency
            
            # 更新对话统计
            conversation.message_count = (conversation.message_count or 0) + 2
            conversation.total_tokens = (conversation.total_tokens or 0) + response.total_tokens
            
            # 更新应用统计
            app.message_count = (app.message_count or 0) + 2
            
            # 自动生成标题
            if conversation.message_count == 2:
                self._generate_title(conversation, content)
            
            await self._db.commit()
            
            return user_message, assistant_message
            
        except Exception as e:
            logger.exception(f'发送消息失败: {e}')
            assistant_message.status = 'failed'
            assistant_message.error_message = str(e)
            await self._db.commit()
            raise
    
    async def send_message_stream(
        self,
        conversation_id: str,
        content: str,
    ) -> AsyncGenerator[Dict, None]:
        """
        流式发送消息
        
        Yields:
            {"type": "content", "content": "..."} 或
            {"type": "done", "message": {...}}
        """
        conversation = await self.get_conversation(conversation_id)
        if not conversation:
            raise ValueError('对话不存在')
        
        # 获取应用
        app_result = await self._db.execute(
            select(AIApp).where(AIApp.id == conversation.app_id)
        )
        app = app_result.scalar_one_or_none()
        if not app:
            raise ValueError('应用不存在')
        
        # 获取模型
        model = await self._get_effective_model(conversation, app)
        if not model:
            raise ValueError('未配置模型')
        
        # 创建用户消息
        user_message = Message(
            conversation_id=conversation_id,
            role='user',
            content=content,
            status='completed',
        )
        self._db.add(user_message)
        await self._db.flush()
        
        # 构建消息列表
        messages = await self._build_messages(conversation, app, content)
        
        # 创建助手消息
        assistant_message = Message(
            conversation_id=conversation_id,
            role='assistant',
            content='',
            status='pending',
            model_name=model.model_name,
        )
        self._db.add(assistant_message)
        await self._db.flush()
        
        try:
            start_time = time.time()
            full_content = ''
            total_tokens = 0
            prompt_tokens = 0
            completion_tokens = 0
            
            # 使用异步流式方法
            async for chunk in self.llm_service.chat_stream(
                model_id=str(model.id),
                messages=messages,
                temperature=app.temperature or 0.7,
                max_tokens=app.max_tokens or 2048,
            ):
                if chunk.content:
                    full_content += chunk.content
                    yield {'type': 'content', 'content': chunk.content}
                
                if chunk.is_finished:
                    prompt_tokens = chunk.prompt_tokens
                    completion_tokens = chunk.completion_tokens
                    total_tokens = chunk.total_tokens
            
            latency = int((time.time() - start_time) * 1000)
            
            # 更新助手消息
            assistant_message.content = full_content
            assistant_message.status = 'completed'
            assistant_message.prompt_tokens = prompt_tokens
            assistant_message.completion_tokens = completion_tokens
            assistant_message.total_tokens = total_tokens
            assistant_message.latency = latency
            
            # 更新对话统计
            conversation.message_count = (conversation.message_count or 0) + 2
            conversation.total_tokens = (conversation.total_tokens or 0) + total_tokens
            
            await self._db.commit()
            
            yield {
                'type': 'done',
                'message': {
                    'id': str(assistant_message.id),
                    'content': full_content,
                    'tokens': total_tokens,
                    'latency': latency,
                },
            }
            
        except Exception as e:
            logger.exception(f'流式发送消息失败: {e}')
            assistant_message.status = 'failed'
            assistant_message.error_message = str(e)
            await self._db.commit()
            yield {'type': 'error', 'error': str(e)}
    
    async def _get_effective_model(self, conversation: Conversation, app: AIApp) -> Optional[LLMModel]:
        """获取有效的模型"""
        model_id = conversation.model_override_id or app.model_id
        if not model_id:
            return None
        
        result = await self._db.execute(
            select(LLMModel).where(
                LLMModel.id == model_id,
                LLMModel.is_active == True,
                LLMModel.is_deleted == False
            )
        )
        return result.scalar_one_or_none()
    
    async def _build_messages(
        self,
        conversation: Conversation,
        app: AIApp,
        user_content: str,
    ) -> List[Dict[str, str]]:
        """构建消息列表"""
        messages = []
        
        # 系统提示词
        if app.system_prompt:
            messages.append({
                'role': 'system',
                'content': app.system_prompt,
            })
        
        # 历史消息
        result = await self._db.execute(
            select(Message).where(
                Message.conversation_id == conversation.id,
                Message.is_deleted == False,
                Message.status == 'completed'
            ).order_by(Message.sys_create_datetime.desc()).limit(20)
        )
        history = result.scalars().all()
        
        for msg in reversed(list(history)):
            messages.append({
                'role': msg.role,
                'content': msg.content,
            })
        
        # 当前用户消息
        messages.append({
            'role': 'user',
            'content': user_content,
        })
        
        return messages
    
    def _generate_title(self, conversation: Conversation, first_message: str):
        """自动生成对话标题"""
        # 简单截取前 20 个字符作为标题
        title = first_message[:20]
        if len(first_message) > 20:
            title += '...'
        conversation.title = title
