# from gevent import monkey
# monkey.patch_all()

import os
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from typing import Dict, List, Any, Optional
import asyncio
import aiosqlite
import sqlite3
from contextlib import closing

from dotenv import load_dotenv
import json
import time
import datetime
import logging
# 加载环境变量
load_dotenv()

def main(query: str) -> dict:
    query = query.strip()
    escaped_query = json.dumps(query, ensure_ascii=False)
    return {
        "format_query": escaped_query,
    }


import sys
sys.path.append(os.getcwd())
from rag2_0.dify.DifyQueryRetrieval import DifyQueryRetrieval

# 定义数据库路径
DATA_DIR = os.path.join(os.getcwd(), "data")
DB_DIR = os.path.join(DATA_DIR, "db")
DB_FILE = os.path.join(DB_DIR, "answer_logs.db")

# 创建异步日志队列和后台任务
log_queue: asyncio.Queue = asyncio.Queue()
worker_task: asyncio.Task | None = None

# 初始化数据库（异步）
async def init_database():
    os.makedirs(DB_DIR, exist_ok=True)
    async with aiosqlite.connect(DB_FILE) as conn:
        await conn.execute('''
            CREATE TABLE IF NOT EXISTS query_types (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                query_type TEXT NOT NULL,
                workflow_run_id TEXT NOT NULL,
                timestamp TEXT NOT NULL
            )
        ''')
        await conn.execute('''
            CREATE TABLE IF NOT EXISTS dislike_reasons (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                dislike_reason TEXT NOT NULL,
                workflow_run_id TEXT NOT NULL,
                timestamp TEXT NOT NULL
            )
        ''')
        await conn.commit()
    logger.info("数据库初始化完成")

# 后台异步工作协程
async def log_worker():
    try:
        while True:
            data = await log_queue.get()
            if data is None:
                # 排空剩余数据后退出
                while not log_queue.empty():
                    pending = log_queue.get_nowait()
                    if pending is None:
                        continue
                    await process_log_data(pending)
                break
            await process_log_data(data)
            log_queue.task_done()
    except asyncio.CancelledError:
        logger.info("日志工作任务被取消，尝试优雅退出...")
    except Exception as e:
        logger.error(f"保存查询数据时出错: {str(e)}", exc_info=True)

# 提取数据处理逻辑到单独异步函数
async def process_log_data(data):
    try:
        async with aiosqlite.connect(DB_FILE) as conn:
            if "dislike_reason" in data:
                await conn.execute(
                    "INSERT INTO dislike_reasons (dislike_reason, workflow_run_id, timestamp) VALUES (?, ?, ?)",
                    (data["dislike_reason"], data["workflow_run_id"], data["timestamp"])
                )
                table_name = "dislike_reasons"
            else:
                await conn.execute(
                    "INSERT INTO query_types (query_type, workflow_run_id, timestamp) VALUES (?, ?, ?)",
                    (data["query_type"], data["workflow_run_id"], data["timestamp"])
                )
                table_name = "query_types"
            await conn.commit()
        logger.info(f"成功保存数据到表: {table_name}")
    except Exception as e:
        logger.error(f"处理日志数据时出错: {str(e)}", exc_info=True)

# 创建数据目录
os.makedirs(DATA_DIR, exist_ok=True)

# 配置日志 - 同时输出到控制台和文件
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# 创建日志目录
LOG_DIR = os.path.join(DATA_DIR, "logs")
os.makedirs(LOG_DIR, exist_ok=True)

# 创建控制台处理器
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)

# 创建文件处理器
file_handler = logging.FileHandler(
    os.path.join(LOG_DIR, "answer_type_service.log"),
    encoding='utf-8'
)
file_handler.setLevel(logging.INFO)

# 创建日志格式
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
console_handler.setFormatter(formatter)
file_handler.setFormatter(formatter)

# 添加处理器到日志器
logger.addHandler(console_handler)
logger.addHandler(file_handler)

# 设置其他库的日志级别
logging.getLogger('httpx').setLevel(logging.WARNING)
logging.getLogger('openai').setLevel(logging.WARNING)

# 定义请求模型
class AnswerTypeRequest(BaseModel):
    query: str
    query_type: str

# 创建FastAPI应用
app = FastAPI(
    title="提问数据类型",
    description="收集用户提问数据类型",
    version="1.0"
)

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 应用启动事件
@app.on_event("startup")
async def startup_event():
    global worker_task
    # 初始化数据库
    await init_database()
    # 启动后台异步任务
    worker_task = asyncio.create_task(log_worker())
    logger.info("后台日志工作任务已启动")

# 应用关闭事件
@app.on_event("shutdown")
async def shutdown_event():
    global worker_task
    if worker_task:
        # 发送退出信号并等待任务结束
        await log_queue.put(None)
        await asyncio.sleep(0)  # 让出执行权给worker处理退出
        try:
            await asyncio.wait_for(worker_task, timeout=10.0)
            logger.info("后台日志工作任务已停止")
        except asyncio.TimeoutError:
            worker_task.cancel()
            logger.warning("工作任务未在超时时间内退出，已取消")

# 添加健康检查端点
@app.get("/health", summary="健康检查")
async def health_check():
    return {"status": "ok"}

