"""
Platform Billing Webhook
接收 LiteLLM spend 回调，写入 token_usage_logs 并扣减 user_balances
"""
import logging
from datetime import datetime
from fastapi import APIRouter, Request, HTTPException
from pydantic import BaseModel
from typing import Optional
from sqlalchemy import text
from app.database import SessionLocal

logger = logging.getLogger("aiflow.billing")

router = APIRouter()


class SpendEvent(BaseModel):
    """LiteLLM spend log event"""
    request_id: Optional[str] = None
    api_key: Optional[str] = None
    model: Optional[str] = None
    total_tokens: int = 0
    prompt_tokens: int = 0
    completion_tokens: int = 0
    spend: float = 0.0
    user: Optional[str] = None  # platform user_id passed via metadata
    metadata: Optional[dict] = None
    startTime: Optional[str] = None
    endTime: Optional[str] = None


@router.post("/webhook")
async def billing_webhook(event: SpendEvent, request: Request):
    """
    接收 LiteLLM spend webhook 回调
    记录 token 使用并扣减用户余额
    """
    # 验证来源
    api_key = event.api_key or ""
    metadata = event.metadata or {}

    # 只处理 mumuainovel 的用量
    app_name = metadata.get("app", "")
    if app_name != "mumuainovel" and "mumu" not in api_key:
        # 也检查是否有自定义的 X-Source header
        source = request.headers.get("X-Source", "")
        if source != "mumuainovel":
            return {"status": "skipped", "reason": "not_mumuainovel"}

    user_id = event.user or metadata.get("user_id", "mumu-user")
    model_name = event.model or "unknown"
    spend = max(event.spend, 0)
    total_tokens = event.total_tokens
    prompt_tokens = event.prompt_tokens
    completion_tokens = event.completion_tokens

    if total_tokens == 0 and spend == 0:
        return {"status": "skipped", "reason": "zero_usage"}

    db = SessionLocal()
    try:
        # 写入 token_usage_logs
        db.execute(text("""
            INSERT INTO token_usage_logs (id, user_id, model_name, app_id, app_name,
                prompt_tokens, completion_tokens, total_tokens, cost, created_at)
            VALUES (gen_random_uuid(), :uid, :model, :app_id, :app_name,
                :pt, :ct, :tt, :cost, :ts)
        """), {
            "uid": user_id,
            "model": model_name,
            "app_id": "mumuainovel-001",
            "app_name": "MuMuAINovel",
            "pt": prompt_tokens,
            "ct": completion_tokens,
            "tt": total_tokens,
            "cost": spend,
            "ts": datetime.utcnow(),
        })

        # 扣减余额
        result = db.execute(text("""
            UPDATE user_balances
            SET balance = balance - :cost,
                total_used = total_used + :cost,
                updated_at = NOW()
            WHERE user_id = :uid AND balance >= :cost
            RETURNING id
        """), {"cost": spend, "uid": user_id})

        if result.rowcount == 0:
            # 余额不足或无记录，创建新余额记录
            # 检查是否存在
            exists = db.execute(text(
                "SELECT id FROM user_balances WHERE user_id = :uid"
            ), {"uid": user_id}).fetchone()
            if not exists:
                db.execute(text("""
                    INSERT INTO user_balances (id, user_id, balance, total_used, created_at, updated_at)
                    VALUES (gen_random_uuid(), :uid, :balance, :cost, NOW(), NOW())
                """), {"uid": user_id, "balance": 10.0 - spend, "cost": spend})

        db.commit()
        logger.info(f"Billing: user={user_id} model={model_name} tokens={total_tokens} spend=${spend:.6f}")

    except Exception as e:
        db.rollback()
        logger.error(f"Billing webhook error: {e}")
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        db.close()

    return {"status": "ok", "user_id": user_id, "spend": spend}


@router.get("/webhook/health")
async def billing_webhook_health():
    return {"status": "ok", "service": "billing-webhook"}