from fastapi import FastAPI, HTTPException, Header
import sqlite3
from pathlib import Path
from datetime import datetime
import requests
import time

app = FastAPI(title="即时派单 API")

API_KEY = "tuhui_dispatch_key_2026"
DESIGNER_DB_PATH = Path("/root/tuhui/backend/designer_status.db")
DISPATCH_DB_PATH = Path("/root/tuhui/backend/dispatch.db")
WECHAT_WEBHOOK = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=cc88bdef-a13f-4d7e-bdb6-ee51b68b8205"

@app.get("/health")
def health():
    return {"status": "ok", "timestamp": datetime.now().isoformat()}

@app.get("/online")
def get_online_designers(x_api_key: str = Header(None)):
    """查询当前在线设计师"""
    if x_api_key != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid API Key")
    
    conn = sqlite3.connect(str(DESIGNER_DB_PATH), timeout=30)
    conn.row_factory = sqlite3.Row
    cursor = conn.cursor()
    
    cursor.execute("SELECT real_name, last_seen FROM designer_status WHERE status='online' ORDER BY last_seen DESC")
    designers = [row["real_name"] for row in cursor.fetchall()]
    
    conn.close()
    
    return {
        "count": len(designers),
        "online": designers,
        "timestamp": datetime.now().isoformat()
    }

@app.get("/assign")
def assign_task(x_api_key: str = Header(None)):
    """
    即时派单接口 - 轮询均匀分配版
    1. 查询在线设计师
    2. 统计每个设计师的已分配任务数
    3. 选择任务最少的（轮询）
    4. 发送企业微信通知
    5. 返回派单结果
    """
    if x_api_key != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid API Key")
    
    try:
        # 1. 查询在线设计师
        conn = sqlite3.connect(str(DESIGNER_DB_PATH), timeout=30)
        dispatch_conn = sqlite3.connect(str(DISPATCH_DB_PATH), timeout=30)
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        dispatch_cursor = dispatch_conn.cursor()
        
        cursor.execute("""
            SELECT real_name FROM designer_status 
            WHERE status = 'online' 
            ORDER BY last_seen DESC
        """)
        online_designers = [row["real_name"] for row in cursor.fetchall()]
        
        if not online_designers:
            conn.close()
            dispatch_conn.close()
            raise HTTPException(status_code=400, detail="暂无在线设计师")
        
        # 2. 统计每个在线设计师的已分配任务数（今天）
        designer_workload = {}
        for designer in online_designers:
            dispatch_cursor.execute('''
                SELECT COUNT(*) as count FROM dispatch_tasks 
                WHERE assigned_to = ? 
                AND date(assigned_at) = date('now')
            ''', (designer,))
            count = dispatch_cursor.fetchone()[0]
            designer_workload[designer] = count
        
        # 3. 选择任务最少的（轮询均匀分配）
        selected_designer = min(designer_workload, key=designer_workload.get)
        min_tasks = designer_workload[selected_designer]
        
        # 4. 创建任务
        task_id = f"task_{int(datetime.now().timestamp())}"
        
        dispatch_cursor.execute('''
            INSERT INTO dispatch_tasks (id, task_name, task_description, task_type, priority, status, assigned_to, assigned_at)
            VALUES (?, ?, ?, ?, ?, 'assigned', ?, ?)
        ''', (task_id, "临时任务", "即时分配", "design", 1, selected_designer, datetime.now()))
        
        dispatch_conn.commit()
        
        # 5. 发送企业微信通知
        message = f"""📋 新任务分配

👤 设计师：{selected_designer}
📊 今日已分配：{min_tasks + 1} 个任务
📝 任务：临时任务
⏰ 时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}

请及时处理！"""
        
        try:
            requests.post(WECHAT_WEBHOOK, json={
                "msgtype": "markdown",
                "markdown": {"content": message}
            }, timeout=10)
            notification_sent = True
        except:
            notification_sent = False
        
        conn.close()
        dispatch_conn.close()
        
        return {
            "success": True,
            "task_id": task_id,
            "assigned_to": selected_designer,
            "workload": designer_workload,
            "online_count": len(online_designers),
            "notification_sent": notification_sent,
            "timestamp": datetime.now().isoformat()
        }
        
    except sqlite3.OperationalError as e:
        if "locked" in str(e):
            time.sleep(0.5)  # 等待 0.5 秒重试
            return assign_task(x_api_key)
        raise

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8006)
