app/request_log_middleware.py

"""请求日志中间件：捕获每个业务请求的完整入参/出参/耗时/错误。

记录范围：/v1/* 和 /admin/api/*
- 流式响应：只记录前 N 个 chunk 摘要 + 总 chunk 数 + 总字节数（避免内存爆炸）
- 请求体：缓存原始 bytes，供下游读取（FastAPI 的 request.body() 会消费）
- 敏感标头/key 自动脱敏
- 失败不影响主流程（记录本身失败也静默）

注意：BaseHTTPMiddleware 会消费 request body，需要用 request.state 透传缓存。
这里采用读 body → 存 request.state.cached_body → 下游通过依赖重新注入的方式。
但 FastAPI 的 Request.body() 内部有缓存，如果中间件先读了，下游再读会拿到缓存。
所以直接 await request.body() 即可，下游能拿到相同结果。
"""
from __future__ import annotations

import json
import time
from typing import Any, Awaitable, Callable

from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import Response, StreamingResponse

from .services import request_log_store


# 记录范围
_RECORDED_PREFIXES = ("/v1/", "/admin/api/", "/u/api/")
# 排除项（避免日志刷屏：admin HTML 轮询、health）
_EXCLUDED_PATHS = {"/admin/api/request-log"}  # 避免查询日志本身被记录导致递归刷屏

# 流式响应采样上限
_MAX_STREAM_CHUNKS_LOG = 20
_MAX_STREAM_BYTES_LOG = 8 * 1024


def _should_record(path: str) -> bool:
    if not path:
        return False
    if path in _EXCLUDED_PATHS:
        return False
    return any(path.startswith(p) for p in _RECORDED_PREFIXES)


def _extract_access_key(request: Request) -> str:
    auth = request.headers.get("authorization")
    if auth:
        a = auth.strip()
        if a.lower().startswith("bearer "):
            return a[7:].strip()
        if a.lower().startswith("basic "):
            return a[6:].strip()
        return a
    return (
        request.headers.get("x-xtc-access-key")
        or request.headers.get("x-xtc-access-token")
        or ""
    )


def _extract_device_id(request: Request) -> str:
    """从请求头 / query 参数提取设备 ID。

    识别顺序：
    1. ``x-xtc-device-id`` 头（推荐，前端 fetch 可直接设置）
    2. ``x-device-id`` 头（兼容通用 SDK）
    3. ``device_id`` query 参数（兜底）
    """
    val = (
        request.headers.get("x-xtc-device-id")
        or request.headers.get("x-device-id")
    )
    if val:
        s = str(val).strip()
        if s:
            return s
    qid = request.query_params.get("device_id")
    if qid:
        s = str(qid).strip()
        if s:
            return s
    return ""


def _safe_json_loads(s: str) -> Any:
    try:
        return json.loads(s)
    except Exception:
        return s