@app.get("/query_type", summary="异步检索API")
async def query_type(query_type: str, workflow_run_id:str):
    try:
        # 记录请求
        logger.info(f"接收到请求: 类型: {query_type}, workflow_run_id: {workflow_run_id}")
        
        # 准备数据
        timestamp = datetime.datetime.now().isoformat()
        query_data = {
            "query_type": query_type,
            "timestamp": timestamp,
            "workflow_run_id": workflow_run_id
        }
        
        # 将数据放入异步队列
        try:
            await log_queue.put(query_data)
            success = True
            logger.info(f"查询数据已加入队列，当前队列大小: {log_queue.qsize()}")
        except Exception as e:
            success = False
            logger.error(f"加入队列时出错: {str(e)}", exc_info=True)
        
        # 返回响应
        content = f"<strong>问题类型</strong>: {query_type}<br><strong>操作是否成功</strong>: {'成功' if success else '失败'}"
        return HTMLResponse(content=content)
    except Exception as e:
        logger.error(f"处理请求时出错: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"处理请求时出错: {str(e)}")

@app.get("/dislike_reason", summary="记录点踩原因")
async def dislike_reason(reason: str, workflow_run_id: str):
    try:
        # 记录请求
        logger.info(f"接收到点踩原因: {reason}, workflow_run_id: {workflow_run_id}")
        
        # 准备数据
        timestamp = datetime.datetime.now().isoformat()
        dislike_data = {
            "dislike_reason": reason,
            "timestamp": timestamp,
            "workflow_run_id": workflow_run_id
        }
        
        # 将数据放入异步队列
        try:
            await log_queue.put(dislike_data)
            success = True
            logger.info(f"点踩原因数据已加入队列，当前队列大小: {log_queue.qsize()}")
        except Exception as e:
            success = False
            logger.error(f"加入队列时出错: {str(e)}", exc_info=True)
        
        # 返回响应
        content = f"<strong>点踩原因</strong>: {reason}<br><strong>操作是否成功</strong>: {'成功' if success else '失败'}"
        return HTMLResponse(content=content)
    except Exception as e:
        logger.error(f"处理点踩原因请求时出错: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"处理点踩原因请求时出错: {str(e)}")

# 添加数据查询API
@app.get("/stats", summary="查询统计数据")
async def get_stats():
    try:
        async with aiosqlite.connect(DB_FILE) as conn:
            conn.row_factory = sqlite3.Row
            # 查询类型统计
            cursor = await conn.execute("SELECT COUNT(*) as count FROM query_types")
            row = await cursor.fetchone()
            query_count = row['count'] if row else 0
            # 点踩原因统计
            cursor = await conn.execute("SELECT COUNT(*) as count FROM dislike_reasons")
            row = await cursor.fetchone()
            dislike_count = row['count'] if row else 0
            # 最近5条查询记录
            cursor = await conn.execute(
                """
                SELECT query_type, workflow_run_id, timestamp 
                FROM query_types 
                ORDER BY id DESC LIMIT 5
                """
            )
            recent_queries = [dict(r) for r in await cursor.fetchall()]
            # 最近5条点踩记录
            cursor = await conn.execute(
                """
                SELECT dislike_reason, workflow_run_id, timestamp 
                FROM dislike_reasons 
                ORDER BY id DESC LIMIT 5
                """
            )
            recent_dislikes = [dict(r) for r in await cursor.fetchall()]
        return {
            "statistics": {
                "total_queries": query_count,
                "total_dislikes": dislike_count
            },
            "recent_data": {
                "queries": recent_queries,
                "dislikes": recent_dislikes
            }
        }
    except Exception as e:
        logger.error(f"获取统计数据时出错: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"获取统计数据时出错: {str(e)}")

@app.get("/query_by_workflow_id", summary="根据工作流ID获取查询类型数据")
async def get_query_by_workflow_id(workflow_run_id: str):
    try:
        async with aiosqlite.connect(DB_FILE) as conn:
            conn.row_factory = sqlite3.Row
            cursor = await conn.execute(
                """
                SELECT id, query_type, workflow_run_id, timestamp 
                FROM query_types 
                WHERE workflow_run_id = ?
                ORDER BY id DESC
                """,
                (workflow_run_id,)
            )
            results = [dict(r) for r in await cursor.fetchall()]
        if not results:
            return {"message": f"未找到工作流ID为 {workflow_run_id} 的查询类型数据", "data": []}
        return {"data": results}
    except Exception as e:
        logger.error(f"根据工作流ID获取查询类型数据时出错: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"根据工作流ID获取查询类型数据时出错: {str(e)}")

@app.get("/dislike_by_workflow_id", summary="根据工作流ID获取点踩原因数据")
async def get_dislike_by_workflow_id(workflow_run_id: str):
    try:
        async with aiosqlite.connect(DB_FILE) as conn:
            conn.row_factory = sqlite3.Row
            cursor = await conn.execute(
                """
                SELECT id, dislike_reason, workflow_run_id, timestamp 
                FROM dislike_reasons 
                WHERE workflow_run_id = ?
                ORDER BY id DESC
                """,
                (workflow_run_id,)
            )
            results = [dict(r) for r in await cursor.fetchall()]
        if not results:
            return {"message": f"未找到工作流ID为 {workflow_run_id} 的点踩原因数据", "data": []}
        return {"data": results}
    except Exception as e:
        logger.error(f"根据工作流ID获取点踩原因数据时出错: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"根据工作流ID获取点踩原因数据时出错: {str(e)}")

if __name__ == "__main__":
    # 使用Uvicorn运行FastAPI应用
    import uvicorn
    uvicorn.run("rag2_0.api.AnswerType_api:app", host="0.0.0.0", port=8003, reload=False, workers=1, log_level="info") 
    # 生产环境可以使用以下命令启动:
    # uvicorn rag2_0.api.AnswerType_api:app --host 0.0.0.0 --port 8003 --workers 1 