app/parse_api.py

import hashlib
import ipaddress
import json
import os
import re
import threading
import time
from typing import Any
from urllib.parse import quote

from fastapi import APIRouter, Request
from fastapi.responses import JSONResponse

from . import parse_core, pipeline_dispatch, source_files
from .doc_store import parse_queue_docs
from .document_identity import content_hash
from .parse_core import ParseRequest
from .storage_lock import serving_lock


router = APIRouter()
GITHUB_TOKEN = os.environ.get("BHA_PROOFREAD_GITHUB_TOKEN", "")
PIPELINE_REPOSITORY = os.environ.get("BHA_PROOFREAD_PIPELINE_REPOSITORY", "anftm/pipeline")
PARSE_ARCHIVE_ID = 25
PARSE_QUEUE_PREFIX = "【文章待拆分】"
PARSE_OCR_BASE_URL = os.environ.get(
    "BHA_PARSE_OCR_BASE_URL",
    "https://raw.githubusercontent.com/banned-historical-archives/banned-historical-archives25/ocr_cache",
).rstrip("/")
PARSE_OCR_MAX_BYTES = 2 * 1024 * 1024
PARSE_MAX_PAYLOAD_CHARACTERS = 2_000_000
PARSE_RATE_LIMIT = 30
PARSE_RATE_WINDOW_SECONDS = 60 * 60
TRUST_PROXY = os.environ.get("BHA_TRUST_PROXY", "1" if os.environ.get("SPACE_ID") else "0") == "1"
try:
    PROXY_NETWORKS = tuple(
        ipaddress.ip_network(value.strip())
        for value in os.environ.get("BHA_TRUST_PROXY_CIDRS", "").split(",")
        if value.strip()
    )
except ValueError:
    PROXY_NETWORKS = ()
try:
    PROXY_HOPS = max(1, int(os.environ.get("BHA_TRUST_PROXY_HOPS", "1")))
except ValueError:
    PROXY_HOPS = 1
parse_requests: dict[str, list[float]] = {}
parse_rate_lock = threading.Lock()
parse_dispatches: dict[str, dict[str, Any]] = {}
parse_dispatch_lock = threading.Lock()


async def dispatch_parse(payload: dict[str, Any]) -> None:
    await pipeline_dispatch.dispatch(payload, "submit-parse", GITHUB_TOKEN, PIPELINE_REPOSITORY)


def parse_client_key(request: Request) -> str:
    peer = request.client.host if request.client else "unknown"
    trusted_peer = False
    try:
        peer_address = ipaddress.ip_address(peer)
        trusted_peer = TRUST_PROXY and bool(PROXY_NETWORKS) and any(peer_address in network for network in PROXY_NETWORKS)
    except ValueError:
        pass
    forwarded = request.headers.get("x-forwarded-for", "") if trusted_peer else ""
    candidates = [value.strip() for value in forwarded.split(",") if value.strip()]
    if candidates and len(candidates) > PROXY_HOPS:
        candidate = candidates[-PROXY_HOPS - 1]
        try:
            ipaddress.ip_address(candidate)
            return candidate
        except ValueError:
            pass
    return peer


def parse_rate_error(request: Request) -> JSONResponse | None:
    key = parse_client_key(request)
    now = time.monotonic()
    with parse_rate_lock:
        recent = [value for value in parse_requests.get(key, []) if now - value < PARSE_RATE_WINDOW_SECONDS]
        if len(recent) >= PARSE_RATE_LIMIT:
            return JSONResponse({"error": "parse rate limit exceeded"}, status_code=429)
        recent.append(now)
        parse_requests[key] = recent
        if len(parse_requests) > 10000:
            for identity in list(parse_requests):
                parse_requests[identity] = [value for value in parse_requests[identity] if now - value < PARSE_RATE_WINDOW_SECONDS]
                if not parse_requests[identity]:
                    parse_requests.pop(identity, None)
            while len(parse_requests) > 10000:
                oldest = min(parse_requests, key=lambda identity: min(parse_requests[identity]))
                parse_requests.pop(oldest, None)
    request.state.parse_rate_key = key
    request.state.parse_rate_timestamp = now
    return None


def release_parse_rate(request: Request) -> None:
    key = getattr(request.state, "parse_rate_key", None)
    timestamp = getattr(request.state, "parse_rate_timestamp", None)
    if key is None or timestamp is None:
        return
    with parse_rate_lock:
        values = parse_requests.get(key, [])
        if timestamp in values:
            values.remove(timestamp)
        if values:
            parse_requests[key] = values
        else:
            parse_requests.pop(key, None)


def duplicate_parse_response(request_id: str) -> JSONResponse | None:
    now = time.time()
    with parse_dispatch_lock:
        for key, value in list(parse_dispatches.items()):
            if now - float(value["created_at"]) > 86400:
                parse_dispatches.pop(key, None)
        existing = parse_dispatches.get(request_id)
        if existing is not None:
            return JSONResponse({
                "ok": True,
                "status": "duplicate",
                "original_status": existing["status"],
                "request_id": request_id,
                "tasks": 1,
            })
        if len(parse_dispatches) >= 1000:
            oldest = min(parse_dispatches, key=lambda key: parse_dispatches[key]["created_at"])
            parse_dispatches.pop(oldest, None)
        parse_dispatches[request_id] = {"created_at": now, "status": "processing"}
    return None


