"""身份证识别 API 服务（极限低配版：1 核 CPU / 2G 内存）。

POST /api/v1/idcard/recognize        上传图片（multipart）
POST /api/v1/idcard/recognize/base64 base64 JSON
POST /api/v1/auth/token              用 client_secret 换短期 token（可选）
GET  /api/v1/idcard/health           健康检查（含内存读数）
GET  /                               服务信息

鉴权三层（全部标准库实现，无额外依赖）：
  1. 静态 token：API_TOKENS=名称:token,名称2:token2，每个调用方一把，可单独吊销
  2. 请求头：Authorization: Bearer <token>（标准），兼容 X-API-Key
  3. 短期 token：设 CLIENT_SECRET 后由 /api/v1/auth/token 签发，HMAC-SHA256 签名带过期

低配适配策略见 README，核心五点：
  1. 单 worker + 引擎全局单例，绝不加载多份模型
  2. 线程三重钉死：OMP/OPENBLAS/MKL=1 + cv2.setNumThreads(0)
     —— Docker 的 cpus 限额不改变容器内 os.cpu_count()，onnxruntime 会按
        宿主机核数建线程池，1 核机器上直接线程风暴，必须显式限死
  3. 关闭方向分类器（身份证基本正向）
  4. 输入图最长边 1200px + det_limit_side_len=960
  5. content-length 预检 + 惰性尺寸检查，防大图/解压炸弹打爆内存
"""
from __future__ import annotations

import base64
import binascii
import hashlib
import hmac
import io
import json
import logging
import os
import time
from typing import Optional

# —— 线程限制必须早于 numpy / onnxruntime import 才会生效 ——
# setdefault：不覆盖用户显式设置，未设置时兜底为 1
for _v in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS"):
    os.environ.setdefault(_v, "1")

import numpy as np
from fastapi import FastAPI, File, Header, HTTPException, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from PIL import Image
from pydantic import BaseModel, Field

try:  # OpenCV 自建线程池，1 核机器上是纯负担
    import cv2

    cv2.setNumThreads(0)
except Exception:  # pragma: no cover - headless 缺失时不影响主流程
    cv2 = None

from . import parser as id_parser

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("idcard-api")

MAX_IMAGE_BYTES = int(os.environ.get("MAX_IMAGE_BYTES", 5 * 1024 * 1024))
MAX_SIDE = int(os.environ.get("MAX_SIDE", 1200))              # 输入图最长边，超过则等比缩小
MAX_DIMENSION = int(os.environ.get("MAX_DIMENSION", 6000))    # header 尺寸硬上限，超过直接拒
DET_LIMIT_SIDE_LEN = int(os.environ.get("DET_LIMIT_SIDE_LEN", 960))
REC_BATCH_NUM = int(os.environ.get("REC_BATCH_NUM", 1))       # 低内存必须为 1
MODEL_DIR = os.environ.get("RAPIDOCR_MODEL_DIR") or None
# 允许跨域的前端来源，逗号分隔；默认 * 方便本地前端开发。
# 生产环境建议收紧为实际地址，例如 http://10.1.1.163:5173,http://nas.local
CORS_ORIGINS = os.environ.get("CORS_ORIGINS", "*").strip()

# —— 鉴权配置 ——
# 1) 静态 token（主力用法）：每个调用方发一把，可单独吊销。
#    格式：名称:token,名称2:token2   例如 frontend:abc123,win-service:def456
API_TOKENS_RAW = os.environ.get("API_TOKENS", "").strip()
# 2) 兼容旧配置：单个 X-API-Key
API_KEY = os.environ.get("API_KEY", "").strip()
# 3) 可选短期 token：设了 CLIENT_SECRET 后启用 /api/v1/auth/token 换 token
CLIENT_ID = os.environ.get("CLIENT_ID", "default").strip()
CLIENT_SECRET = os.environ.get("CLIENT_SECRET", "").strip()
TOKEN_TTL = int(os.environ.get("TOKEN_TTL", 3600))            # 短期 token 有效期（秒）


def _parse_tokens() -> dict:
    """解析 API_TOKENS 为 {token: 调用方名称}。token 里不要含逗号/冒号。"""
    tokens = {}
    for item in API_TOKENS_RAW.split(","):
        item = item.strip()
        if not item:
            continue
        if ":" in item:
            name, tok = item.split(":", 1)
            if tok.strip():
                tokens[tok.strip()] = name.strip() or "client"
        else:
            tokens[item] = "client"
    if API_KEY:
        tokens[API_KEY] = "legacy-key"
    return tokens


