pdf/gateway/app/routers/documents/crud.py

import re
import httpx
from urllib.parse import urlparse
from fastapi import APIRouter, File, Form, HTTPException, UploadFile, status
from pydantic import BaseModel

from app.schemas.document import (
    DocumentInfoResponse,
    PageInfoResponse,
    PermissionsResponse,
    ProtectDocumentRequest,
)
from app.services import engine
from app.services.store import document_store

router = APIRouter(prefix="/documents", tags=["documents"])

import io
from PIL import Image


class RemoteImportRequest(BaseModel):
    stream_url: str
    upload_url: str
    auth_token: str | None = None
    resource_id: str | None = None


@router.post("/import-remote", response_model=DocumentInfoResponse, status_code=status.HTTP_201_CREATED)
async def import_remote_document(req: RemoteImportRequest) -> DocumentInfoResponse:
    if not engine.is_available():
        raise HTTPException(
            status_code=status.HTTP_501_NOT_IMPLEMENTED,
            detail="Engine bridge (bindings/python) not yet available.",
        )

    stream_parsed = urlparse(req.stream_url)
    upload_parsed = urlparse(req.upload_url)
    if stream_parsed.scheme not in ("http", "https") or upload_parsed.scheme not in ("http", "https"):
        raise HTTPException(status_code=400, detail="Invalid URL scheme. Only http/https are allowed.")

    headers = {}
    if req.auth_token:
        headers["Authorization"] = f"Bearer {req.auth_token}"

    MAX_SIZE = 50 * 1024 * 1024  # 50 MB limit
    async with httpx.AsyncClient(timeout=60.0, follow_redirects=True) as client:
        try:
            async with client.stream("GET", req.stream_url, headers=headers) as res:
                if res.status_code != 200:
                    await res.aread()
                    raise HTTPException(
                        status_code=res.status_code,
                        detail=f"Failed to stream remote document: {res.text[:200]}",
                    )
                
                content_length = res.headers.get("Content-Length")
                if content_length and int(content_length) > MAX_SIZE:
                    raise HTTPException(status_code=400, detail="Document exceeds maximum allowed size (50MB).")
                
                bytes_data_arr = bytearray()
                async for chunk in res.aiter_bytes():
                    bytes_data_arr.extend(chunk)
                    if len(bytes_data_arr) > MAX_SIZE:
                        raise HTTPException(status_code=400, detail="Document exceeds maximum allowed size (50MB).")
                
                bytes_data = bytes(bytes_data_arr)
        except Exception as err:
            if isinstance(err, HTTPException):
                raise err
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"Failed to connect to remote stream_url: {err!s}",
            )

        cd = res.headers.get("Content-Disposition", "")
        match = re.search(r'filename="?([^";]+)"?', cd)
        filename = match.group(1) if match else f"remote_document_{req.resource_id or 'file'}.pdf"

    try:
        pdfengine = engine.require()
        doc = pdfengine.PdfDocument.load_from_memory(bytes_data, "")
        info = document_store.add_document(filename, bytes_data, doc)
        info["remote_context"] = {
            "stream_url": req.stream_url,
            "upload_url": req.upload_url,
            "auth_token": req.auth_token,
            "resource_id": req.resource_id,
        }
        return make_document_response(info)
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Failed to load remote PDF into engine: {e!s}",
        )


def make_document_response(d: dict) -> DocumentInfoResponse:
    pages_list = []
    if "doc_instance" in d:
        doc = d["doc_instance"]
        for i in range(doc.page_count):
            try:
                page = doc.get_page(i)
                pages_list.append(PageInfoResponse(index=i, width=page.width, height=page.height))
            except Exception:
                pass
    perms = d.get("permissions")
    return DocumentInfoResponse(
        id=d["id"],
        filename=d["filename"],
        sizeBytes=d["sizeBytes"],
        totalPages=d["totalPages"],
        pageWidth=d.get("pageWidth", 612.0),
        pageHeight=d.get("pageHeight", 792.0),
        uploadedAt=d["uploadedAt"],
        status=d["status"],
        pages=pages_list,
        permissions=PermissionsResponse(**perms) if perms else PermissionsResponse(),
    )


