BadNote/server/badnote_server/routers/documents_router.py

"""Documents router for BadNote — upload, download, annotations, bookmarks."""

import json
from datetime import datetime, timezone
from uuid import uuid4

from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Form, status
from fastapi.responses import FileResponse

from ..auth import get_current_user
from ..database import get_db
from ..models import (
    AnnotationResponse,
    AnnotationUpdate,
    BookmarkCreate,
    BookmarkResponse,
    DocumentResponse,
)
from ..storage import delete_document, get_document_path, save_document

router = APIRouter()


def _row_to_doc(row) -> DocumentResponse:
    return DocumentResponse(
        id=row["id"],
        user_id=row["user_id"],
        filename=row["filename"],
        doc_type=row["doc_type"],
        page_count=row["page_count"],
        created_at=row["created_at"],
        updated_at=row["updated_at"],
    )


# ── Documents ───────────────────────────────────────────────────────────────


@router.post("/upload", response_model=DocumentResponse, status_code=status.HTTP_201_CREATED)
async def upload_document(
    file: UploadFile = File(...),
    doc_type: str = Form("pdf"),
    page_count: int = Form(0),
    user_id: str = Depends(get_current_user),
) -> DocumentResponse:
    """Upload a document file."""
    doc_id = str(uuid4())
    now = datetime.now(timezone.utc).isoformat()
    filename = file.filename or "document"
    file_bytes = await file.read()
    file_path = save_document(file_bytes, doc_id, filename)

    db = await get_db()
    await db.execute(
        """INSERT INTO documents (id, user_id, filename, doc_type, file_path, page_count, created_at, updated_at)
           VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
        (doc_id, user_id, filename, doc_type, file_path, page_count, now, now),
    )
    await db.commit()

    cursor = await db.execute("SELECT * FROM documents WHERE id = ?", (doc_id,))
    row = await cursor.fetchone()
    return _row_to_doc(row)


@router.get("", response_model=list[DocumentResponse])
async def list_documents(
    user_id: str = Depends(get_current_user),
) -> list[DocumentResponse]:
    """List all documents for the current user."""
    db = await get_db()
    cursor = await db.execute(
        "SELECT * FROM documents WHERE user_id = ? ORDER BY created_at DESC",
        (user_id,),
    )
    rows = await cursor.fetchall()
    return [_row_to_doc(r) for r in rows]


@router.get("/{doc_id}", response_model=DocumentResponse)
async def get_document(
    doc_id: str,
    user_id: str = Depends(get_current_user),
) -> DocumentResponse:
    """Get document metadata."""
    db = await get_db()
    cursor = await db.execute(
        "SELECT * FROM documents WHERE id = ? AND user_id = ?", (doc_id, user_id)
    )
    row = await cursor.fetchone()
    if row is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document not found")
    return _row_to_doc(row)


@router.get("/{doc_id}/file")
async def download_document(
    doc_id: str,
    user_id: str = Depends(get_current_user),
) -> FileResponse:
    """Stream document file download."""
    db = await get_db()
    cursor = await db.execute(
        "SELECT * FROM documents WHERE id = ? AND user_id = ?", (doc_id, user_id)
    )
    row = await cursor.fetchone()
    if row is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document not found")

    file_path = row["file_path"]
    return FileResponse(path=file_path, filename=row["filename"], media_type="application/octet-stream")


@router.delete("/{doc_id}", status_code=status.HTTP_200_OK)
async def delete_document_endpoint(
    doc_id: str,
    user_id: str = Depends(get_current_user),
) -> dict:
    """Delete document, its file, annotations, and bookmarks."""
    db = await get_db()
    cursor = await db.execute(
        "SELECT id FROM documents WHERE id = ? AND user_id = ?", (doc_id, user_id)
    )
    if await cursor.fetchone() is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document not found")

    delete_document(doc_id)
    await db.execute("DELETE FROM documents WHERE id = ?", (doc_id,))
    await db.commit()
    return {"deleted": doc_id}


# ── Annotations ─────────────────────────────────────────────────────────────


@router.get("/{doc_id}/annotations", response_model=list[AnnotationResponse])
async def list_annotations(
    doc_id: str,
    user_id: str = Depends(get_current_user),
) -> list[AnnotationResponse]:
    """Get all annotations for a document."""
    db = await get_db()
    # Verify document ownership
    cursor = await db.execute(
        "SELECT id FROM documents WHERE id = ? AND user_id = ?", (doc_id, user_id)
    )
    if await cursor.fetchone() is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document not found")

    cursor = await db.execute(
        "SELECT * FROM annotations WHERE document_id = ? ORDER BY page_number",
        (doc_id,),
    )
    rows = await cursor.fetchall()
    return [
        AnnotationResponse(
            id=r["id"],
            document_id=r["document_id"],
            page_number=r["page_number"],
            annotation_json=json.loads(r["annotation_json"]),
            created_at=r["created_at"],
            updated_at=r["updated_at"],
        )
        for r in rows
    ]


@router.put("/{doc_id}/annotations/{page}", response_model=AnnotationResponse, status_code=status.HTTP_200_OK)
async def update_annotation(
    doc_id: str,
    page: int,
    body: AnnotationUpdate,
    user_id: str = Depends(get_current_user),
) -> AnnotationResponse:
    """Create or replace annotations for a page."""
    db = await get_db()
    cursor = await db.execute(
        "SELECT id FROM documents WHERE id = ? AND user_id = ?", (doc_id, user_id)
    )
    if await cursor.fetchone() is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document not found")

    now = datetime.now(timezone.utc).isoformat()
    annotation_json = json.dumps(body.annotation_json)

    # Check if annotation for this page already exists
    cursor = await db.execute(
        "SELECT id FROM annotations WHERE document_id = ? AND page_number = ?",
        (doc_id, page),
    )
    existing = await cursor.fetchone()

    if existing:
        ann_id = existing["id"]
        await db.execute(
            "UPDATE annotations SET annotation_json = ?, updated_at = ? WHERE id = ?",
            (annotation_json, now, ann_id),
        )
    else:
        ann_id = str(uuid4())
        await db.execute(
            """INSERT INTO annotations (id, document_id, page_number, annotation_json, created_at, updated_at)
               VALUES (?, ?, ?, ?, ?, ?)""",
            (ann_id, doc_id, page, annotation_json, now, now),
        )
    await db.commit()

    cursor = await db.execute("SELECT * FROM annotations WHERE id = ?", (ann_id,))
    row = await cursor.fetchone()
    return AnnotationResponse(
        id=row["id"],
        document_id=row["document_id"],
        page_number=row["page_number"],
        annotation_json=json.loads(row["annotation_json"]),
        created_at=row["created_at"],
        updated_at=row["updated_at"],
    )


# ── Bookmarks ───────────────────────────────────────────────────────────────


@router.get("/{doc_id}/bookmarks", response_model=list[BookmarkResponse])
async def list_bookmarks(
    doc_id: str,
    user_id: str = Depends(get_current_user),
) -> list[BookmarkResponse]:
    """Get all bookmarks for a document."""
    db = await get_db()
    cursor = await db.execute(
        "SELECT id FROM documents WHERE id = ? AND user_id = ?", (doc_id, user_id)
    )
    if await cursor.fetchone() is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document not found")

    cursor = await db.execute(
        "SELECT * FROM bookmarks WHERE document_id = ? ORDER BY page_number",
        (doc_id,),
    )
    rows = await cursor.fetchall()
    return [
        BookmarkResponse(
            id=r["id"],
            document_id=r["document_id"],
            page_number=r["page_number"],
            label=r["label"],
            color=r["color"],
            created_at=r["created_at"],
        )
        for r in rows
    ]


@router.post("/{doc_id}/bookmarks", response_model=BookmarkResponse, status_code=status.HTTP_201_CREATED)
async def create_bookmark(
    doc_id: str,
    body: BookmarkCreate,
    user_id: str = Depends(get_current_user),
) -> BookmarkResponse:
    """Add a bookmark to a document."""
    db = await get_db()
    cursor = await db.execute(
        "SELECT id FROM documents WHERE id = ? AND user_id = ?", (doc_id, user_id)
    )
    if await cursor.fetchone() is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document not found")

    bookmark_id = str(uuid4())
    now = datetime.now(timezone.utc).isoformat()
    await db.execute(
        """INSERT INTO bookmarks (id, document_id, page_number, label, color, created_at)
           VALUES (?, ?, ?, ?, ?, ?)""",
        (bookmark_id, doc_id, body.page_number, body.label, body.color, now),
    )
    await db.commit()

    return BookmarkResponse(
        id=bookmark_id,
        document_id=doc_id,
        page_number=body.page_number,
        label=body.label,
        color=body.color,
        created_at=now,
    )


@router.delete("/{doc_id}/bookmarks/{bookmark_id}", status_code=status.HTTP_200_OK)
async def delete_bookmark(
    doc_id: str,
    bookmark_id: str,
    user_id: str = Depends(get_current_user),
) -> dict:
    """Delete a bookmark."""
    db = await get_db()
    cursor = await db.execute(
        "SELECT id FROM documents WHERE id = ? AND user_id = ?", (doc_id, user_id)
    )
    if await cursor.fetchone() is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document not found")

    cursor = await db.execute(
        "SELECT id FROM bookmarks WHERE id = ? AND document_id = ?",
        (bookmark_id, doc_id),
    )
    if await cursor.fetchone() is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Bookmark not found")

    await db.execute("DELETE FROM bookmarks WHERE id = ?", (bookmark_id,))
    await db.commit()
    return {"deleted": bookmark_id}