Crawl/backend/app/routers/group_keys.py

"""
分组-Key号池管理路由
支持一个分组绑定多个key，形成号池
"""
from __future__ import annotations

from datetime import datetime
from typing import List, Optional

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel

from app.db import get_pool

router = APIRouter(tags=["group_keys"])


class GroupKeyIn(BaseModel):
    """添加key到分组"""
    group_id: int
    api_key_id: int
    priority: Optional[int] = 0
    weight: Optional[int] = 1


class GroupKeyUpdate(BaseModel):
    """更新分组-key关系"""
    priority: Optional[int] = None
    weight: Optional[int] = None
    is_active: Optional[bool] = None
    status: Optional[str] = None
    status_msg: Optional[str] = None


class GroupKeyOut(BaseModel):
    id: int
    group_id: int
    group_name: str
    api_key_id: int
    api_key_name: str
    api_key_masked: str
    priority: int
    weight: int
    is_active: bool
    status: str
    status_msg: Optional[str]
    last_used_at: Optional[datetime]
    total_tokens: int
    request_count: int
    error_count: int
    created_at: datetime
    updated_at: datetime


class BatchAddKeysIn(BaseModel):
    """批量添加key到分组"""
    group_id: int
    api_key_ids: List[int]
    priority: Optional[int] = 0
    weight: Optional[int] = 1


def _mask_key(v: str) -> str:
    if len(v) <= 10:
        return '••••••••'
    return v[:6] + '••••••••' + v[-4:]


@router.get("/group-keys", response_model=List[GroupKeyOut])
async def list_group_keys(group_id: Optional[int] = None) -> List[GroupKeyOut]:
    """获取分组-key列表"""
    pool = get_pool()

    query = """
        SELECT gk.id, gk.group_id, mg.name AS group_name,
               gk.api_key_id, ak.name AS api_key_name, ak.key_value,
               gk.priority, gk.weight, gk.is_active, gk.status, gk.status_msg,
               gk.last_used_at, gk.total_tokens, gk.request_count, gk.error_count,
               gk.created_at, gk.updated_at
        FROM group_keys gk
        JOIN model_groups mg ON mg.id = gk.group_id
        JOIN api_keys ak ON ak.id = gk.api_key_id
    """

    if group_id is not None:
        query += " WHERE gk.group_id = $1"
        rows = await pool.fetch(query, group_id)
    else:
        rows = await pool.fetch(query)

    return [
        GroupKeyOut(**{
            **dict(r),
            "api_key_masked": _mask_key(r["key_value"])
        })
        for r in rows
    ]


@router.post("/group-keys", response_model=GroupKeyOut, status_code=201)
async def add_key_to_group(body: GroupKeyIn) -> GroupKeyOut:
    """添加key到分组（号池）"""
    pool = get_pool()

    # 检查分组是否存在
    group = await pool.fetchrow("SELECT id, name FROM model_groups WHERE id = $1", body.group_id)
    if not group:
        raise HTTPException(status_code=404, detail="分组不存在")

    # 检查key是否存在
    key = await pool.fetchrow("SELECT id, name, key_value FROM api_keys WHERE id = $1", body.api_key_id)
    if not key:
        raise HTTPException(status_code=404, detail="API Key不存在")

    # 检查是否已存在
    existing = await pool.fetchrow(
        "SELECT id FROM group_keys WHERE group_id = $1 AND api_key_id = $2",
        body.group_id, body.api_key_id
    )
    if existing:
        raise HTTPException(status_code=400, detail="该key已在此分组中")

    row = await pool.fetchrow(
        """
        INSERT INTO group_keys (group_id, api_key_id, priority, weight)
        VALUES ($1, $2, $3, $4)
        RETURNING id, group_id, api_key_id, priority, weight, is_active, status, status_msg,
                  last_used_at, total_tokens, request_count, error_count, created_at, updated_at
        """,
        body.group_id, body.api_key_id, body.priority or 0, body.weight or 1
    )

    return GroupKeyOut(**{
        **dict(row),
        "group_name": group["name"],
        "api_key_name": key["name"],
        "api_key_masked": _mask_key(key["key_value"]),
    })