class RequestLogMiddleware(BaseHTTPMiddleware):
    """记录业务请求的完整日志。"""

    async def dispatch(
        self,
        request: Request,
        call_next: Callable[[Request], Awaitable[Response]],
    ) -> Response:
        path = request.url.path
        if not _should_record(path):
            return await call_next(request)

        start = time.time()
        method = request.method
        query = str(request.url.query) if request.url.query else None

        # 读请求体（缓存到 state，下游可复用）
        request_body_raw: bytes = b""
        try:
            request_body_raw = await request.body()
        except Exception:
            pass

        # 请求标头脱敏
        req_headers = request_log_store._redact_headers(dict(request.headers))

        # 请求体脱敏
        req_body_redacted: Any = None
        if request_body_raw:
            ct = (request.headers.get("content-type") or "").lower()
            if "application/json" in ct:
                try:
                    req_body_redacted = request_log_store._redact_body(json.loads(request_body_raw))
                except Exception:
                    req_body_redacted = request_body_raw.decode("utf-8", "replace")
            elif "multipart/form-data" in ct:
                req_body_redacted = "[multipart form-data]"
            else:
                req_body_redacted = request_body_raw.decode("utf-8", "replace")

        access_key = _extract_access_key(request)
        client_ip = request.client.host if request.client else None
        user_agent = request.headers.get("user-agent")
        device_id = _extract_device_id(request)

        # 公共日志构造
        def _build_record(status_code, elapsed_ms, ok, provider, model, is_stream,
                          error_code, error_message, resp_body, resp_headers):
            return {
                "method": method, "path": path, "query": query,
                "status_code": status_code, "elapsed_ms": elapsed_ms,
                "access_key": access_key, "client_ip": client_ip,
                "user_agent": user_agent,
                "provider": provider, "model": model, "stream": is_stream,
                "ok": ok, "error_code": error_code, "error_message": error_message,
                "request_headers": json.dumps(req_headers, ensure_ascii=False),
                "request_body": json.dumps(req_body_redacted, ensure_ascii=False) if req_body_redacted is not None else None,
                "response_body": json.dumps(resp_body, ensure_ascii=False) if resp_body is not None else None,
                "response_headers": json.dumps(resp_headers, ensure_ascii=False),
                "device_id": device_id,
            }

        # 调用下游
        response: Response
        try:
            response = await call_next(request)
        except Exception as e:
            elapsed_ms = int((time.time() - start) * 1000)
            try:
                request_log_store.insert(_build_record(
                    500, elapsed_ms, False, None, None, False,
                    "internal_error", str(e), None, None,
                ))
            except Exception:
                pass
            raise

        # 从响应头提取业务信息
        provider = response.headers.get("x-xtc-provider")
        model = response.headers.get("x-xtc-model")
        is_stream = response.headers.get("content-type", "").startswith("text/event-stream")
        status_code = response.status_code
        ok = 200 <= status_code < 400
        resp_headers = request_log_store._redact_headers(dict(response.headers))

        if isinstance(response, StreamingResponse):
            # 流式响应：包装 iterator，在迭代过程中采集 chunk，迭代结束后写日志
            return self._wrap_streaming(
                response, start, _build_record,
                provider, model, is_stream, status_code, ok, resp_headers,
            )

        # 普通响应：直接读 body
        error_code = None
        error_message = None
        resp_body_str: Any = None
        try:
            body_bytes = b""
            async for chunk in response.body_iterator:
                body_bytes += chunk
            response.body_iterator = _iter_bytes(body_bytes)
            ct = (response.headers.get("content-type") or "").lower()
            if "application/json" in ct:
                try:
                    parsed = json.loads(body_bytes)
                    resp_body_str = parsed
                    if isinstance(parsed, dict) and parsed.get("ok") is False:
                        err = parsed.get("error") or {}
                        error_code = err.get("code")
                        error_message = err.get("message")
                except Exception:
                    resp_body_str = body_bytes.decode("utf-8", "replace")
            else:
                resp_body_str = body_bytes.decode("utf-8", "replace")
                if len(resp_body_str) > 4096:
                    resp_body_str = resp_body_str[:4096] + "...[truncated]"
        except Exception as e:
            resp_body_str = f"[capture failed: {e}]"

        elapsed_ms = int((time.time() - start) * 1000)
        try:
            request_log_store.insert(_build_record(
                status_code, elapsed_ms, ok, provider, model, is_stream,
                error_code, error_message, resp_body_str, resp_headers,
            ))
        except Exception as e:
            print(f"[request_log] insert failed: {e}")

        return response

    def _wrap_streaming(
        self, response: StreamingResponse, start: float, _build_record,
        provider, model, is_stream, status_code, ok, resp_headers,
    ) -> StreamingResponse:
        """包装流式响应：边发送边采集，发送结束后写日志。"""
        original_iterator = response.body_iterator
        chunks_collected: list[str] = []
        total_bytes = 0
        total_chunks = 0
        # 流式错误捕获
        error_code = None
        error_message = None

        async def wrapped():
            nonlocal total_bytes, total_chunks, error_code, error_message
            try:
                async for chunk in original_iterator:
                    if isinstance(chunk, str):
                        chunk_bytes = chunk.encode("utf-8")
                    else:
                        chunk_bytes = chunk
                    total_bytes += len(chunk_bytes)
                    total_chunks += 1
                    # 采集前 N 个 chunk
                    if total_chunks <= _MAX_STREAM_CHUNKS_LOG:
                        text = chunk_bytes.decode("utf-8", "replace")
                        chunks_collected.append(text)
                        # 尝试从 SSE data 中提取错误
                        if not ok and text.startswith("data:"):
                            payload = text[5:].strip()
                            if payload and payload != "[DONE]":
                                try:
                                    obj = json.loads(payload)
                                    if isinstance(obj, dict) and obj.get("error"):
                                        err = obj["error"]
                                        error_code = err.get("code") if isinstance(err, dict) else None
                                        error_message = err.get("message") if isinstance(err, dict) else str(err)
                                except Exception:
                                    pass
                    yield chunk
            finally:
                # 流结束后写日志
                elapsed_ms = int((time.time() - start) * 1000)
                summary = {
                    "type": "streaming",
                    "total_chunks": total_chunks,
                    "total_bytes": total_bytes,
                    "sampled_chunks": chunks_collected,
                }
                try:
                    request_log_store.insert(_build_record(
                        status_code, elapsed_ms, ok, provider, model, is_stream,
                        error_code, error_message, summary, resp_headers,
                    ))
                except Exception as e:
                    print(f"[request_log] streaming insert failed: {e}")

        response.body_iterator = wrapped()
        return response


async def _iter_bytes(data: bytes):
    yield data