www_site/services/api/app/main.py

from datetime import datetime, timedelta, timezone
from typing import Any
from uuid import UUID

import psycopg
from fastapi import Depends, FastAPI, HTTPException, Request, Response, status
from fastapi.responses import JSONResponse
from psycopg import Connection
from psycopg.types.json import Json

from .config import settings
from .db import get_connection
from .schemas import (
    AuthResponse,
    LoginRequest,
    MembershipRequest,
    PasswordChangeRequest,
    QuotaAdjustmentRequest,
    RegisterRequest,
    StatusRequest,
    UserPublic,
    VerificationConfirmRequest,
    VerificationSendRequest,
    normalize_email,
)
from .security import (
    check_origin,
    clear_session_cookie,
    code_digest,
    current_user,
    hash_password,
    new_csrf_token,
    new_token,
    password_hasher,
    rate_limiter,
    require_admin,
    require_csrf,
    set_csrf_cookie,
    set_session_cookie,
    token_digest,
    utc_now,
    verify_password,
)

app = FastAPI(title="Kaotings Business API", version="0.2.0", docs_url=None if settings.app_env == "production" else "/docs")


def error(code: str, message: str, http_status: int) -> HTTPException:
    return HTTPException(status_code=http_status, detail={"code": code, "message": message})


def period_bounds(now: datetime | None = None) -> tuple[datetime, datetime]:
    current = now or utc_now()
    start = current.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
    if start.month == 12:
        end = start.replace(year=start.year + 1, month=1)
    else:
        end = start.replace(month=start.month + 1)
    return start, end


def effective_plan(connection: Connection, user: dict) -> str:
    if user["role"] == "admin":
        return user["plan"]
    grant = connection.execute(
        """
        SELECT plan FROM membership_grants
        WHERE user_id = %s AND revoked_at IS NULL AND starts_at <= now()
          AND (expires_at IS NULL OR expires_at > now())
        ORDER BY starts_at DESC LIMIT 1
        """,
        (user["id"],),
    ).fetchone()
    return grant["plan"] if grant else "free"


def user_public(connection: Connection, user: dict) -> UserPublic:
    return UserPublic(
        id=user["id"], email=user["email"], phone=user["phone"], role=user["role"], plan=effective_plan(connection, user),
        status=user["status"], email_verified=user["email_verified"], phone_verified=user["phone_verified"],
        created_at=user["created_at"], last_login_at=user["last_login_at"],
    )


def ensure_quota(connection: Connection, user_id: UUID, plan: str) -> dict:
    start, end = period_bounds()
    policy = connection.execute("SELECT period_limit FROM plan_policies WHERE code = %s", (plan,)).fetchone()
    if not policy:
        raise error("POLICY_MISSING", "权益策略未配置", 500)
    connection.execute(
        """
        INSERT INTO quota_accounts(user_id, period_start, period_end, limit_snapshot)
        VALUES (%s, %s, %s, %s)
        ON CONFLICT (user_id, period_start, period_end) DO NOTHING
        """,
        (user_id, start, end, policy["period_limit"]),
    )
    row = connection.execute(
        "SELECT * FROM quota_accounts WHERE user_id = %s AND period_start = %s AND period_end = %s",
        (user_id, start, end),
    ).fetchone()
    return row


def usage_view(quota: dict) -> dict[str, Any]:
    available = max(0, quota["limit_snapshot"] + quota["adjustment"] - quota["used"] - quota["reserved"])
    return {
        "period_start": quota["period_start"], "period_end": quota["period_end"], "limit": quota["limit_snapshot"],
        "adjustment": quota["adjustment"], "used": quota["used"], "reserved": quota["reserved"], "available": available,
    }


@app.exception_handler(HTTPException)
async def http_error_handler(_: Request, exc: HTTPException):
    detail = exc.detail if isinstance(exc.detail, dict) else {"code": "REQUEST_FAILED", "message": str(exc.detail)}
    return JSONResponse(status_code=exc.status_code, content={"error": detail})


@app.get("/healthz")
def health(connection: Connection = Depends(get_connection)):
    connection.execute("SELECT 1")
    return {"status": "ok", "database": "ok", "service": "api"}


@app.get("/api/v1/auth/csrf")
def csrf(response: Response, request: Request):
    check_origin(request)
    token = request.cookies.get("kaotings_csrf") or new_csrf_token()
    set_csrf_cookie(response, token)
    return {"csrf_token": token}


@app.post("/api/v1/auth/register", response_model=AuthResponse, status_code=201)
def register(payload: RegisterRequest, request: Request, connection: Connection = Depends(get_connection)):
    check_origin(request)
    key = f"register:{request.client.host if request.client else 'unknown'}"
    if not rate_limiter.allow(key, 5, 3600):
        raise error("RATE_LIMITED", "请求过于频繁", 429)
    email = normalize_email(str(payload.email))
    try:
        user = connection.execute(
            """
            INSERT INTO users(email, password_hash) VALUES (%s, %s)
            RETURNING *
            """,
            (email, hash_password(payload.password)),
        ).fetchone()
        ensure_quota(connection, user["id"], "free")
        connection.commit()
    except psycopg.errors.UniqueViolation:
        connection.rollback()
        raise error("REGISTRATION_FAILED", "注册信息不可用", 409)
    return {"user": user_public(connection, user)}