@router.post("/group-keys/batch-add", response_model=List[GroupKeyOut], status_code=201)
async def batch_add_keys_to_group(body: BatchAddKeysIn) -> List[GroupKeyOut]:
    """批量添加key到分组"""
    pool = get_pool()

    # 检查分组
    group = await pool.fetchrow("SELECT id, name FROM model_groups WHERE id = $1", body.group_id)
    if not group:
        raise HTTPException(status_code=404, detail="分组不存在")

    results = []
    for key_id in body.api_key_ids:
        key = await pool.fetchrow("SELECT id, name, key_value FROM api_keys WHERE id = $1", key_id)
        if not key:
            continue

        # 检查是否已存在
        existing = await pool.fetchrow(
            "SELECT id FROM group_keys WHERE group_id = $1 AND api_key_id = $2",
            body.group_id, key_id
        )
        if existing:
            continue

        row = await pool.fetchrow(
            """
            INSERT INTO group_keys (group_id, api_key_id, priority, weight)
            VALUES ($1, $2, $3, $4)
            RETURNING id, group_id, api_key_id, priority, weight, is_active, status, status_msg,
                      last_used_at, total_tokens, request_count, error_count, created_at, updated_at
            """,
            body.group_id, key_id, body.priority or 0, body.weight or 1
        )

        results.append(GroupKeyOut(**{
            **dict(row),
            "group_name": group["name"],
            "api_key_name": key["name"],
            "api_key_masked": _mask_key(key["key_value"]),
        }))

    return results


@router.put("/group-keys/{id}", response_model=GroupKeyOut)
async def update_group_key(id: int, body: GroupKeyUpdate) -> GroupKeyOut:
    """更新分组-key关系"""
    pool = get_pool()

    existing = await pool.fetchrow(
        """
        SELECT gk.*, mg.name AS group_name, ak.name AS api_key_name, ak.key_value
        FROM group_keys gk
        JOIN model_groups mg ON mg.id = gk.group_id
        JOIN api_keys ak ON ak.id = gk.api_key_id
        WHERE gk.id = $1
        """,
        id
    )
    if not existing:
        raise HTTPException(status_code=404, detail="记录不存在")

    updates = []
    values = []
    param_idx = 1

    if body.priority is not None:
        updates.append(f"priority = ${param_idx}")
        values.append(body.priority)
        param_idx += 1
    if body.weight is not None:
        updates.append(f"weight = ${param_idx}")
        values.append(body.weight)
        param_idx += 1
    if body.is_active is not None:
        updates.append(f"is_active = ${param_idx}")
        values.append(body.is_active)
        param_idx += 1
    if body.status is not None:
        updates.append(f"status = ${param_idx}")
        values.append(body.status)
        param_idx += 1
    if body.status_msg is not None:
        updates.append(f"status_msg = ${param_idx}")
        values.append(body.status_msg)
        param_idx += 1

    if updates:
        updates.append("updated_at = NOW()")
        values.append(id)
        await pool.execute(
            f"UPDATE group_keys SET {', '.join(updates)} WHERE id = ${param_idx}",
            *values
        )

    # 重新获取
    row = await pool.fetchrow(
        """
        SELECT gk.id, gk.group_id, mg.name AS group_name,
               gk.api_key_id, ak.name AS api_key_name, ak.key_value,
               gk.priority, gk.weight, gk.is_active, gk.status, gk.status_msg,
               gk.last_used_at, gk.total_tokens, gk.request_count, gk.error_count,
               gk.created_at, gk.updated_at
        FROM group_keys gk
        JOIN model_groups mg ON mg.id = gk.group_id
        JOIN api_keys ak ON ak.id = gk.api_key_id
        WHERE gk.id = $1
        """,
        id
    )

    return GroupKeyOut(**{
        **dict(row),
        "api_key_masked": _mask_key(row["key_value"])
    })


