app/proofread_api.py

import asyncio
import hashlib
import ipaddress
import json
import os
import threading
import time
from typing import Any

from fastapi import APIRouter, Request
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field

from . import pipeline_dispatch, proofread_core, source_files
from .document_identity import content_hash


router = APIRouter()
GITHUB_TOKEN = os.environ.get("BHA_PROOFREAD_GITHUB_TOKEN", "")
PIPELINE_REPOSITORY = os.environ.get("BHA_PROOFREAD_PIPELINE_REPOSITORY", "anftm/pipeline")
RATE_LIMIT = 30
RATE_WINDOW_SECONDS = 60 * 60
MAX_PAYLOAD_CHARACTERS = 2_000_000
DISPATCH_MAX_CHARACTERS = 60_000
DISPATCH_MAX_PROPERTIES = 10
BLOB_CHUNK_CHARACTERS = 60_000
CHANGE_PREVIEW_BUDGET_CHARACTERS = 48_000
TAG_TYPES = {"文稿大类", "文稿类型", "地点", "人物", "出版方/发行方", "主题/事件", "记录", "审核", "翻译", "翻印/传抄"}
PART_TYPES = {
    "appellation", "title", "authors", "place", "subtitle", "subtitle2", "subtitle3",
    "subtitle4", "subtitle5", "subdate", "paragraph", "quotation", "signature", "image",
    "image_description",
}
AUTO_MERGE_MAX_DIFF_COST = 500
AUTO_MERGE_MAX_NET_PARAGRAPH_DELTA = 3
AUTO_MERGE_POLICY = {
    "max_diff_cost": AUTO_MERGE_MAX_DIFF_COST,
    "max_net_paragraph_delta": AUTO_MERGE_MAX_NET_PARAGRAPH_DELTA,
}
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
proofread_requests: dict[str, list[float]] = {}
proofread_rate_lock = threading.Lock()
proofread_dispatches: dict[str, dict[str, Any]] = {}
proofread_dispatch_lock = threading.Lock()


class ProofreadRequest(BaseModel):
    doc_id: str = Field(min_length=1, max_length=512)
    patch: dict[str, Any] | None = None
    metadata: dict[str, Any] | None = None
    description: str = Field(default="", max_length=2000)
    base_content_hash: str | None = Field(default=None, min_length=64, max_length=64, pattern=r"[0-9a-f]{64}")


def proofread_rate_error(request: Request) -> JSONResponse | None:
    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()]
    forwarded_client = ""
    if candidates and len(candidates) > PROXY_HOPS:
        candidate = candidates[-PROXY_HOPS - 1]
        try:
            if ipaddress.ip_address(candidate):
                forwarded_client = candidate
        except ValueError:
            pass
    key = forwarded_client or peer
    now = time.monotonic()
    with proofread_rate_lock:
        recent = [value for value in proofread_requests.get(key, []) if now - value < RATE_WINDOW_SECONDS]
        if len(recent) >= RATE_LIMIT:
            return JSONResponse({"error": "proofread rate limit exceeded"}, status_code=429)
        recent.append(now)
        proofread_requests[key] = recent
        if len(proofread_requests) > 10000:
            expired = []
            for identity, values in proofread_requests.items():
                values[:] = [value for value in values if now - value < RATE_WINDOW_SECONDS]
                if not values:
                    expired.append(identity)
            for identity in expired:
                proofread_requests.pop(identity, None)
            while len(proofread_requests) > 10000:
                oldest = min(proofread_requests, key=lambda identity: min(proofread_requests[identity]))
                proofread_requests.pop(oldest, None)
    try:
        request.state.proofread_rate_key = key
        request.state.proofread_rate_timestamp = now
    except Exception:
        pass
    return None

def release_proofread_rate(request: Request) -> None:
    try:
        key = request.state.proofread_rate_key
        timestamp = request.state.proofread_rate_timestamp
    except Exception:
        return
    with proofread_rate_lock:
        values = proofread_requests.get(key, [])
        try:
            values.remove(timestamp)
        except ValueError:
            return
        if values:
            proofread_requests[key] = values
        else:
            proofread_requests.pop(key, None)

def clean_string_list(value: Any, field: str, limit: int) -> list[str]:
    if not isinstance(value, list) or len(value) > limit:
        raise ValueError(f"{field} must be a list with at most {limit} entries")
    if any(not isinstance(item, str) for item in value):
        raise ValueError(f"{field} entries must be strings")
    cleaned = [item.strip() for item in value]
    if any(not item or len(item) > 300 for item in cleaned):
        raise ValueError(f"{field} contains an invalid entry")
    return cleaned