@app.post("/api/v1/auth/login", response_model=AuthResponse)
def login(payload: LoginRequest, request: Request, response: Response, connection: Connection = Depends(get_connection)):
    check_origin(request)
    key = f"login:{request.client.host if request.client else 'unknown'}:{normalize_email(str(payload.email))}"
    if not rate_limiter.allow(key, 10, 300):
        raise error("RATE_LIMITED", "请求过于频繁", 429)
    user = connection.execute("SELECT * FROM users WHERE email = %s", (normalize_email(str(payload.email)),)).fetchone()
    if not user or not verify_password(payload.password, user["password_hash"]):
        raise error("LOGIN_FAILED", "邮箱或密码错误", 401)
    if user["status"] != "active":
        raise error("ACCOUNT_DISABLED", "账户不可用", 403)
    token = new_token()
    connection.execute("INSERT INTO sessions(user_id, token_hash, expires_at) VALUES (%s, %s, %s)", (user["id"], token_digest(token), utc_now() + timedelta(seconds=settings.session_ttl_seconds)))
    user = connection.execute("UPDATE users SET last_login_at = now(), updated_at = now() WHERE id = %s RETURNING *", (user["id"],)).fetchone()
    ensure_quota(connection, user["id"], effective_plan(connection, user))
    connection.commit()
    set_session_cookie(response, token)
    return {"user": user_public(connection, user)}


@app.post("/api/v1/auth/logout", status_code=204, dependencies=[Depends(require_csrf)])
def logout(response: Response, request: Request, connection: Connection = Depends(get_connection)):
    token = request.cookies.get(settings.session_cookie_name)
    if token:
        connection.execute("UPDATE sessions SET revoked_at = now() WHERE token_hash = %s", (token_digest(token),))
        connection.commit()
    clear_session_cookie(response)
    return Response(status_code=204)


@app.get("/api/v1/auth/me", response_model=UserPublic)
def me(user: dict = Depends(current_user), connection: Connection = Depends(get_connection)):
    return user_public(connection, user)


@app.post("/api/v1/auth/password/change", dependencies=[Depends(require_csrf)])
def change_password(payload: PasswordChangeRequest, request: Request, connection: Connection = Depends(get_connection), user: dict = Depends(current_user)):
    if not verify_password(payload.current_password, user["password_hash"]):
        raise error("PASSWORD_INVALID", "当前密码错误", 400)
    connection.execute("UPDATE users SET password_hash = %s, updated_at = now() WHERE id = %s", (hash_password(payload.new_password), user["id"]))
    connection.execute("UPDATE sessions SET revoked_at = now() WHERE user_id = %s AND id <> %s", (user["id"], user["session_id"]))
    connection.commit()
    return {"status": "ok"}


@app.get("/api/v1/account/usage")
def account_usage(user: dict = Depends(current_user), connection: Connection = Depends(get_connection)):
    quota = ensure_quota(connection, user["id"], effective_plan(connection, user))
    connection.commit()
    return {"plan": effective_plan(connection, user), **usage_view(quota)}


@app.post("/api/v1/auth/verification/send")
def verification_send(payload: VerificationSendRequest, request: Request, user: dict = Depends(current_user)):
    check_origin(request)
    enabled = settings.email_verification_enabled if payload.channel == "email" else settings.phone_verification_enabled
    if not enabled:
        raise error("VERIFICATION_NOT_ENABLED", "验证渠道尚未启用", 503)
    raise error("VERIFICATION_PROVIDER_UNAVAILABLE", "验证渠道尚未配置", 503)


@app.post("/api/v1/auth/verification/confirm", dependencies=[Depends(require_csrf)])
def verification_confirm(payload: VerificationConfirmRequest, request: Request, connection: Connection = Depends(get_connection), user: dict = Depends(current_user)):
    check_origin(request)
    challenge = connection.execute("SELECT * FROM verification_challenges WHERE id = %s AND user_id = %s", (payload.challenge_id, user["id"])).fetchone()
    if not challenge or challenge["consumed_at"] or challenge["expires_at"] <= utc_now() or challenge["attempt_count"] >= 5:
        raise error("VERIFICATION_INVALID", "验证挑战无效或已过期", 400)
    if not __import__("hmac").compare_digest(challenge["code_digest"], code_digest(payload.code)):
        connection.execute("UPDATE verification_challenges SET attempt_count = attempt_count + 1 WHERE id = %s", (payload.challenge_id,))
        connection.commit()
        raise error("VERIFICATION_INVALID", "验证码错误", 400)
    connection.execute("UPDATE verification_challenges SET consumed_at = now() WHERE id = %s", (payload.challenge_id,))
    if challenge["channel"] == "email":
        connection.execute("UPDATE users SET email_verified = true, updated_at = now() WHERE id = %s", (user["id"],))
    else:
        connection.execute("UPDATE users SET phone_verified = true, updated_at = now() WHERE id = %s", (user["id"],))
    connection.commit()
    return {"status": "ok"}


