from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session
from sqlalchemy import text
from database import get_db
import models
from core.deps import get_current_admin_user
import httpx
import json
import asyncio
from datetime import datetime
from typing import List, Dict, Any

router = APIRouter()

def generate_chart_config(data: List[Dict[str, Any]], chart_type: str, title: str = "数据图表") -> Dict:
    """
    根据数据生成ECharts图表配置
    
    Args:
        data: 查询结果数据
        chart_type: 图表类型 (bar, line, pie, scatter)
        title: 图表标题
    
    Returns:
        ECharts配置字典
    """
    print(f"[DEBUG] generate_chart_config 被调用: chart_type={chart_type}, title={title}, data_len={len(data) if data else 0}")
    
    if not data or len(data) == 0:
        print(f"[ERROR] 数据为空，无法生成图表")
        return None
    
    # 获取数据的键
    keys = list(data[0].keys())
    print(f"[DEBUG] 数据列: {keys}")
    print(f"[DEBUG] 第一行数据: {data[0]}")
    
    # 尝试识别x轴和y轴数据
    # 假设第一列是x轴（标签），其他数值列是y轴（数据）
    x_key = keys[0]
    y_keys = [k for k in keys[1:] if isinstance(data[0].get(k), (int, float))]
    
    if not y_keys:
        # 如果没有数值列，尝试使用第二列
        y_keys = [keys[1]] if len(keys) > 1 else []
    
    print(f"[DEBUG] x_key={x_key}, y_keys={y_keys}")
    
    # 提取数据
    x_data = [str(row[x_key]) for row in data]
    
    # 根据图表类型生成配置
    if chart_type == "pie":
        # 饼图：使用第一个数值列
        y_key = y_keys[0] if y_keys else keys[1]
        pie_data = [{"name": str(row[x_key]), "value": row[y_key]} for row in data]
        
        config = {
            "title": {"text": title, "left": "center"},
            "tooltip": {"trigger": "item", "formatter": "{a} <br/>{b}: {c} ({d}%)"},
            "legend": {"orient": "vertical", "left": "left"},
            "series": [{
                "name": title,
                "type": "pie",
                "radius": "50%",
                "data": pie_data,
                "emphasis": {
                    "itemStyle": {
                        "shadowBlur": 10,
                        "shadowOffsetX": 0,
                        "shadowColor": "rgba(0, 0, 0, 0.5)"
                    }
                }
            }]
        }
    
    elif chart_type == "line":
        # 折线图
        series = []
        for y_key in y_keys:
            y_data = [row[y_key] for row in data]
            series.append({
                "name": y_key,
                "type": "line",
                "data": y_data,
                "smooth": True
            })
        
        config = {
            "title": {"text": title, "left": "center"},
            "tooltip": {"trigger": "axis"},
            "legend": {"data": y_keys, "top": "bottom"},
            "xAxis": {"type": "category", "data": x_data, "boundaryGap": False},
            "yAxis": {"type": "value"},
            "series": series
        }
    
    elif chart_type == "scatter":
        # 散点图
        series = []
        for y_key in y_keys:
            scatter_data = [[i, row[y_key]] for i, row in enumerate(data)]
            series.append({
                "name": y_key,
                "type": "scatter",
                "data": scatter_data
            })
        
        config = {
            "title": {"text": title, "left": "center"},
            "tooltip": {"trigger": "item"},
            "legend": {"data": y_keys, "top": "bottom"},
            "xAxis": {"type": "value"},
            "yAxis": {"type": "value"},
            "series": series
        }
    
    else:  # bar (默认)
        # 柱状图
        series = []
        for y_key in y_keys:
            y_data = [row[y_key] for row in data]
            series.append({
                "name": y_key,
                "type": "bar",
                "data": y_data
            })
        
        config = {
            "title": {"text": title, "left": "center"},
            "tooltip": {"trigger": "axis", "axisPointer": {"type": "shadow"}},
            "legend": {"data": y_keys, "top": "bottom"},
            "xAxis": {"type": "category", "data": x_data},
            "yAxis": {"type": "value"},
            "series": series
        }
    
    print(f"[DEBUG] 图表配置生成完成: type={chart_type}, series_count={len(config.get('series', []))}")
    return config