def clean_metadata(value: dict[str, Any]) -> dict[str, Any]:
    allowed_groups = {"article", "source"}
    if set(value) - allowed_groups:
        raise ValueError("metadata contains unsupported groups")
    result: dict[str, Any] = {}
    if "article" in value:
        article = value["article"]
        if not isinstance(article, dict) or set(article) - {"title", "authors", "dates", "tags"}:
            raise ValueError("article metadata contains unsupported fields")
        cleaned: dict[str, Any] = {}
        if "title" in article:
            if not isinstance(article["title"], str):
                raise ValueError("title is invalid")
            title = article["title"].strip()
            if not title or len(title) > 500:
                raise ValueError("title is invalid")
            cleaned["title"] = title
        if "authors" in article:
            cleaned["authors"] = clean_string_list(article["authors"], "authors", 50)
        if "dates" in article:
            dates = article["dates"]
            if not isinstance(dates, list) or len(dates) > 50:
                raise ValueError("dates must be a list with at most 50 entries")
            for date in dates:
                if not isinstance(date, dict) or set(date) - {"year", "month", "day"}:
                    raise ValueError("date contains unsupported fields")
                for key, minimum, maximum in (("year", -9999, 9999), ("month", 1, 12), ("day", 1, 31)):
                    if key in date and (type(date[key]) is not int or not minimum <= date[key] <= maximum):
                        raise ValueError(f"date {key} is invalid")
                if not any(key in date for key in ("year", "month", "day")):
                    raise ValueError("date must contain a value")
            cleaned["dates"] = [dict(date) for date in dates]
        if "tags" in article:
            tags = article["tags"]
            if not isinstance(tags, list) or len(tags) > 100:
                raise ValueError("tags must be a list with at most 100 entries")
            for tag in tags:
                if not isinstance(tag, dict) or set(tag) != {"name", "type"}:
                    raise ValueError("each tag requires name and type")
                if not isinstance(tag["name"], str) or not isinstance(tag["type"], str):
                    raise ValueError("tag name and type must be strings")
                if not tag["name"].strip() or not tag["type"].strip():
                    raise ValueError("tag name and type cannot be empty")
                if tag["type"].strip() not in TAG_TYPES:
                    raise ValueError("tag type is invalid")
            cleaned["tags"] = [{"name": tag["name"].strip(), "type": tag["type"].strip()} for tag in tags]
        result["article"] = cleaned
    if "source" in value:
        source = value["source"]
        if not isinstance(source, dict) or set(source) - {"name", "author", "type", "files"}:
            raise ValueError("source metadata contains unsupported fields")
        cleaned = {}
        for key in ("name", "author", "type"):
            if key in source:
                if not isinstance(source[key], str):
                    raise ValueError(f"source {key} is invalid")
                text = source[key].strip()
                if len(text) > 500 or (key in {"name", "type"} and not text):
                    raise ValueError(f"source {key} is invalid")
                cleaned[key] = text
        if "files" in source:
            files = clean_string_list(source["files"], "source files", source_files.MAX_SOURCE_FILES)
            if any(source_files.validate_source_url(url) for url in files):
                raise ValueError("source files contain an invalid URL")
            cleaned["files"] = files
        result["source"] = cleaned
    if not any(result.values()):
        raise ValueError("metadata contains no changes")
    return result

def metadata_is_safe_to_merge(metadata: dict[str, Any] | None, _article: dict[str, Any]) -> bool:
    return metadata is None

async def dispatch_proofread(payload: dict[str, Any]) -> None:
    await pipeline_dispatch.dispatch(
        payload, "submit-proofread", GITHUB_TOKEN, PIPELINE_REPOSITORY,
        DISPATCH_MAX_CHARACTERS, DISPATCH_MAX_PROPERTIES,
        BLOB_CHUNK_CHARACTERS,
    )

def proofread_request_id(doc_id: str, patch: dict[str, Any] | None, metadata: dict[str, Any] | None, description: str, base_content_hash: str) -> str:
    value = {
        "doc_id": doc_id,
        "patch": patch,
        "metadata": metadata,
        "description": description,
        "base_content_hash": base_content_hash,
    }
    return hashlib.sha256(json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode()).hexdigest()

def duplicate_proofread_response(request_id: str) -> JSONResponse | None:
    now = time.time()
    with proofread_dispatch_lock:
        expired = [key for key, item in proofread_dispatches.items() if now - float(item["created_at"]) > 86400]
        for key in expired:
            proofread_dispatches.pop(key, None)
        existing = proofread_dispatches.get(request_id)
        if existing is None:
            if len(proofread_dispatches) >= 1000:
                oldest = min(proofread_dispatches, key=lambda key: proofread_dispatches[key]["created_at"])
                proofread_dispatches.pop(oldest, None)
            proofread_dispatches[request_id] = {"created_at": now, "status": "processing"}
            return None
        return JSONResponse({
            "ok": True,
            "status": "duplicate",
            "original_status": existing["status"],
            "request_id": request_id,
            "tasks": 1,
            "auto_merge": existing.get("auto_merge", False),
        }, status_code=200)

