# 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 threading
import queue
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 = queue.Queue()
worker_thread = None
db_lock = threading.Lock()  # 数据库操作锁

# 初始化数据库
def init_database():
    os.makedirs(DB_DIR, exist_ok=True)
    
    with db_lock:
        with closing(sqlite3.connect(DB_FILE)) as conn:
            cursor = conn.cursor()
            
            # 创建查询类型表
            cursor.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
            )
            ''')
            
            # 创建点踩原因表
            cursor.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
            )
            ''')
            
            conn.commit()
    
    logger.info("数据库初始化完成")

# 后台工作线程函数
def log_worker():
    while True:
        try:
            # 从队列获取数据，设置超时以允许线程退出
            data = log_queue.get(timeout=1.0)
            if data is None:  # 接收到退出信号
                # 处理剩余数据后再退出
                while not log_queue.empty():
                    data = log_queue.get_nowait()
                    if data is None:  # 跳过额外的停止信号
                        continue
                    process_log_data(data)
                break
                
            process_log_data(data)
            log_queue.task_done()
        except queue.Empty:
            continue
        except Exception as e:
            logger.error(f"保存查询数据时出错: {str(e)}", exc_info=True)

# 提取数据处理逻辑到单独函数
def process_log_data(data):
    try:
        with db_lock:
            with closing(sqlite3.connect(DB_FILE)) as conn:
                cursor = conn.cursor()
                
                if "dislike_reason" in data:
                    # 保存点踩原因
                    cursor.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:
                    # 保存查询类型
                    cursor.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"
                
                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_thread
    # 初始化数据库
    init_database()
    
    # 启动后台工作线程
    worker_thread = threading.Thread(target=log_worker, daemon=True)
    worker_thread.start()
    logger.info("后台日志工作线程已启动")

# 应用关闭事件
@app.on_event("shutdown")
def shutdown_event():
    global worker_thread
    if worker_thread:
        # 发送退出信号
        log_queue.put(None)
        # 等待工作线程处理剩余数据
        worker_thread.join(timeout=10.0)
        if worker_thread.is_alive():
            logger.warning("工作线程未在超时时间内退出")
        else:
            logger.info("后台日志工作线程已停止")

# 添加健康检查端点
@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:
            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:
            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:
        with db_lock:
            with closing(sqlite3.connect(DB_FILE)) as conn:
                conn.row_factory = sqlite3.Row  # 启用字典行工厂
                cursor = conn.cursor()
                
                # 查询类型统计
                cursor.execute("SELECT COUNT(*) as count FROM query_types")
                query_count = cursor.fetchone()['count']
                
                # 点踩原因统计
                cursor.execute("SELECT COUNT(*) as count FROM dislike_reasons")
                dislike_count = cursor.fetchone()['count']
                
                # 最近5条查询记录
                cursor.execute("""
                    SELECT query_type, workflow_run_id, timestamp 
                    FROM query_types 
                    ORDER BY id DESC LIMIT 5
                """)
                recent_queries = [dict(row) for row in cursor.fetchall()]
                
                # 最近5条点踩记录
                cursor.execute("""
                    SELECT dislike_reason, workflow_run_id, timestamp 
                    FROM dislike_reasons 
                    ORDER BY id DESC LIMIT 5
                """)
                recent_dislikes = [dict(row) for row in 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:
        with db_lock:
            with closing(sqlite3.connect(DB_FILE)) as conn:
                conn.row_factory = sqlite3.Row  # 启用字典行工厂
                cursor = conn.cursor()
                
                # 查询指定工作流ID的查询类型数据
                cursor.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(row) for row in 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:
        with db_lock:
            with closing(sqlite3.connect(DB_FILE)) as conn:
                conn.row_factory = sqlite3.Row  # 启用字典行工厂
                cursor = conn.cursor()
                
                # 查询指定工作流ID的点踩原因数据
                cursor.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(row) for row in 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.dify.AnswerType:app", host="0.0.0.0", port=8003, reload=False, workers=1, log_level="info") 
    # 生产环境可以使用以下命令启动:
    # uvicorn rag2_0.dify.AnswerType:app --host 0.0.0.0 --port 8003 --workers 1 