# 数据库表结构信息
DATABASE_SCHEMA = """
# 数据库表结构

## users 表（用户表）
- id: INTEGER (主键)
- username: STRING (用户名，唯一)
- email: STRING (邮箱，唯一)
- is_active: BOOLEAN (是否激活)
- is_admin: BOOLEAN (是否管理员)

## chat_rooms 表（聊天室表）
- id: INTEGER (主键)
- name: STRING (聊天室名称)
- owner_id: INTEGER (群主ID，外键关联 users.id)
- is_private: BOOLEAN (是否私聊)

## messages 表（消息表）
- id: INTEGER (主键)
- room_id: INTEGER (聊天室ID，外键关联 chat_rooms.id)
- sender_id: INTEGER (发送者ID，外键关联 users.id)
- content: TEXT (消息内容)
- message_type: STRING (消息类型: text, emoji, image)
- timestamp: DATETIME (发送时间)
- is_sensitive: BOOLEAN (是否敏感消息)

## room_members 表（聊天室成员表）
- id: INTEGER (主键)
- room_id: INTEGER (聊天室ID，外键关联 chat_rooms.id)
- user_id: INTEGER (用户ID，外键关联 users.id)

## digital_employees 表（数字员工表）
- id: INTEGER (主键)
- name: STRING (员工名称)
- model_id: INTEGER (模型ID)
- system_prompt: TEXT (系统提示词)

## room_employees 表（聊天室数字员工关联表）
- id: INTEGER (主键)
- room_id: INTEGER (聊天室ID)
- employee_id: INTEGER (数字员工ID)

## 常用查询示例：
1. 查询某个用户发送的所有消息：
   SELECT * FROM messages WHERE sender_id = ?

2. 查询某个聊天室的所有消息：
   SELECT m.*, u.username as sender_name 
   FROM messages m 
   JOIN users u ON m.sender_id = u.id 
   WHERE m.room_id = ?

3. 查询包含特定关键词的消息：
   SELECT * FROM messages WHERE content LIKE '%关键词%'

4. 查询敏感消息：
   SELECT * FROM messages WHERE is_sensitive = 1

5. 统计每个用户的消息数量：
   SELECT sender_id, COUNT(*) as count 
   FROM messages 
   GROUP BY sender_id

注意：
- 只能执行 SELECT 查询，不允许 INSERT, UPDATE, DELETE 等修改操作
- 查询结果最多返回 100 条记录
- 时间字段使用 UTC 时间
"""

def execute_sql_query(db: Session, sql: str) -> dict:
    """
    安全地执行 SQL 查询
    只允许 SELECT 语句
    """
    # 安全检查：只允许 SELECT 语句
    sql_upper = sql.strip().upper()
    if not sql_upper.startswith('SELECT'):
        return {
            "success": False,
            "error": "只允许执行 SELECT 查询语句"
        }
    
    # 禁止的关键词
    forbidden_keywords = ['DROP', 'DELETE', 'UPDATE', 'INSERT', 'ALTER', 'CREATE', 'TRUNCATE']
    for keyword in forbidden_keywords:
        if keyword in sql_upper:
            return {
                "success": False,
                "error": f"查询中不允许包含 {keyword} 关键词"
            }
    
    try:
        # 执行查询 - 使用 text() 包装 SQL 字符串
        result = db.execute(text(sql))
        rows = result.fetchall()
        
        # 限制返回数量
        if len(rows) > 100:
            rows = rows[:100]
            truncated = True
        else:
            truncated = False
        
        # 转换为字典列表
        columns = result.keys()
        data = []
        for row in rows:
            row_dict = {}
            for i, col in enumerate(columns):
                value = row[i]
                # 处理日期时间类型
                if isinstance(value, datetime):
                    value = value.isoformat()
                row_dict[col] = value
            data.append(row_dict)
        
        return {
            "success": True,
            "data": data,
            "count": len(data),
            "truncated": truncated,
            "sql": sql
        }
    except Exception as e:
        return {
            "success": False,
            "error": f"SQL 执行错误: {str(e)}",
            "sql": sql
        }