_TOKEN_MAP = _parse_tokens()

# 弱口令提醒：示例占位符（REPLACE_ME_*）或过短的 token 不该出现在正式环境。
# 只告警不阻止启动，避免影响联调；换成真随机串后这条日志自动消失。
_weak_clients = sorted(
    {name for tok, name in _TOKEN_MAP.items() if "REPLACE_ME" in tok.upper() or len(tok) < 12}
)
if _weak_clients:
    log.warning(
        "token 过弱或仍是示例占位符，请尽快更换：%s（生成：openssl rand -hex 24）",
        ", ".join(_weak_clients),
    )

_engine = None


def get_engine():
    """懒加载 OCR 引擎（单例）。

    rapidocr-onnxruntime 不同版本的构造参数略有差异，这里逐级降级尝试，
    避免因某个参数不存在直接启动失败。生产环境请锁定 requirements.txt 版本。
    """
    global _engine
    if _engine is None:
        from rapidocr_onnxruntime import RapidOCR

        log.info("loading RapidOCR engine (model_dir=%s)", MODEL_DIR)
        base = {"model_dir": MODEL_DIR} if MODEL_DIR else {}
        for kwargs in (
            {
                **base,
                "rec_batch_num": REC_BATCH_NUM,
                "det_limit_side_len": DET_LIMIT_SIDE_LEN,
                "use_angle_cls": False,
            },
            {**base, "rec_batch_num": REC_BATCH_NUM, "use_angle_cls": False},
            {**base, "rec_batch_num": REC_BATCH_NUM},
            base,
        ):
            try:
                _engine = RapidOCR(**kwargs)
                log.info("engine ready with kwargs=%s", sorted(kwargs))
                break
            except TypeError as ex:
                log.warning("RapidOCR rejected %s: %s", sorted(kwargs), ex)
        if _engine is None:
            raise RuntimeError("cannot initialize RapidOCR engine")
    return _engine


app = FastAPI(
    title="IDCard Recognition API",
    version="2.0.0",
    description="离线身份证信息识别服务，返回结构化 JSON。极限低配硬件友好（1C/2G）。",
)

# CORS：浏览器里的前端页面跨源调用时必须开启，否则预检（OPTIONS）直接失败。
# 后端到后端调用（curl / requests / 其他容器）不受 CORS 影响，开不开都能用。
# 鉴权走 X-API-Key 请求头，不用 cookie，所以 allow_credentials 保持 False
# （allow_origins=["*"] 与 credentials=True 是非法组合，浏览器会拒绝）。
_cors_origins = [o.strip() for o in CORS_ORIGINS.split(",") if o.strip()] or ["*"]
app.add_middleware(
    CORSMiddleware,
    allow_origins=_cors_origins,
    allow_credentials=False,
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["*"],
    expose_headers=["*"],
    max_age=600,
)
if "*" in _cors_origins and not _TOKEN_MAP and not CLIENT_SECRET:
    log.warning(
        "CORS 允许全部来源且未配置任何 token：局域网内任意网页都能读取识别结果。"
        " 生产环境请设置 API_TOKENS 并收紧 CORS_ORIGINS。"
    )


@app.middleware("http")
async def _limit_payload(request: Request, call_next):
    """第一道防线：按 content-length 预检，别让超大请求先读进内存。"""
    cl = request.headers.get("content-length")
    if cl:
        try:
            if int(cl) > MAX_IMAGE_BYTES:
                return JSONResponse(
                    status_code=413,
                    content={"detail": "payload too large, max %d bytes" % MAX_IMAGE_BYTES},
                )
        except ValueError:
            pass
    return await call_next(request)


class Base64Body(BaseModel):
    image_base64: str = Field(..., description="图片 base64，可带 data:image/... 前缀")
    side: str = Field("auto", description="auto / front / back")
    mask: bool = Field(False, description="是否脱敏返回")


class TokenRequest(BaseModel):
    client_id: str = Field("default", description="调用方标识")
    client_secret: str = Field(..., description="客户端密钥（与 CLIENT_SECRET 一致）")


def _b64u(raw: bytes) -> str:
    return base64.urlsafe_b64encode(raw).rstrip(b"=").decode()