def _convert_image_to_pdf_bytes(image_bytes: bytes) -> bytes:
    img = Image.open(io.BytesIO(image_bytes))
    if img.mode in ("RGBA", "LA", "P"):
        background = Image.new("RGB", img.size, (255, 255, 255))
        if img.mode == "P":
            img = img.convert("RGBA")
        if "A" in img.mode:
            background.paste(img, mask=img.split()[-1])
        else:
            background.paste(img)
        img = background
    elif img.mode != "RGB":
        img = img.convert("RGB")

    pdf_buffer = io.BytesIO()
    img.save(pdf_buffer, format="PDF")
    return pdf_buffer.getvalue()


@router.post("", response_model=DocumentInfoResponse, status_code=status.HTTP_201_CREATED)
async def upload_document(file: UploadFile = File(...), password: str = "") -> DocumentInfoResponse:
    if not engine.is_available():
        raise HTTPException(
            status_code=status.HTTP_501_NOT_IMPLEMENTED,
            detail="Engine bridge (bindings/python) not yet available.",
        )

    bytes_data = await file.read()
    filename = file.filename or "uploaded_file.pdf"
    filename_lower = filename.lower()
    image_exts = [".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tiff", ".gif"]
    
    is_image = any(filename_lower.endswith(ext) for ext in image_exts) or \
               bytes_data.startswith(b"\x89PNG") or \
               bytes_data.startswith(b"\xff\xd8") or \
               bytes_data.startswith(b"RIFF") or \
               bytes_data.startswith(b"BM")

    if is_image:
        try:
            bytes_data = _convert_image_to_pdf_bytes(bytes_data)
            if not filename_lower.endswith(".pdf"):
                filename = f"{filename}.pdf"
        except Exception as e:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"Failed to convert image to PDF: {e!s}"
            )

    try:
        pdfengine = engine.require()
        doc = pdfengine.PdfDocument.load_from_memory(bytes_data, password)
        info = document_store.add_document(filename, bytes_data, doc, password=password)
        return make_document_response(info)
    except ValueError as e:
        detail = str(e)
        if "Password required" in detail:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED, detail="Password required"
            )
        elif "Invalid password" in detail:
            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid password")
        else:
            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=detail)
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST, detail=f"Failed to load PDF: {e!s}"
        )


@router.get("", response_model=list[DocumentInfoResponse])
def list_documents() -> list[DocumentInfoResponse]:
    if not engine.is_available():
        raise HTTPException(
            status_code=status.HTTP_501_NOT_IMPLEMENTED,
            detail="Engine bridge (bindings/python) not yet available.",
        )

    docs = document_store.list_documents()
    return [make_document_response(d) for d in docs]


@router.get("/{document_id}", response_model=DocumentInfoResponse)
def get_document(document_id: str) -> DocumentInfoResponse:
    if not engine.is_available():
        raise HTTPException(
            status_code=status.HTTP_501_NOT_IMPLEMENTED,
            detail="Engine bridge (bindings/python) not yet available.",
        )

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

    return make_document_response(d)


@router.delete("/{document_id}")
def delete_document(document_id: str):
    if not engine.is_available():
        raise HTTPException(
            status_code=status.HTTP_501_NOT_IMPLEMENTED,
            detail="Engine bridge (bindings/python) not yet available.",
        )

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

    return {"success": True}


@router.post("/{document_id}/protect", response_model=DocumentInfoResponse)
async def protect_document(
    document_id: str,
    req: ProtectDocumentRequest,
) -> DocumentInfoResponse:
    if not engine.is_available():
        raise HTTPException(
            status_code=status.HTTP_501_NOT_IMPLEMENTED,
            detail="Engine bridge (bindings/python) not yet available.",
        )

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

    if not req.userPassword or not req.userPassword.strip():
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="User password cannot be empty.",
        )

    if req.confirmPassword is not None and req.userPassword != req.confirmPassword:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="User password and confirmation password do not match.",
        )

    owner_pwd = req.ownerPassword if req.ownerPassword else req.userPassword

    try:
        doc_instance = d["doc_instance"]
        unencrypted_bytes = doc_instance.save_full_for_export()
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to extract document bytes for encryption: {e!s}",
        )

    try:
        pdfengine = engine.require()
        protected_bytes = pdfengine.protect_pdf(
            unencrypted_bytes,
            req.userPassword,
            owner_pwd,
            req.permissions.model_dump(),
        )
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to encrypt document: {e!s}",
        )

    try:
        protected_doc = pdfengine.PdfDocument.load_from_memory(protected_bytes, req.userPassword)
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to verify encrypted PDF integrity: {e!s}",
        )

    updated_info = document_store.update_document_bytes(
        doc_id=document_id,
        bytes_data=protected_bytes,
        doc_instance=protected_doc,
        password=req.userPassword,
    )
    if not updated_info:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document not found during update")

    return make_document_response(updated_info)