@router.post("/chat")
async def chat_with_assistant(
    request: dict,
    db: Session = Depends(get_db),
    admin_user: models.User = Depends(get_current_admin_user)
):
    """
    对话助手 - 支持流式输出
    """
    user_message = request.get("message")
    model_id = request.get("model_id")
    
    if not user_message:
        raise HTTPException(status_code=400, detail="消息不能为空")
    
    # 获取 AI 模型
    if model_id:
        ai_model = db.query(models.AIModel).filter(models.AIModel.id == model_id).first()
    else:
        # 使用默认模型
        ai_model = db.query(models.AIModel).filter(models.AIModel.is_default == True).first()
        if not ai_model:
            ai_model = db.query(models.AIModel).first()
    
    if not ai_model:
        raise HTTPException(status_code=404, detail="没有可用的AI模型")
    
    # 定义 SQL 查询工具
    sql_tool = {
        "type": "function",
        "function": {
            "name": "execute_sql_query",
            "description": "执行 SQL 查询来获取聊天消息数据。只能执行 SELECT 查询，不允许修改数据。",
            "parameters": {
                "type": "object",
                "properties": {
                    "sql": {
                        "type": "string",
                        "description": "要执行的 SQL 查询语句，必须是 SELECT 语句"
                    },
                    "explanation": {
                        "type": "string",
                        "description": "对这个查询的简短说明"
                    },
                    "need_chart": {
                        "type": "boolean",
                        "description": "是否需要生成图表。统计、对比、趋势类查询应该生成图表，简单查询不需要"
                    }
                },
                "required": ["sql", "explanation", "need_chart"]
            }
        }
    }
    
    # 定义图表参数生成工具
    chart_params_tool = {
        "type": "function",
        "function": {
            "name": "generate_chart_params",
            "description": "根据SQL查询结果，生成图表所需的参数（类型、标题、说明）",
            "parameters": {
                "type": "object",
                "properties": {
                    "chart_type": {
                        "type": "string",
                        "enum": ["bar", "line", "pie", "scatter"],
                        "description": "图表类型：bar(柱状图-对比), line(折线图-趋势), pie(饼图-占比), scatter(散点图-分布)"
                    },
                    "title": {
                        "type": "string",
                        "description": "图表标题，简洁明了"
                    },
                    "explanation": {
                        "type": "string",
                        "description": "为什么选择这种图表类型，如何解读数据"
                    }
                },
                "required": ["chart_type", "title", "explanation"]
            }
        }
    }
    
    # 系统提示词
    system_prompt = f"""你是一个智能数据分析助手，专门帮助管理员查询和分析聊天消息数据。

你可以使用以下工具：
1. execute_sql_query - 执行SQL查询获取数据（需要指定是否需要图表）
2. generate_chart_params - 根据查询结果生成图表参数

{DATABASE_SCHEMA}

**工作流程**：
1. 理解用户查询意图
2. 调用 execute_sql_query：
   - 生成SQL语句
   - 判断是否需要图表（need_chart: true/false）
   - 统计、对比、趋势类查询 → need_chart=true
   - 简单查询、单条记录 → need_chart=false
3. 如果 need_chart=true，调用 generate_chart_params：
   - 分析查询结果数据
   - 选择最合适的图表类型
   - 生成图表标题和说明
4. 用清晰、友好的语言总结结果

**图表类型选择指南**：
- **柱状图(bar)**: 对比不同类别 → "各用户消息数量"、"不同聊天室活跃度"
- **折线图(line)**: 时间趋势 → "每日消息量变化"、"每小时活跃度"
- **饼图(pie)**: 占比分布 → "消息类型占比"、"用户活跃度分布"
- **散点图(scatter)**: 数据分布 → "消息长度分布"、"发送时间分布"

**示例1 - 需要图表**：
用户："统计每个用户的消息数量"
1. execute_sql_query(sql="SELECT sender_id, COUNT(*) as count FROM messages GROUP BY sender_id", need_chart=true)
2. generate_chart_params(chart_type="bar", title="用户消息数量统计")
3. 总结数据

**示例2 - 不需要图表**：
用户："查询ID为5的消息内容"
1. execute_sql_query(sql="SELECT * FROM messages WHERE id=5", need_chart=false)
2. 总结结果（无需图表）

注意事项：
- 只能使用 SELECT 语句
- 查询结果最多返回 100 条
- 时间使用 UTC 时区
- 总结时要突出关键信息和数据洞察

示例对话：
用户："统计每个用户的消息数量"
助手：我来帮你统计每个用户的消息数量...
[调用 execute_sql_query 查询]
[调用 generate_chart 生成柱状图]
助手：查询完成！我为您生成了用户消息数量统计图表。从图表可以看出...
"""
    
    async def generate_stream():
        try:
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_message}
            ]
            
            headers = {
                "Authorization": f"Bearer {ai_model.api_key}",
                "Content-Type": "application/json"
            }
            
            payload = {
                "model": ai_model.name,
                "messages": messages,
                "tools": [sql_tool, chart_params_tool],
                "tool_choice": "auto",
                "stream": False
            }
            
            # 发送第一次请求
            async with httpx.AsyncClient(timeout=60.0) as client:
                response = await client.post(
                    f"{ai_model.base_url}/chat/completions",
                    headers=headers,
                    json=payload
                )
                
                if response.status_code != 200:
                    yield f"data: {json.dumps({'type': 'error', 'content': f'API 错误: {response.status_code}'})}\n\n"
                    return
                
                result = response.json()
                
                if "choices" not in result or len(result["choices"]) == 0:
                    yield f"data: {json.dumps({'type': 'error', 'content': 'AI 响应格式错误'})}\n\n"
                    return
                
                assistant_message = result["choices"][0]["message"]
                
                # 检查是否有工具调用
                if "tool_calls" in assistant_message and assistant_message["tool_calls"]:
                    # 用于存储查询结果和是否需要图表
                    query_result = None
                    need_chart = False
                    
                    # 先将assistant_message添加到消息历史
                    messages.append(assistant_message)
                    
                    # 第一阶段：处理SQL查询
                    for tool_call in assistant_message["tool_calls"]:
                        function_name = tool_call["function"]["name"]
                        function_args_str = tool_call["function"]["arguments"]
                        
                        try:
                            function_args = json.loads(function_args_str)
                        except json.JSONDecodeError:
                            function_args = {}
                        
                        if function_name == "execute_sql_query":
                            sql = function_args.get("sql", "")
                            explanation = function_args.get("explanation", "")
                            need_chart = function_args.get("need_chart", False)
                            
                            print(f"[DEBUG] SQL查询: {sql}")
                            print(f"[DEBUG] 是否需要图表: {need_chart}")
                            
                            # 发送工具调用通知
                            yield f"data: {json.dumps({'type': 'tool_call', 'tool': 'SQL查询', 'explanation': explanation, 'sql': sql, 'need_chart': need_chart}, ensure_ascii=False)}\n\n"
                            
                            # 执行 SQL 查询
                            query_result = execute_sql_query(db, sql)
                            
                            # 发送查询结果通知
                            if query_result["success"]:
                                yield f"data: {json.dumps({'type': 'tool_result', 'success': True, 'count': query_result['count']}, ensure_ascii=False)}\n\n"
                                print(f"[DEBUG] SQL查询成功，返回 {query_result['count']} 条记录")
                            else:
                                yield f"data: {json.dumps({'type': 'tool_result', 'success': False, 'error': query_result['error']}, ensure_ascii=False)}\n\n"
                                print(f"[ERROR] SQL查询失败: {query_result['error']}")
                            
                            # 将工具结果添加到消息历史
                            messages.append({
                                "role": "tool",
                                "tool_call_id": tool_call["id"],
                                "name": function_name,
                                "content": json.dumps(query_result, ensure_ascii=False)
                            })
                            
                            break  # 只处理第一个SQL查询
                    
                    # 第二阶段：如果需要图表且查询成功，让AI生成图表参数
                    if need_chart and query_result and query_result.get("success") and query_result.get("data"):
                        print(f"[DEBUG] 需要生成图表，准备第二次调用AI")
                        
                        # 构建消息：让AI根据查询结果生成图表参数
                        chart_request_messages = [
                            {"role": "system", "content": "你是图表参数生成专家。根据SQL查询结果，选择最合适的图表类型和标题。"},
                            {"role": "user", "content": f"SQL查询结果：\n{json.dumps(query_result['data'][:5], ensure_ascii=False, indent=2)}\n\n请调用 generate_chart_params 工具生成图表参数。"}
                        ]
                        
                        # 第二次调用AI生成图表参数
                        chart_payload = {
                            "model": ai_model.name,
                            "messages": chart_request_messages,
                            "tools": [chart_params_tool],
                            "tool_choice": {"type": "function", "function": {"name": "generate_chart_params"}},
                            "stream": False
                        }
                        
                        chart_response = await client.post(
                            f"{ai_model.base_url}/chat/completions",
                            headers=headers,
                            json=chart_payload
                        )
                        
                        if chart_response.status_code == 200:
                            chart_result = chart_response.json()
                            if "choices" in chart_result and len(chart_result["choices"]) > 0:
                                chart_message = chart_result["choices"][0]["message"]
                                
                                if "tool_calls" in chart_message and chart_message["tool_calls"]:
                                    chart_tool_call = chart_message["tool_calls"][0]
                                    chart_args_str = chart_tool_call["function"]["arguments"]
                                    
                                    try:
                                        chart_args = json.loads(chart_args_str)
                                        chart_type = chart_args.get("chart_type", "bar")
                                        chart_title = chart_args.get("title", "数据图表")
                                        chart_explanation = chart_args.get("explanation", "")
                                        
                                        print(f"[DEBUG] AI生成图表参数: type={chart_type}, title={chart_title}")
                                        
                                        # 发送图表生成通知
                                        yield f"data: {json.dumps({'type': 'chart_call', 'chart_type': chart_type, 'title': chart_title, 'explanation': chart_explanation}, ensure_ascii=False)}\n\n"
                                        
                                        # 生成图表配置
                                        chart_config = generate_chart_config(
                                            query_result["data"],
                                            chart_type,
                                            chart_title
                                        )
                                        
                                        if chart_config:
                                            print(f"[DEBUG] 图表配置生成成功")
                                            yield f"data: {json.dumps({'type': 'chart_result', 'success': True, 'config': chart_config}, ensure_ascii=False)}\n\n"
                                        else:
                                            print(f"[ERROR] 图表配置生成失败")
                                            yield f"data: {json.dumps({'type': 'chart_result', 'success': False, 'error': '无法生成图表配置'}, ensure_ascii=False)}\n\n"
                                    
                                    except json.JSONDecodeError as e:
                                        print(f"[ERROR] 解析图表参数失败: {e}")
                                        yield f"data: {json.dumps({'type': 'chart_result', 'success': False, 'error': '图表参数解析失败'}, ensure_ascii=False)}\n\n"
                        else:
                            print(f"[ERROR] 图表参数生成请求失败: {chart_response.status_code}")
                    
                    print(f"[DEBUG] 工具调用完成，准备生成总结")
                    
                    # 打印消息历史用于调试
                    print(f"[DEBUG] 消息历史:")
                    for i, msg in enumerate(messages):
                        print(f"  [{i}] role={msg.get('role')}, has_content={bool(msg.get('content'))}, has_tool_calls={bool(msg.get('tool_calls'))}")
                    
                    # 重新构建消息历史，只保留system和user消息
                    # 只包含SQL查询结果，不包含图表信息（图表已经在前端显示）
                    converted_messages = [messages[0]]  # 保留system消息
                    
                    # 添加原始用户消息
                    converted_messages.append(messages[1])
                    
                    # 只收集SQL查询结果，忽略图表工具结果
                    sql_results = []
                    for msg in messages[2:]:  # 跳过system和第一个user消息
                        if msg.get("role") == "tool":
                            tool_content = msg.get("content", "")
                            tool_name = msg.get("name", "")
                            
                            # 只处理SQL查询结果
                            if tool_name == "execute_sql_query":
                                try:
                                    result_data = json.loads(tool_content)
                                    if result_data.get("success"):
                                        data_preview = result_data['data'][:5] if len(result_data['data']) > 5 else result_data['data']
                                        sql_results.append(f"SQL查询成功，返回{result_data.get('count', 0)}条记录。\n数据示例：\n{json.dumps(data_preview, ensure_ascii=False, indent=2)}")
                                    else:
                                        sql_results.append(f"SQL查询失败：{result_data.get('error', '未知错误')}")
                                except:
                                    sql_results.append(tool_content)
                    
                    # 将SQL查询结果合并为一个user消息
                    if sql_results:
                        converted_messages.append({
                            "role": "user",
                            "content": "以下是查询结果，请用简洁友好的语言总结关键信息和数据洞察（图表已自动生成，无需提及）：\n\n" + "\n\n".join(sql_results)
                        })
                        print(f"[DEBUG] 合并了{len(sql_results)}个SQL查询结果")
                    
                    print(f"[DEBUG] 转换后的消息数: {len(converted_messages)}")
                    for i, msg in enumerate(converted_messages):
                        print(f"  [{i}] role={msg.get('role')}, content_length={len(str(msg.get('content', '')))}")
                    
                    # 第三次调用：生成总结（流式响应）
                    print(f"[DEBUG] 准备生成总结，使用流式响应")
                    payload["messages"] = converted_messages
                    payload["stream"] = True
                    # 移除 tools 参数
                    if "tools" in payload:
                        del payload["tools"]
                    if "tool_choice" in payload:
                        del payload["tool_choice"]
                    
                    print(f"[DEBUG] 发送流式请求到模型...")
                    
                    # 使用真正的流式请求
                    async with client.stream(
                        "POST",
                        f"{ai_model.base_url}/chat/completions",
                        headers=headers,
                        json=payload,
                        timeout=60.0
                    ) as stream_response:
                        if stream_response.status_code != 200:
                            print(f"[ERROR] 流式请求失败: {stream_response.status_code}")
                            error_text = await stream_response.aread()
                            print(f"[ERROR] 错误详情: {error_text.decode('utf-8')[:500]}")
                            yield f"data: {json.dumps({'type': 'error', 'content': f'AI响应错误: {stream_response.status_code}'}, ensure_ascii=False)}\n\n"
                        else:
                            print(f"[DEBUG] 流式响应开始")
                            buffer = ""
                            async for chunk in stream_response.aiter_bytes():
                                buffer += chunk.decode('utf-8')
                                lines = buffer.split('\n')
                                buffer = lines[-1]  # 保留最后一行（可能不完整）
                                
                                for line in lines[:-1]:
                                    line = line.strip()
                                    if line.startswith('data: '):
                                        data_str = line[6:]  # 移除 'data: ' 前缀
                                        if data_str == '[DONE]':
                                            print(f"[DEBUG] 流式响应完成")
                                            continue
                                        
                                        try:
                                            data = json.loads(data_str)
                                            if "choices" in data and len(data["choices"]) > 0:
                                                delta = data["choices"][0].get("delta", {})
                                                content = delta.get("content", "")
                                                if content:
                                                    # 转发内容到前端
                                                    yield f"data: {json.dumps({'type': 'content', 'content': content}, ensure_ascii=False)}\n\n"
                                        except json.JSONDecodeError:
                                            pass  # 忽略解析错误
                            
                            print(f"[DEBUG] 流式内容发送完成")
                
                else:
                    # 没有工具调用，直接返回内容（流式）
                    content = assistant_message.get("content", "")
                    # 逐字输出
                    for char in content:
                        yield f"data: {json.dumps({'type': 'content', 'content': char}, ensure_ascii=False)}\n\n"
                        await asyncio.sleep(0.01)
                
                # 发送完成信号
                yield f"data: {json.dumps({'type': 'done'})}\n\n"
                
        except Exception as e:
            yield f"data: {json.dumps({'type': 'error', 'content': f'错误: {str(e)}'})}\n\n"
    
    return StreamingResponse(
        generate_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
        }
    )

@router.get("/test-sql")
async def test_sql_query(
    sql: str,
    db: Session = Depends(get_db),
    admin_user: models.User = Depends(get_current_admin_user)
):
    """
    测试 SQL 查询（用于调试）
    """
    result = execute_sql_query(db, sql)
    return result
