app.py

import json
import os
import shutil
import subprocess
import tempfile
import zipfile
from pathlib import Path

from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from starlette.background import BackgroundTask


ROOT = Path(__file__).parent
ALLOWED_ORIGINS = [origin.strip() for origin in os.getenv("ALLOWED_ORIGINS", "*").split(",") if origin.strip()]
EXTRACTION_TIMEOUT = int(os.getenv("EXTRACTION_TIMEOUT_SECONDS", "900"))

app = FastAPI(title="JEE PDF Extraction API", version="1.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOWED_ORIGINS,
    allow_credentials=False,
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["*"],
)


@app.get("/health")
def health():
    return {"ok": True, "service": "jee-pdf-extractor"}


@app.post("/extract/questions")
async def extract_questions(pdf: UploadFile = File(...), settings_json: str = Form("{}")):
    if not pdf.filename.lower().endswith(".pdf"):
        raise HTTPException(status_code=400, detail="Upload a PDF file.")

    with tempfile.TemporaryDirectory() as temp_dir_name:
        temp_dir = Path(temp_dir_name)
        pdf_path = temp_dir / "paper.pdf"
        output_dir = temp_dir / "questions"
        zip_path = temp_dir / "questions.zip"
        await save_upload(pdf, pdf_path)

        result = run_extractor([
            "questions",
            "--pdf",
            str(pdf_path),
            "--settings-json",
            settings_json,
            "--output-dir",
            str(output_dir),
        ])

        questions = []
        for question in result.get("questions", []):
            image_path = Path(question["imagePath"])
            image_filename = f"q-{question['questionNo']}.png"
            archive_path = output_dir / image_filename
            question.pop("imagePath", None)
            question["imageFilename"] = image_filename
            if image_path.resolve() != archive_path.resolve():
                shutil.copyfile(image_path, archive_path)
            questions.append(question)

        with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_STORED) as archive:
            archive.writestr("metadata.json", json.dumps({"questions": questions}))
            for image_path in sorted(output_dir.glob("q-*.png")):
                archive.write(image_path, image_path.name)

        final_fd, final_name = tempfile.mkstemp(suffix=".zip")
        os.close(final_fd)
        final_zip = Path(final_name)
        shutil.copyfile(zip_path, final_zip)

    return FileResponse(
        final_zip,
        media_type="application/zip",
        filename="questions.zip",
        background=BackgroundTask(final_zip.unlink, missing_ok=True),
    )


@app.post("/extract/key")
async def extract_key(pdf: UploadFile = File(...), questions_json: str = Form("[]")):
    if not pdf.filename.lower().endswith(".pdf"):
        raise HTTPException(status_code=400, detail="Upload a PDF file.")

    with tempfile.TemporaryDirectory() as temp_dir_name:
        temp_dir = Path(temp_dir_name)
        pdf_path = temp_dir / "answer-key.pdf"
        await save_upload(pdf, pdf_path)
        result = run_extractor([
            "key",
            "--pdf",
            str(pdf_path),
            "--questions-json",
            questions_json,
        ])
    return JSONResponse(result)


async def save_upload(upload: UploadFile, destination: Path):
    with destination.open("wb") as handle:
        while chunk := await upload.read(1024 * 1024):
            handle.write(chunk)
    if destination.stat().st_size <= 0:
        raise HTTPException(status_code=400, detail="Uploaded PDF was empty.")
    with destination.open("rb") as handle:
        if handle.read(5) != b"%PDF-":
            raise HTTPException(status_code=400, detail="Uploaded file is not a valid PDF.")


def run_extractor(args):
    command = ["python3", str(ROOT / "server_extract.py"), *args]
    try:
        completed = subprocess.run(command, check=False, capture_output=True, text=True, timeout=EXTRACTION_TIMEOUT)
    except subprocess.TimeoutExpired as error:
        raise HTTPException(status_code=504, detail=f"Extraction timed out after {EXTRACTION_TIMEOUT}s") from error

    try:
        parsed = json.loads(completed.stdout or "{}")
    except json.JSONDecodeError as error:
        raise HTTPException(status_code=500, detail=completed.stderr or "Extractor returned invalid JSON") from error

    if completed.returncode != 0 or parsed.get("error"):
        raise HTTPException(status_code=422, detail=parsed.get("error") or completed.stderr or "Extractor failed")
    return parsed