pdf/gateway/app/routers/edits.py

import json
from typing import Annotated, Literal

from fastapi import APIRouter, HTTPException, status
from pydantic import BaseModel, Field

from app.services import engine
from app.services.store import document_store

router = APIRouter(prefix="/documents/{document_id}/edits", tags=["edits"])
compat_router = APIRouter(tags=["edits"])


class TextOverlayData(BaseModel):
    text: str
    x: float
    y: float
    width: float
    height: float
    fontSize: float = Field(..., gt=0)
    fontFamily: str
    color: str


class RedactionData(BaseModel):
    x: float
    y: float
    width: float
    height: float
    fillColor: str = "#000000"


class ImageOverlayData(BaseModel):
    x: float
    y: float
    width: float
    height: float
    imageData: str


class HighlightQuadPoint(BaseModel):
    x1: float
    y1: float
    x2: float
    y2: float
    x3: float
    y3: float
    x4: float
    y4: float


class HighlightData(BaseModel):
    quadPoints: list[HighlightQuadPoint]
    color: str
    opacity: float = 1.0
    author: str
    content: str | None = None


class FreeTextData(BaseModel):
    x: float
    y: float
    width: float
    height: float
    text: str
    fontSize: float = Field(12.0, gt=0)
    color: str = "#000000"


class StickyNoteData(BaseModel):
    x: float
    y: float
    author: str
    content: str
    timestamp: str | None = None


class FreehandPoint(BaseModel):
    x: float
    y: float


class FreehandData(BaseModel):
    paths: list[list[FreehandPoint]]
    color: str
    thickness: float


class PageRotationData(BaseModel):
    rotation: Literal[0, 90, 180, 270]


class TextOverlayOperation(BaseModel):
    id: str
    type: Literal["text_overlay"]
    pageIndex: int = Field(..., ge=0)
    data: TextOverlayData


class RedactionOperation(BaseModel):
    id: str
    type: Literal["redaction"]
    pageIndex: int = Field(..., ge=0)
    data: RedactionData


class ImageOverlayOperation(BaseModel):
    id: str
    type: Literal["image_overlay"]
    pageIndex: int = Field(..., ge=0)
    data: ImageOverlayData


class HighlightOperation(BaseModel):
    id: str
    type: Literal["highlight"]
    pageIndex: int = Field(..., ge=0)
    data: HighlightData


class FreeTextOperation(BaseModel):
    id: str
    type: Literal["free_text"]
    pageIndex: int = Field(..., ge=0)
    data: FreeTextData


class CommentOperation(BaseModel):
    id: str
    type: Literal["comment"]
    pageIndex: int = Field(..., ge=0)
    data: StickyNoteData


class FreehandOperation(BaseModel):
    id: str
    type: Literal["freehand"]
    pageIndex: int = Field(..., ge=0)
    data: FreehandData


class PageRotationOperation(BaseModel):
    id: str
    type: Literal["page_rotation"]
    pageIndex: int = Field(..., ge=0)
    data: PageRotationData


class PageDeletionData(BaseModel):
    pass


class PageDeletionOperation(BaseModel):
    id: str
    type: Literal["page_deletion"]
    pageIndex: int
    data: PageDeletionData


class PageReorderData(BaseModel):
    destPageIndex: int


class PageReorderOperation(BaseModel):
    id: str
    type: Literal["page_reorder"]
    pageIndex: int
    data: PageReorderData


class EditTextData(BaseModel):
    # Target line/run bbox in PDF bottom-left page space (frontend does the Y-flip).
    x: float
    y: float
    width: float = Field(..., gt=0)
    height: float = Field(..., gt=0)
    newText: str
    originalText: str | None = None
    fontSize: float | None = Field(default=None, gt=0)
    color: str | None = None
    fallbackFont: str = "Helvetica"


class EditTextOperation(BaseModel):
    id: str
    type: Literal["edit_text"]
    pageIndex: int = Field(..., ge=0)
    data: EditTextData


class UpdateFieldData(BaseModel):
    value: str | bool
    annotationId: str | None = None


class UpdateFieldOperation(BaseModel):
    id: str
    type: Literal["update_field"]
    pageIndex: int
    data: UpdateFieldData


class DeleteAnnotationData(BaseModel):
    annotationId: str