@router.post("/api/proofread", status_code=202)
async def submit_proofread(body: ProofreadRequest, request: Request):
    if not GITHUB_TOKEN:
        return JSONResponse({"error": "proofreading is not configured"}, status_code=503)
    if len(json.dumps(body.model_dump(), ensure_ascii=False)) > MAX_PAYLOAD_CHARACTERS:
        return JSONResponse({"error": "proofread payload is too large"}, status_code=413)
    source = await asyncio.to_thread(source_files.get_source_doc, body.doc_id)
    if not source:
        return JSONResponse({"error": "document not found"}, status_code=404)
    current_content_hash = content_hash(source)
    if body.base_content_hash and body.base_content_hash != current_content_hash:
        return JSONResponse({"error": "document changed since preview", "content_hash": current_content_hash}, status_code=409)
    payload: dict[str, Any] = {
        "doc_id": body.doc_id,
        "archive_id": source.get("archive_id"),
        "article_id": source.get("article_id"),
        "publication_id": source.get("publication_id"),
        "title": f"校订 {source.get('title') or source.get('article_id')}",
        "kind": "proofread",
        "body": {"description": body.description},
    }
    body_fields = payload["body"]
    if body.patch:
        try:
            normalized_patch = proofread_core.validate_proofread_patch(
                body.patch, source.get("article") or {}, PART_TYPES,
            )
            body_fields["patch"] = normalized_patch
        except (TypeError, ValueError) as exc:
            return JSONResponse({"error": str(exc)}, status_code=400)
    if body.metadata:
        try:
            metadata = clean_metadata(body.metadata)
        except ValueError as exc:
            return JSONResponse({"error": str(exc)}, status_code=400)
        article = source.get("article") or {}
        body_fields["metadata"] = metadata
        body_fields["locator"] = {
            "title": article.get("title"),
            "authors": article.get("authors") or [],
            "dates": article.get("dates") or [],
            "is_range_date": bool(article.get("is_range_date")),
            "page_start": article.get("page_start"),
            "page_end": article.get("page_end"),
        }
    if "patch" not in body_fields and "metadata" not in body_fields:
        return JSONResponse({"error": "no proofreading changes supplied"}, status_code=400)
    request_id = proofread_request_id(body.doc_id, body_fields.get("patch"), body_fields.get("metadata"), body.description, current_content_hash)
    payload["request_id"] = request_id
    payload["base_content_hash"] = current_content_hash
    payload["changed"] = proofread_core.proofread_change_preview(
        body_fields.get("patch"), source.get("article") or {}, body_fields.get("metadata"), source,
    )
    if not payload["changed"]:
        return JSONResponse({"error": "no effective proofreading changes supplied"}, status_code=400)
    if duplicate := duplicate_proofread_response(request_id):
        return duplicate
    if len(json.dumps(payload, ensure_ascii=False)) > CHANGE_PREVIEW_BUDGET_CHARACTERS:
        payload["changed"] = proofread_core.proofread_change_preview(
            body_fields.get("patch"), source.get("article") or {}, body_fields.get("metadata"), source,
            clip=2000,
        )
    if error := proofread_rate_error(request):
        with proofread_dispatch_lock:
            proofread_dispatches.pop(request_id, None)
        return error
    if proofread_core.patch_is_safe_to_merge(
        body_fields.get("patch"), source.get("article") or {},
        AUTO_MERGE_MAX_DIFF_COST, AUTO_MERGE_MAX_NET_PARAGRAPH_DELTA,
    ) and metadata_is_safe_to_merge(
        body_fields.get("metadata"), source.get("article") or {},
    ):
        payload["auto_merge"] = True
        payload["auto_merge_policy"] = AUTO_MERGE_POLICY
    try:
        await dispatch_proofread(payload)
    except Exception as exc:
        with proofread_dispatch_lock:
            proofread_dispatches.pop(request_id, None)
        release_proofread_rate(request)
        return JSONResponse({"error": str(exc)}, status_code=502)
    response = {"ok": True, "status": "accepted", "request_id": request_id, "tasks": 1, "auto_merge": payload.get("auto_merge") is True}
    with proofread_dispatch_lock:
        proofread_dispatches[request_id] = {"created_at": time.time(), "status": "accepted", "auto_merge": response["auto_merge"]}
    return response