@app.get("/api/v1/admin/users")
def admin_users(user: dict = Depends(require_admin), connection: Connection = Depends(get_connection)):
    rows = connection.execute("SELECT id, email, phone, role, plan, status, email_verified, phone_verified, created_at, last_login_at FROM users ORDER BY created_at DESC LIMIT 100").fetchall()
    return {"items": rows}


@app.patch("/api/v1/admin/users/{user_id}/status", dependencies=[Depends(require_csrf)])
def admin_status(user_id: UUID, payload: StatusRequest, connection: Connection = Depends(get_connection), actor: dict = Depends(require_admin)):
    target = connection.execute("SELECT * FROM users WHERE id = %s", (user_id,)).fetchone()
    if not target:
        raise error("NOT_FOUND", "用户不存在", 404)
    if target["role"] == "admin" and payload.status == "disabled":
        count = connection.execute("SELECT count(*) FROM users WHERE role = 'admin' AND status = 'active'").fetchone()["count"]
        if count <= 1:
            raise error("LAST_ADMIN_PROTECTED", "不能禁用最后一个可用管理员", 409)
    updated = connection.execute("UPDATE users SET status = %s, updated_at = now() WHERE id = %s RETURNING id, status", (payload.status, user_id)).fetchone()
    connection.execute("INSERT INTO admin_audit_logs(actor_id, action, target_type, target_id, before_value, after_value, reason) VALUES (%s, 'user_status', 'user', %s, %s, %s, %s)", (actor["id"], user_id, Json({"status": target["status"]}), Json({"status": updated["status"]}), payload.reason))
    connection.commit()
    return updated


@app.put("/api/v1/admin/users/{user_id}/membership", dependencies=[Depends(require_csrf)])
def admin_membership(user_id: UUID, payload: MembershipRequest, connection: Connection = Depends(get_connection), actor: dict = Depends(require_admin)):
    target = connection.execute("SELECT * FROM users WHERE id = %s", (user_id,)).fetchone()
    if not target:
        raise error("NOT_FOUND", "用户不存在", 404)
    starts_at = payload.starts_at or utc_now()
    connection.execute("UPDATE membership_grants SET revoked_at = now() WHERE user_id = %s AND revoked_at IS NULL", (user_id,))
    grant = connection.execute("INSERT INTO membership_grants(user_id, plan, starts_at, expires_at, created_by, reason) VALUES (%s, %s, %s, %s, %s, %s) RETURNING id, plan, starts_at, expires_at", (user_id, payload.plan, starts_at, payload.expires_at, actor["id"], payload.reason)).fetchone()
    connection.execute("UPDATE users SET plan = %s, updated_at = now() WHERE id = %s", (payload.plan, user_id))
    connection.execute("INSERT INTO admin_audit_logs(actor_id, action, target_type, target_id, before_value, after_value, reason) VALUES (%s, 'membership', 'user', %s, %s, %s, %s)", (actor["id"], user_id, Json({"plan": target["plan"]}), Json({"plan": payload.plan, "expires_at": payload.expires_at.isoformat() if payload.expires_at else None}), payload.reason))
    connection.commit()
    return grant


@app.post("/api/v1/admin/users/{user_id}/quota-adjustments", dependencies=[Depends(require_csrf)])
def admin_quota(user_id: UUID, payload: QuotaAdjustmentRequest, connection: Connection = Depends(get_connection), actor: dict = Depends(require_admin)):
    target = connection.execute("SELECT * FROM users WHERE id = %s", (user_id,)).fetchone()
    if not target:
        raise error("NOT_FOUND", "用户不存在", 404)
    quota = ensure_quota(connection, user_id, effective_plan(connection, target))
    existing = connection.execute("SELECT id FROM usage_records WHERE user_id = %s AND idempotency_key = %s", (user_id, payload.idempotency_key)).fetchone()
    if existing:
        raise error("IDEMPOTENCY_CONFLICT", "该调整已提交", 409)
    connection.execute("UPDATE quota_accounts SET adjustment = adjustment + %s, version = version + 1 WHERE id = %s", (payload.amount, quota["id"]))
    connection.execute("INSERT INTO usage_records(user_id, quota_account_id, type, amount, idempotency_key) VALUES (%s, %s, 'adjust', %s, %s)", (user_id, quota["id"], payload.amount, payload.idempotency_key))
    connection.execute("INSERT INTO admin_audit_logs(actor_id, action, target_type, target_id, after_value, reason) VALUES (%s, 'quota_adjustment', 'user', %s, %s, %s)", (actor["id"], user_id, Json({"amount": payload.amount}), payload.reason))
    connection.commit()
    return {"status": "ok", "amount": payload.amount}