def _b64u_decode(text: str) -> bytes:
    return base64.urlsafe_b64decode(text + "=" * (-len(text) % 4))


def _sign(payload: str) -> str:
    """HMAC-SHA256 签名（标准库，无额外依赖）。"""
    digest = hmac.new(CLIENT_SECRET.encode(), payload.encode(), hashlib.sha256).digest()
    return _b64u(digest)


def issue_token(subject: str) -> dict:
    """签发短期 token：base64(payload).hmac  —— 结构类似 JWT，但不依赖 PyJWT。"""
    payload = _b64u(
        json.dumps(
            {"sub": subject, "iat": int(time.time()), "exp": int(time.time()) + TOKEN_TTL},
            separators=(",", ":"),
        ).encode()
    )
    return {"token": payload + "." + _sign(payload), "expires_in": TOKEN_TTL, "token_type": "Bearer"}


def _verify_short_token(token: str) -> Optional[str]:
    """校验短期 token 的签名和有效期，成功返回 subject。"""
    if not CLIENT_SECRET or "." not in token:
        return None
    payload, sig = token.rsplit(".", 1)
    if not hmac.compare_digest(sig, _sign(payload)):   # 常数时间比较，防时序侧信道
        return None
    try:
        data = json.loads(_b64u_decode(payload))
    except Exception:
        return None
    # 到达 exp 时刻即失效（用 >=，避免整数截断让 token 多活最多 1 秒）
    if int(time.time()) >= int(data.get("exp", 0)):
        return None
    return str(data.get("sub") or "token")


def _authenticate(authorization: Optional[str], x_api_key: Optional[str]) -> str:
    """统一鉴权，返回调用方名称用于日志。

    支持两种请求头（同时兼容）：
      Authorization: Bearer <token>     ← 标准写法，推荐
      X-API-Key: <token>                ← 旧写法，保留兼容
    token 可以是 API_TOKENS 里的静态 token，也可以是 /auth/token 签发的短期 token。
    完全没配置鉴权时放行（返回 anonymous），方便本地调试。
    """
    if not _TOKEN_MAP and not CLIENT_SECRET:
        return "anonymous"

    token = ""
    if authorization and authorization.lower().startswith("bearer "):
        token = authorization[7:].strip()
    elif x_api_key:
        token = x_api_key.strip()

    if not token:
        raise HTTPException(
            status_code=401,
            detail="missing token",
            headers={"WWW-Authenticate": "Bearer"},
        )
    if token in _TOKEN_MAP:
        return _TOKEN_MAP[token]
    subject = _verify_short_token(token)
    if subject:
        return subject
    raise HTTPException(
        status_code=401,
        detail="invalid or expired token",
        headers={"WWW-Authenticate": "Bearer"},
    )


def _load_image(data: bytes) -> Image.Image:
    """读图 + 尺寸收敛。低配关键：图越大，解码缓冲和 det 张量越吃内存。"""
    if len(data) > MAX_IMAGE_BYTES:
        raise HTTPException(status_code=413, detail="image too large, max %d bytes" % MAX_IMAGE_BYTES)
    try:
        img = Image.open(io.BytesIO(data))  # 惰性：只读 header，不解码
    except Exception:
        raise HTTPException(status_code=400, detail="invalid image file")
    # 解压炸弹防护：解码前按 header 尺寸拦截。
    # 6000px 是硬上限（解码后约 108MB 缓冲，可接受）；手机常见 3000-4000px 拍照
    # 必须放行、交给下面的 thumbnail 缩小，不能用 MAX_SIDE 当阈值。
    if max(img.size) > MAX_DIMENSION:
        raise HTTPException(
            status_code=413, detail="image dimension too large, max %d px" % MAX_DIMENSION
        )
    try:
        img.load()
    except Exception:
        raise HTTPException(status_code=400, detail="invalid image file")
    if img.mode != "RGB":
        img = img.convert("RGB")
    if max(img.size) > MAX_SIDE:
        img.thumbnail((MAX_SIDE, MAX_SIDE), Image.LANCZOS)
    return img


def _run_ocr(img: Image.Image) -> dict:
    engine = get_engine()
    started = time.perf_counter()
    result, _elapse = engine(np.array(img))
    lines = id_parser._lines_from_result(result)
    data = id_parser.parse(lines)
    data["elapsed_ms"] = round((time.perf_counter() - started) * 1000, 1)
    if not lines:
        raise HTTPException(status_code=422, detail="no text detected in image")
    if data["field_count"] == 0:
        raise HTTPException(status_code=422, detail="no id card fields detected")
    return data


