bid-doctor/server/src/bid_doctor/document/api.py

from __future__ import annotations

import json
from typing import Annotated, Any, Literal
from uuid import UUID

from fastapi import APIRouter, Body, Query, Request
from fastapi.responses import Response
from pydantic import BaseModel, Field
from sqlalchemy import select

from bid_doctor.core import keys
from bid_doctor.core.config import get_settings
from bid_doctor.core.crypto import open_object, sign_download, verify_download
from bid_doctor.core.db import TenantContext, tenant_tx
from bid_doctor.core.errors import AppError
from bid_doctor.core.storage import storage
from bid_doctor.document import service, uploads
from bid_doctor.document.models import DocumentPage, SourceDocument
from bid_doctor.identity.deps import Ctx

router = APIRouter(tags=["uploads", "documents"])


class UploadIn(BaseModel):
    filename: str = Field(max_length=255)
    size_bytes: int = Field(ge=1)
    sha256: str | None = Field(default=None, pattern=r"^[0-9a-f]{64}$")
    purpose: Literal["project_document", "qualification_file"]


class DocumentIn(BaseModel):
    upload_id: UUID
    role: Literal["tender", "addendum", "clarification", "bid_volume"]
    volume_name: str | None = Field(default=None, max_length=40)


class ParseIn(BaseModel):
    expected_charge: int | None = None


@router.post("/uploads", status_code=201)
async def create_upload(body: UploadIn, ctx: Ctx) -> dict[str, Any]:
    return await uploads.create(ctx, body.filename, body.size_bytes, body.purpose, body.sha256)


@router.get("/uploads/{upload_id}")
async def get_upload(upload_id: UUID, ctx: Ctx) -> dict[str, Any]:
    return await uploads.describe(ctx, upload_id)


@router.put("/uploads/{upload_id}/parts/{part_no}", status_code=204)
async def put_part(upload_id: UUID, part_no: int, request: Request, ctx: Ctx) -> None:
    await uploads.put_part(ctx, upload_id, part_no, await request.body())


@router.post("/uploads/{upload_id}/complete")
async def complete_upload(upload_id: UUID, ctx: Ctx) -> dict[str, Any]:
    return await uploads.complete(ctx, upload_id)


@router.delete("/uploads/{upload_id}", status_code=204)
async def abort_upload(upload_id: UUID, ctx: Ctx) -> None:
    await uploads.discard(ctx, upload_id)


@router.get("/projects/{project_id}/documents")
async def list_documents(project_id: UUID, ctx: Ctx) -> dict[str, Any]:
    async with tenant_tx(ctx) as s:
        docs = (
            (
                await s.execute(
                    select(SourceDocument)
                    .where(SourceDocument.project_id == project_id)
                    .order_by(SourceDocument.created_at)
                )
            )
            .scalars()
            .all()
        )
        return {"items": [await service.doc_to_api(s, d) for d in docs]}


@router.post("/projects/{project_id}/documents", status_code=201)
async def create_document(project_id: UUID, body: DocumentIn, ctx: Ctx) -> dict[str, Any]:
    async with tenant_tx(ctx) as s:
        doc = await service.create(s, ctx, project_id, body.upload_id, body.role, body.volume_name)
        return await service.doc_to_api(s, doc)


@router.get("/documents/{document_id}")
async def get_document(document_id: UUID, ctx: Ctx) -> dict[str, Any]:
    async with tenant_tx(ctx) as s:
        return await service.doc_to_api(s, await service.get(s, document_id))


@router.get("/documents/{document_id}/pages")
async def list_pages(document_id: UUID, ctx: Ctx) -> dict[str, Any]:
    async with tenant_tx(ctx) as s:
        doc = await service.get(s, document_id)
        pages = (
            (
                await s.execute(
                    select(DocumentPage).where(DocumentPage.document_id == doc.id).order_by(DocumentPage.page_no)
                )
            )
            .scalars()
            .all()
        )
    return {
        "normalized_sha256": doc.normalized_sha256.hex() if doc.normalized_sha256 else None,
        "items": [
            {
                "page_no": p.page_no,
                "width_pt": p.width_pt,
                "height_pt": p.height_pt,
                "rotation": p.rotation,
                "text_source": p.text_source,
            }
            for p in pages
        ],
    }


@router.get("/documents/{document_id}/pages/{page_no}/sentences")
async def list_sentences(document_id: UUID, page_no: int, ctx: Ctx) -> dict[str, Any]:
    async with tenant_tx(ctx) as s:
        doc = await service.get(s, document_id)
        dk = await keys.data_key(s, ctx, doc.project_id)
    if not doc.text_layer_key:
        raise AppError("document_not_ready")
    pages = json.loads(open_object(dk, storage().get(doc.text_layer_key), doc.text_layer_key))
    page = next((p for p in pages if p["page_no"] == page_no), None)
    if page is None:
        raise AppError("not_found")
    return {"items": page["sentences"]}


@router.get("/documents/{document_id}/file")
async def get_file_url(
    document_id: UUID, ctx: Ctx, variant: Literal["normalized", "original"] = Query()
) -> dict[str, Any]:
    async with tenant_tx(ctx) as s:
        doc = await service.get(s, document_id)
    if variant == "normalized" and not doc.normalized_key:
        raise AppError("document_not_ready")
    payload = f"{ctx.tenant_id}:{doc.id}:{variant}"
    sig, exp = sign_download(payload, get_settings().download_url_ttl_s)
    from datetime import UTC, datetime

    url = f"{get_settings().public_base_url}/api/v1/files/{ctx.tenant_id}/{doc.id}/{variant}?exp={exp}&sig={sig}"
    return {"url": url, "expires_at": datetime.fromtimestamp(exp, UTC)}


@router.get("/files/{tenant_id}/{document_id}/{variant}", include_in_schema=False)
async def download(
    tenant_id: UUID, document_id: UUID, variant: Literal["normalized", "original"], exp: int, sig: str
) -> Response:
    """签名下载（本地磁盘驱动用）。不带鉴权头，靠短时签名；S3 驱动时直接返回对象存储的预签名地址。"""
    if not verify_download(f"{tenant_id}:{document_id}:{variant}", exp, sig):
        raise AppError("not_found")
    ctx = TenantContext(tenant_id=tenant_id, user_id=UUID(int=0))
    async with tenant_tx(ctx) as s:
        doc = await service.get(s, document_id)
        dk = await keys.data_key(s, ctx, doc.project_id)
    key = doc.normalized_key if variant == "normalized" else doc.original_key
    assert key is not None
    media = "application/pdf" if variant == "normalized" else "application/octet-stream"
    return Response(
        open_object(dk, storage().get(key), key), media_type=media, headers={"Cache-Control": "private, max-age=300"}
    )


@router.post("/documents/{document_id}/parse-jobs", status_code=201)
async def start_parse(document_id: UUID, ctx: Ctx, body: Annotated[ParseIn | None, Body()] = None) -> dict[str, Any]:
    async with tenant_tx(ctx) as s:
        job = await service.start_parse(s, ctx, document_id, body.expected_charge if body else None)
        return service.job_to_api(job)


@router.get("/parse-jobs/{job_id}")
async def get_parse_job(job_id: UUID, ctx: Ctx) -> dict[str, Any]:
    async with tenant_tx(ctx) as s:
        return service.job_to_api(await service.get_job(s, job_id))


@router.post("/parse-jobs/{job_id}/cancel", status_code=202)
async def cancel_parse_job(job_id: UUID, ctx: Ctx) -> dict[str, Any]:
    async with tenant_tx(ctx) as s:
        return service.job_to_api(await service.request_cancel(s, job_id))