def parse_queue_source(doc_id: str) -> dict[str, Any] | None:
    source = source_files.get_source_doc(doc_id)
    if not parse_core.is_queue_source(source, PARSE_ARCHIVE_ID, PARSE_QUEUE_PREFIX):
        return None
    return source

@router.get("/api/parse/queue")
async def parse_queue(page: int = 1, page_size: int = 20):
    if page < 1 or page_size < 1 or page_size > 100:
        return JSONResponse({"error": "invalid pagination"}, status_code=400)
    with serving_lock():
        total, sources = parse_queue_docs(PARSE_ARCHIVE_ID, PARSE_QUEUE_PREFIX, page, page_size)
    return {
        "archive_id": PARSE_ARCHIVE_ID,
        "prefix": PARSE_QUEUE_PREFIX,
        "total": total,
        "page": page,
        "page_size": page_size,
        "items": [parse_core.queue_item(source) for source in sources],
    }

@router.get("/api/parse/{doc_id}/ocr")
async def parse_ocr_page(doc_id: str, page: int = 1):
    source = parse_queue_source(doc_id)
    if not source:
        return JSONResponse({"error": "parse item not found"}, status_code=404)
    article = source.get("article") or {}
    page_start = int(article.get("page_start") or 0)
    page_end = int(article.get("page_end") or 0)
    if page < page_start or page > page_end:
        return JSONResponse({"error": "page is outside the source range"}, status_code=404)
    publication_id = str(source.get("publication_id") or "")
    if not re.fullmatch(r"[0-9a-f-]{36}", publication_id, re.IGNORECASE):
        return JSONResponse({"error": "invalid publication id"}, status_code=500)
    url = f"{PARSE_OCR_BASE_URL}/{quote(publication_id, safe='')}/{page}.json"
    try:
        response = await source_files.source_client.get(url, follow_redirects=True)
        response.raise_for_status()
        if len(response.content) > PARSE_OCR_MAX_BYTES:
            return JSONResponse({"error": "OCR page is too large"}, status_code=502)
        payload = response.json()
    except Exception as exc:
        return JSONResponse({"error": f"OCR page unavailable: {exc}"}, status_code=502)
    if not isinstance(payload, dict) or not isinstance(payload.get("ocr_results"), list):
        return JSONResponse({"error": "OCR page is invalid"}, status_code=502)
    return JSONResponse({"publication_id": publication_id, "page": page, **payload}, headers={
        "Cache-Control": "public, max-age=300",
    })

@router.post("/api/parse", status_code=202)
async def submit_parse(body: ParseRequest, request: Request):
    if not GITHUB_TOKEN:
        return JSONResponse({"error": "parsing is not configured"}, status_code=503)
    if len(json.dumps(body.model_dump(), ensure_ascii=False)) > PARSE_MAX_PAYLOAD_CHARACTERS:
        return JSONResponse({"error": "parse payload is too large"}, status_code=413)
    source = parse_queue_source(body.doc_id)
    if not source:
        return JSONResponse({"error": "parse item not found or already resolved"}, status_code=404)
    current_hash = content_hash(source)
    if body.base_content_hash and body.base_content_hash != current_hash:
        return JSONResponse({"error": "parse item changed since preview", "content_hash": current_hash}, status_code=409)
    articles, error = parse_core.validate_articles(body.articles, source, PARSE_QUEUE_PREFIX)
    if error:
        return JSONResponse({"error": error}, status_code=400)
    payload: dict[str, Any] = {
        "doc_id": body.doc_id,
        "archive_id": source.get("archive_id"),
        "publication_id": source.get("publication_id"),
        "kind": "parse",
        "body": {
            "articles": articles,
            "locator": {
                "title": (source.get("article") or {}).get("title"),
                "page_start": (source.get("article") or {}).get("page_start"),
                "page_end": (source.get("article") or {}).get("page_end"),
            },
            "source_files": source.get("source_files") or [],
            "description": body.description,
        },
    }
    request_id = hashlib.sha256(json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode()).hexdigest()
    payload["request_id"] = request_id
    if duplicate := duplicate_parse_response(request_id):
        return duplicate
    if error := parse_rate_error(request):
        with parse_dispatch_lock:
            parse_dispatches.pop(request_id, None)
        return error
    try:
        await dispatch_parse(payload)
    except Exception as exc:
        with parse_dispatch_lock:
            parse_dispatches.pop(request_id, None)
        release_parse_rate(request)
        return JSONResponse({"error": str(exc)}, status_code=502)
    with parse_dispatch_lock:
        parse_dispatches[request_id] = {"created_at": time.time(), "status": "accepted"}
    return {"ok": True, "status": "accepted", "request_id": request_id, "tasks": 1}