pdf-quiz-generator/backend/app/routers/documents.py

import os
import uuid
import shutil

from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, status
from sqlalchemy.orm import Session

from app.config import settings
from app.database import get_db
from app.models.pdf_document import PDFDocument
from app.models.section import Section
from app.models.user import User
from app.schemas.document import DocumentResponse, DocumentStatusResponse, SectionCreate, SectionResponse
from app.utils.auth import get_current_user, require_moderator
from app.services import vector_service

router = APIRouter()


@router.post("/upload", response_model=DocumentResponse)
def upload_document(
    file: UploadFile = File(...),
    db: Session = Depends(get_db),
    current_user: User = Depends(require_moderator),
):
    if not file.filename or not file.filename.lower().endswith(".pdf"):
        raise HTTPException(status_code=400, detail="Only PDF files are accepted")

    # Save file to disk streaming (handles large files)
    os.makedirs(settings.UPLOAD_DIR, exist_ok=True)
    safe_name = f"{uuid.uuid4()}_{os.path.basename(file.filename)}"
    file_path = os.path.join(settings.UPLOAD_DIR, safe_name)

    with open(file_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer, length=1024 * 1024)

    # Check file size
    file_size = os.path.getsize(file_path)
    if file_size > settings.MAX_UPLOAD_SIZE:
        os.remove(file_path)
        raise HTTPException(status_code=400, detail=f"File too large. Max size: {settings.MAX_UPLOAD_SIZE} bytes")

    # Create DB record
    doc = PDFDocument(
        user_id=current_user.id,
        filename=safe_name,
        original_filename=file.filename,
        status="processing",
    )
    db.add(doc)
    db.commit()
    db.refresh(doc)

    # Dispatch background processing
    try:
        from app.tasks.pdf_tasks import process_pdf
        process_pdf.delay(doc.id, file_path)
    except Exception:
        # If Celery/Redis not available, process synchronously
        from app.services import pdf_service
        try:
            total_pages = pdf_service.get_page_count(file_path)
            doc.total_pages = total_pages
            pages = pdf_service.extract_text_by_page(file_path)
            if pages:
                vector_service.store_pages(doc.id, pages)
                doc.status = "ready"
            else:
                doc.status = "error"
                doc.error_message = "No text could be extracted"
            db.commit()
        except Exception as e:
            doc.status = "error"
            doc.error_message = str(e)[:500]
            db.commit()

    db.refresh(doc)
    return doc


@router.get("/", response_model=list[DocumentResponse])
def list_documents(
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    if current_user.is_moderator:
        docs = db.query(PDFDocument).order_by(PDFDocument.uploaded_at.desc()).all()
    else:
        docs = db.query(PDFDocument).filter(PDFDocument.user_id == current_user.id).order_by(PDFDocument.uploaded_at.desc()).all()
    return docs


def _get_doc_or_404(document_id: int, current_user: User, db) -> PDFDocument:
    """Fetch document; moderators can access any, users only their own."""
    query = db.query(PDFDocument).filter(PDFDocument.id == document_id)
    if not current_user.is_moderator:
        query = query.filter(PDFDocument.user_id == current_user.id)
    doc = query.first()
    if not doc:
        raise HTTPException(status_code=404, detail="Document not found")
    return doc


@router.get("/{document_id}", response_model=DocumentResponse)
def get_document(
    document_id: int,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    return _get_doc_or_404(document_id, current_user, db)


@router.get("/{document_id}/status", response_model=DocumentStatusResponse)
def get_document_status(
    document_id: int,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    doc = _get_doc_or_404(document_id, current_user, db)
    return DocumentStatusResponse(
        id=doc.id,
        status=doc.status,
        total_pages=doc.total_pages,
        error_message=doc.error_message,
    )


@router.post("/{document_id}/sections", response_model=SectionResponse)
def create_section(
    document_id: int,
    section_data: SectionCreate,
    db: Session = Depends(get_db),
    current_user: User = Depends(require_moderator),
):
    doc = _get_doc_or_404(document_id, current_user, db)
    if not doc:
        raise HTTPException(status_code=404, detail="Document not found")
    if doc.status != "ready":
        raise HTTPException(status_code=400, detail="Document is not ready yet")
    if section_data.start_page < 1 or section_data.end_page < 1:
        raise HTTPException(status_code=400, detail="Page numbers must be positive")
    if section_data.start_page >= section_data.end_page:
        raise HTTPException(status_code=400, detail="start_page must be less than end_page")
    if doc.total_pages and section_data.end_page > doc.total_pages:
        raise HTTPException(status_code=400, detail=f"end_page exceeds document length ({doc.total_pages} pages)")

    section = Section(
        document_id=document_id,
        name=section_data.name,
        start_page=section_data.start_page,
        end_page=section_data.end_page,
    )
    db.add(section)
    db.commit()
    db.refresh(section)
    return section


@router.delete("/{document_id}", status_code=status.HTTP_204_NO_CONTENT)
def delete_document(
    document_id: int,
    db: Session = Depends(get_db),
    current_user: User = Depends(require_moderator),
):
    doc = _get_doc_or_404(document_id, current_user, db)

    # Delete file
    file_path = os.path.join(settings.UPLOAD_DIR, doc.filename)
    if os.path.exists(file_path):
        os.remove(file_path)

    # Delete ChromaDB collection
    vector_service.delete_collection(document_id)

    db.delete(doc)
    db.commit()