@router.delete("/group-keys/{id}", status_code=204)
async def remove_key_from_group(id: int):
    """从分组中移除key"""
    pool = get_pool()
    result = await pool.execute("DELETE FROM group_keys WHERE id = $1", id)
    if result == "DELETE 0":
        raise HTTPException(status_code=404, detail="记录不存在")


@router.get("/groups/{group_id}/pool-summary")
async def get_group_pool_summary(group_id: int):
    """获取分组号池摘要（健康度统计）"""
    pool = get_pool()

    # 检查分组
    group = await pool.fetchrow("SELECT id, name FROM model_groups WHERE id = $1", group_id)
    if not group:
        raise HTTPException(status_code=404, detail="分组不存在")

    # 统计
    stats = await pool.fetchrow(
        """
        SELECT
            COUNT(*) AS total_keys,
            COUNT(*) FILTER (WHERE status = 'active' AND is_active = TRUE) AS active_keys,
            COUNT(*) FILTER (WHERE status = 'rate_limited') AS rate_limited_keys,
            COUNT(*) FILTER (WHERE status = 'banned') AS banned_keys,
            COUNT(*) FILTER (WHERE status = 'failed') AS failed_keys,
            SUM(total_tokens) AS total_tokens,
            SUM(request_count) AS total_requests,
            SUM(error_count) AS total_errors
        FROM group_keys
        WHERE group_id = $1
        """,
        group_id
    )

    # 健康度计算
    total = stats["total_keys"] or 0
    active = stats["active_keys"] or 0
    health_rate = (active / total * 100) if total > 0 else 0

    return {
        "group_id": group_id,
        "group_name": group["name"],
        "total_keys": total,
        "active_keys": active,
        "rate_limited_keys": stats["rate_limited_keys"] or 0,
        "banned_keys": stats["banned_keys"] or 0,
        "failed_keys": stats["failed_keys"] or 0,
        "health_rate": round(health_rate, 2),
        "total_tokens": stats["total_tokens"] or 0,
        "total_requests": stats["total_requests"] or 0,
        "total_errors": stats["total_errors"] or 0,
    }


@router.get("/groups/{group_id}/keys")
async def get_group_keys_detail(group_id: int):
    """获取分组的key详情列表"""
    pool = get_pool()

    rows = await pool.fetch(
        """
        SELECT gk.id, gk.api_key_id, ak.name AS api_key_name, ak.key_value,
               gk.priority, gk.weight, gk.is_active, gk.status, gk.status_msg,
               gk.last_used_at, gk.total_tokens, gk.request_count, gk.error_count,
               gk.created_at, gk.updated_at
        FROM group_keys gk
        JOIN api_keys ak ON ak.id = gk.api_key_id
        WHERE gk.group_id = $1
        ORDER BY gk.priority DESC, gk.weight DESC
        """,
        group_id
    )

    return [
        {
            "id": r["id"],
            "api_key_id": r["api_key_id"],
            "api_key_name": r["api_key_name"],
            "api_key_masked": _mask_key(r["key_value"]),
            "priority": r["priority"],
            "weight": r["weight"],
            "is_active": r["is_active"],
            "status": r["status"],
            "status_msg": r["status_msg"],
            "last_used_at": r["last_used_at"].isoformat() if r["last_used_at"] else None,
            "total_tokens": r["total_tokens"],
            "request_count": r["request_count"],
            "error_count": r["error_count"],
            "created_at": r["created_at"].isoformat() if r["created_at"] else None,
            "updated_at": r["updated_at"].isoformat() if r["updated_at"] else None,
        }
        for r in rows
    ]


# ============================================
# 推送到company_base_pool
# ============================================

@router.post("/groups/{group_id}/push")
async def push_group(group_id: int):
    """手动推送单个分组到company_base_pool"""
    pool = get_pool()
    try:
        await push_group_keys_to_pool(pool, group_id, [])
        return {"success": True, "message": "推送成功"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"推送失败: {str(e)}")


@router.post("/groups/push-all")
async def push_all_groups():
    """手动推送所有分组到company_base_pool"""
    pool = get_pool()
    try:
        await push_all_groups_to_pool(pool)
        return {"success": True, "message": "推送完成"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"推送失败: {str(e)}")