class DeleteAnnotationOperation(BaseModel):
    id: str
    type: Literal["delete_annotation"]
    pageIndex: int = Field(..., ge=0)
    data: DeleteAnnotationData


class UpdateAnnotationData(BaseModel):
    annotationId: str
    x: float | None = None
    y: float | None = None
    width: float | None = None
    height: float | None = None
    color: str | None = None
    thickness: float | None = None
    text: str | None = None


class UpdateAnnotationOperation(BaseModel):
    id: str
    type: Literal["update_annotation"]
    pageIndex: int = Field(..., ge=0)
    data: UpdateAnnotationData


EditOperation = Annotated[
    TextOverlayOperation
    | RedactionOperation
    | ImageOverlayOperation
    | HighlightOperation
    | FreeTextOperation
    | CommentOperation
    | FreehandOperation
    | PageRotationOperation
    | PageDeletionOperation
    | PageReorderOperation
    | UpdateFieldOperation
    | DeleteAnnotationOperation
    | UpdateAnnotationOperation
    | EditTextOperation,
    Field(discriminator="type"),
]


class EditsRequest(BaseModel):
    version: Literal["1.0"]
    operations: list[EditOperation]


def apply_edits_impl(document_id: str, request: EditsRequest):
    if not engine.is_available():
        raise HTTPException(
            status_code=status.HTTP_501_NOT_IMPLEMENTED,
            detail="Engine bridge (bindings/python) not yet available.",
        )

    doc_info = document_store.get_document(document_id)
    if not doc_info:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document not found")

    created_temp_files = []
    try:
        pdfengine = engine.require()
        doc = doc_info["doc_instance"]

        req_dict = request.model_dump()
        for op in req_dict.get("operations", []):
            if op.get("type") == "image_overlay":
                img_data_str = op["data"].get("imageData", "")
                if img_data_str:
                    import base64
                    import io
                    import os
                    import tempfile

                    # pyrefly: ignore [missing-import]
                    from PIL import Image   

                    # Remove data URI header if present
                    if "," in img_data_str:
                        img_data_str = img_data_str.split(",", 1)[1]

                    raw_bytes = base64.b64decode(img_data_str)
                    img = Image.open(io.BytesIO(raw_bytes))
                    img_rgba = img.convert("RGBA")

                    # Convert RGBA to BGRA
                    r, g, b, a = img_rgba.split()
                    img_bgra = Image.merge("RGBA", (b, g, r, a))

                    bgra_bytes = img_bgra.tobytes()

                    # Create a temporary binary file to hold raw pixel data
                    fd, temp_path = tempfile.mkstemp(suffix=".bin", prefix="pdf_pixel_")
                    created_temp_files.append(temp_path)
                    try:
                        with os.fdopen(fd, "wb") as tmp:
                            tmp.write(bgra_bytes)
                    except Exception:
                        os.close(fd)
                        raise

                    op["data"]["pixelDataPath"] = temp_path
                    op["data"]["pixelWidth"] = img.width
                    op["data"]["pixelHeight"] = img.height

                    # Delete base64 strings to keep JSON payload tiny
                    if "imageData" in op["data"]:
                        del op["data"]["imageData"]

        edits_json = json.dumps(req_dict)
        doc.apply_edits(edits_json)

        # Redaction and in-place text rewrites mutate existing objects, which do
        # not round-trip cleanly through an incremental save — force a full save.
        full_save_types = {"redaction", "edit_text"}
        needs_full = any(op.get("type") in full_save_types for op in req_dict.get("operations", []))
        new_bytes = doc.save_full() if needs_full else doc.save_incremental()

        new_doc = pdfengine.PdfDocument.load_from_memory(new_bytes)

        new_info = document_store.add_document(
            filename=doc_info["filename"], bytes_data=new_bytes, doc_instance=new_doc
        )

        return {"success": True, "newDocumentId": new_info["id"]}
    except Exception as e:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)) from e
    finally:
        import contextlib
        import os

        for temp_path in created_temp_files:
            if os.path.exists(temp_path):
                with contextlib.suppress(Exception):
                    os.remove(temp_path)


@router.post("")
def apply_edits(document_id: str, request: EditsRequest):
    return apply_edits_impl(document_id, request)


@compat_router.post("/edits/{document_id}")
def apply_edits_compat(document_id: str, request: EditsRequest):
    return apply_edits_impl(document_id, request)