@router.post("/{document_id}/unlock", response_model=DocumentInfoResponse)
async def unlock_document(document_id: str) -> DocumentInfoResponse:
    if not engine.is_available():
        raise HTTPException(
            status_code=status.HTTP_501_NOT_IMPLEMENTED,
            detail="Engine bridge (bindings/python) not yet available.",
        )

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

    perms = d.get("permissions") or {}
    if not perms.get("isEncrypted", False):
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Document is not password protected.",
        )

    if perms.get("ownerUnlocked") is False and perms.get("canModify") is False:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Owner password authorization required to remove security restrictions.",
        )

    try:
        pdfengine = engine.require()
        pwd = d.get("password", "")
        unencrypted_bytes = pdfengine.unlock_pdf(d["bytes_data"], pwd)
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to generate unencrypted PDF stream: {e!s}",
        )

    try:
        unlocked_doc = pdfengine.PdfDocument.load_from_memory(unencrypted_bytes, "")
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to validate unlocked PDF integrity: {e!s}",
        )

    from app.services.store import extract_permissions
    unlocked_perms = extract_permissions(unlocked_doc) or {}
    if unlocked_perms.get("isEncrypted", True):
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="PDF validation failed: Output document retains encryption dictionary.",
        )

    if unlocked_doc.page_count != d["totalPages"]:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="PDF validation failed: Page count mismatch after unlocking.",
        )

    updated_info = document_store.update_document_bytes(
        doc_id=document_id,
        bytes_data=unencrypted_bytes,
        doc_instance=unlocked_doc,
        permissions=unlocked_perms,
        password="",
    )
    if not updated_info:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document not found during update")

    return make_document_response(updated_info)


import json
from app.services.pdf_merge import merge_pdf_files


@router.post("/merge", response_model=DocumentInfoResponse, status_code=status.HTTP_201_CREATED)
async def merge_documents(
    files: list[UploadFile] = File(...),
    manifest: str = Form(default="[]"),
    output_filename: str = Form(default="merged.pdf"),
    parent_document_id: str = Form(default=""),
) -> DocumentInfoResponse:
    if not engine.is_available():
        raise HTTPException(
            status_code=status.HTTP_501_NOT_IMPLEMENTED,
            detail="Engine bridge (bindings/python) not yet available.",
        )

    if not files or len(files) < 1:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="At least one PDF file must be provided for merging.",
        )

    files_bytes: list[bytes] = []
    for file in files:
        data = await file.read()
        fn_lower = (file.filename or "").lower()
        image_exts = [".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tiff", ".gif"]
        is_img = any(fn_lower.endswith(ext) for ext in image_exts) or \
                 data.startswith(b"\x89PNG") or \
                 data.startswith(b"\xff\xd8") or \
                 data.startswith(b"RIFF") or \
                 data.startswith(b"BM")
        if is_img:
            try:
                data = _convert_image_to_pdf_bytes(data)
            except Exception as e:
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail=f"Failed to convert image {file.filename} to PDF: {e!s}",
                )
        files_bytes.append(data)

    parsed_manifest: list[dict] = []
    if manifest and manifest.strip():
        try:
            parsed = json.loads(manifest)
            if isinstance(parsed, list):
                parsed_manifest = parsed
        except Exception:
            pass

    if not parsed_manifest:
        parsed_manifest = [{"fileIndex": idx, "pages": "all"} for idx in range(len(files))]

    try:
        merged_bytes = merge_pdf_files(files_bytes, parsed_manifest)
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Failed to merge PDF files: {e!s}",
        )

    if not output_filename or not output_filename.lower().endswith(".pdf"):
        output_filename = f"{output_filename or 'merged'}.pdf"

    try:
        pdfengine = engine.require()
        doc = pdfengine.PdfDocument.load_from_memory(merged_bytes, "")
        info = document_store.add_document(output_filename, merged_bytes, doc)
        
        if parent_document_id:
            parent_doc = document_store.get_document(parent_document_id)
            if parent_doc and "remote_context" in parent_doc:
                info["remote_context"] = parent_doc["remote_context"]
                
        return make_document_response(info)
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Failed to load merged PDF into engine: {e!s}",
        )