def _mem_mb() -> Optional[int]:
    """容器/进程内存读数，方便 docker stats 之外直接看接口。"""
    for path in (
        "/sys/fs/cgroup/memory.current",              # cgroup v2
        "/sys/fs/cgroup/memory/memory.usage_in_bytes",  # cgroup v1
    ):
        try:
            with open(path) as f:
                return int(f.read().strip()) // (1024 * 1024)
        except (OSError, ValueError):
            continue
    try:
        with open("/proc/self/status") as f:
            for line in f:
                if line.startswith("VmRSS:"):
                    return int(line.split()[1]) // 1024
    except OSError:
        pass
    return None


@app.get("/")
def root():
    """根路径给个可读的入口说明，避免直接访问域名时看到 404。"""
    return {
        "service": "idcard-api",
        "version": app.version,
        "status": "running",
        "docs": "/docs",
        "endpoints": {
            "health": "GET  /api/v1/idcard/health",
            "recognize": "POST /api/v1/idcard/recognize        (multipart, 字段名 file)",
            "recognize_base64": "POST /api/v1/idcard/recognize/base64  (JSON)",
            "token": "POST /api/v1/auth/token                  (需配置 CLIENT_SECRET)",
        },
        "auth": {
            "enabled": bool(_TOKEN_MAP) or bool(CLIENT_SECRET),
            "scheme": "Authorization: Bearer <token>（兼容 X-API-Key 头）",
            "short_lived_token": bool(CLIENT_SECRET),
        },
    }


@app.get("/api/v1/idcard/health")
def health():
    return {
        "status": "ok",
        "engine_loaded": _engine is not None,
        "mem_mb": _mem_mb(),
        "auth_enabled": bool(_TOKEN_MAP) or bool(CLIENT_SECRET),
    }


@app.post("/api/v1/auth/token")
def create_token(body: TokenRequest):
    """用 client_id + client_secret 换一个短期 token（默认 1 小时）。

    适合前端场景：不在页面里硬编码长期密钥，只持有短时 token。
    """
    if not CLIENT_SECRET:
        raise HTTPException(status_code=400, detail="token endpoint disabled, CLIENT_SECRET not set")
    if not hmac.compare_digest(body.client_secret, CLIENT_SECRET):
        raise HTTPException(status_code=401, detail="invalid client credentials")
    payload = issue_token(body.client_id or CLIENT_ID)
    log.info("token issued for client=%s", body.client_id or CLIENT_ID)
    return payload


@app.post("/api/v1/idcard/recognize")
async def recognize(
    file: UploadFile = File(..., description="身份证图片 jpg/png"),
    side: str = "auto",
    mask: bool = False,
    authorization: Optional[str] = Header(None),
    x_api_key: Optional[str] = Header(None),
):
    caller = _authenticate(authorization, x_api_key)
    if file.content_type and not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="content type must be image/*")
    data = _run_ocr(_load_image(await file.read()))
    if side in ("front", "back"):
        data["side"] = side
    # 只记调用方和耗时，不记姓名/身份证号（合规要求）
    log.info("recognize by=%s side=%s %sms", caller, data["side"], data["elapsed_ms"])
    return _respond(data, mask)


@app.post("/api/v1/idcard/recognize/base64")
def recognize_base64(
    body: Base64Body,
    authorization: Optional[str] = Header(None),
    x_api_key: Optional[str] = Header(None),
):
    caller = _authenticate(authorization, x_api_key)
    raw = body.image_base64
    if "," in raw and raw.strip().startswith("data:"):
        raw = raw.split(",", 1)[1]
    try:
        payload = base64.b64decode(raw, validate=True)
    except (binascii.Error, ValueError):
        raise HTTPException(status_code=400, detail="invalid base64 payload")
    out = _run_ocr(_load_image(payload))
    if body.side in ("front", "back"):
        out["side"] = body.side
    log.info("recognize_b64 by=%s side=%s %sms", caller, out["side"], out["elapsed_ms"])
    return _respond(out, body.mask)


def _respond(data: dict, mask: bool) -> JSONResponse:
    return JSONResponse({"code": 0, "data": id_parser.mask(data) if mask else data})