backend/app/documents/router.py

"""
Documents router — upload, analyze, chat, history.
All endpoints require JWT authentication.
"""
import os
from typing import Optional
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Query
from sqlalchemy.orm import Session
from pydantic import BaseModel

from app.database.database import get_db
from app.database.models import User, UploadedDocument, DocumentMessage, generate_uuid
from app.auth.router import get_current_user
from app.documents.service import extract_text, analyze_document, chat_with_document

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

# ─── Config ───────────────────────────────────────────────────────────────────
MAX_FILE_SIZE = 10 * 1024 * 1024  # 10 MB
ALLOWED_TYPES = {
    "application/pdf": "pdf",
    "application/vnd.openxmlformats-officedocument.wordprocessingml.document": "docx",
    "text/plain": "txt",
    "text/csv": "csv",
    "application/octet-stream": None,  # resolved by extension
}
ALLOWED_EXTENSIONS = {"pdf", "docx", "txt", "csv"}


def _resolve_file_type(filename: str, content_type: str) -> str:
    ext = filename.rsplit(".", 1)[-1].lower() if "." in filename else ""
    if ext in ALLOWED_EXTENSIONS:
        return ext
    mapped = ALLOWED_TYPES.get(content_type)
    if mapped:
        return mapped
    raise HTTPException(status_code=400, detail=f"Unsupported file type: {content_type}. Allowed: PDF, DOCX, TXT, CSV")


# ─── Upload ───────────────────────────────────────────────────────────────────
@router.post("/upload", status_code=201)
async def upload_document(
    file: UploadFile = File(...),
    language: str = Query(default="en"),
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    """Upload a document, extract text, and run AI analysis."""
    file_bytes = await file.read()

    if len(file_bytes) > MAX_FILE_SIZE:
        raise HTTPException(status_code=413, detail="File too large. Maximum size is 10 MB.")

    if not file_bytes:
        raise HTTPException(status_code=400, detail="Empty file.")

    file_type = _resolve_file_type(file.filename or "upload", file.content_type or "")

    # Extract text
    extracted_text = extract_text(file_bytes, file_type)
    if not extracted_text.strip():
        raise HTTPException(
            status_code=422,
            detail=(
                "No readable text could be extracted from this document. "
                "If this is a scanned PDF or image-only statement, please upload a text-based PDF, DOCX, TXT, or CSV file."
            ),
        )

    # AI analysis
    analysis = analyze_document(extracted_text, file.filename or "document", language)

    # Persist
    doc = UploadedDocument(
        id=generate_uuid(),
        user_id=current_user.id,
        filename=file.filename or "upload",
        file_type=file_type,
        file_size=len(file_bytes),
        extracted_text=extracted_text[:50000],  # cap stored text at 50k chars
        ai_summary=analysis["summary"],
        ai_insights=analysis["insights"] + (
            [f"⚠️ Suspicious: {s}" for s in analysis["suspicious"]] if analysis["suspicious"] else []
        ),
    )
    db.add(doc)
    db.commit()
    db.refresh(doc)

    return {
        "id": doc.id,
        "filename": doc.filename,
        "file_type": doc.file_type,
        "file_size": doc.file_size,
        "extracted_length": len(extracted_text),
        "summary": analysis["summary"],
        "insights": analysis["insights"],
        "suspicious": analysis["suspicious"],
        "created_at": doc.created_at.isoformat() if doc.created_at else None,
    }


# ─── Re-analyze ───────────────────────────────────────────────────────────────
@router.post("/analyze/{doc_id}")
def analyze_existing(
    doc_id: str,
    language: str = Query(default="en"),
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    doc = db.query(UploadedDocument).filter(
        UploadedDocument.id == doc_id,
        UploadedDocument.user_id == current_user.id,
    ).first()
    if not doc:
        raise HTTPException(status_code=404, detail="Document not found.")

    analysis = analyze_document(doc.extracted_text or "", doc.filename, language)
    doc.ai_summary = analysis["summary"]
    doc.ai_insights = analysis["insights"] + [f"⚠️ {s}" for s in analysis["suspicious"]]
    db.commit()

    return {
        "id": doc.id,
        "summary": analysis["summary"],
        "insights": analysis["insights"],
        "suspicious": analysis["suspicious"],
    }


# ─── Chat with document ───────────────────────────────────────────────────────
class DocChatRequest(BaseModel):
    question: str
    language: str = "en"


@router.post("/chat/{doc_id}")
def chat_document(
    doc_id: str,
    req: DocChatRequest,
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    doc = db.query(UploadedDocument).filter(
        UploadedDocument.id == doc_id,
        UploadedDocument.user_id == current_user.id,
    ).first()
    if not doc:
        raise HTTPException(status_code=404, detail="Document not found.")

    # Load conversation history for this document
    history_msgs = (
        db.query(DocumentMessage)
        .filter(
            DocumentMessage.document_id == doc_id,
            DocumentMessage.user_id == current_user.id,
        )
        .order_by(DocumentMessage.created_at.asc())
        .limit(20)
        .all()
    )
    history = [{"role": m.role, "content": m.content} for m in history_msgs]

    # Get AI response
    answer = chat_with_document(
        question=req.question,
        extracted_text=doc.extracted_text or "",
        filename=doc.filename,
        history=history,
        language=req.language,
    )

    # Persist both messages
    user_msg = DocumentMessage(
        id=generate_uuid(),
        user_id=current_user.id,
        document_id=doc_id,
        role="user",
        content=req.question,
        language=req.language,
    )
    ai_msg = DocumentMessage(
        id=generate_uuid(),
        user_id=current_user.id,
        document_id=doc_id,
        role="assistant",
        content=answer,
        language=req.language,
    )
    db.add(user_msg)
    db.add(ai_msg)
    db.commit()

    return {
        "question": req.question,
        "answer": answer,
        "document_id": doc_id,
        "language": req.language,
    }


# ─── History ──────────────────────────────────────────────────────────────────
@router.get("/history")
def get_document_history(
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    docs = (
        db.query(UploadedDocument)
        .filter(UploadedDocument.user_id == current_user.id)
        .order_by(UploadedDocument.created_at.desc())
        .limit(20)
        .all()
    )
    return {
        "documents": [
            {
                "id": d.id,
                "filename": d.filename,
                "file_type": d.file_type,
                "file_size": d.file_size,
                "summary": d.ai_summary,
                "insights": d.ai_insights or [],
                "extracted_length": len(d.extracted_text or ""),
                "created_at": d.created_at.isoformat() if d.created_at else None,
            }
            for d in docs
        ]
    }


# ─── Single document + its chat ───────────────────────────────────────────────
@router.get("/{doc_id}")
def get_document(
    doc_id: str,
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    doc = db.query(UploadedDocument).filter(
        UploadedDocument.id == doc_id,
        UploadedDocument.user_id == current_user.id,
    ).first()
    if not doc:
        raise HTTPException(status_code=404, detail="Document not found.")

    messages = (
        db.query(DocumentMessage)
        .filter(
            DocumentMessage.document_id == doc_id,
            DocumentMessage.user_id == current_user.id,
        )
        .order_by(DocumentMessage.created_at.asc())
        .all()
    )

    return {
        "id": doc.id,
        "filename": doc.filename,
        "file_type": doc.file_type,
        "file_size": doc.file_size,
        "summary": doc.ai_summary,
        "insights": doc.ai_insights or [],
        "extracted_length": len(doc.extracted_text or ""),
        "created_at": doc.created_at.isoformat() if doc.created_at else None,
        "messages": [
            {
                "id": m.id,
                "role": m.role,
                "content": m.content,
                "language": m.language,
                "created_at": m.created_at.isoformat() if m.created_at else None,
            }
            for m in messages
        ],
    }


# ─── Delete ───────────────────────────────────────────────────────────────────
@router.delete("/{doc_id}")
def delete_document(
    doc_id: str,
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    doc = db.query(UploadedDocument).filter(
        UploadedDocument.id == doc_id,
        UploadedDocument.user_id == current_user.id,
    ).first()
    if not doc:
        raise HTTPException(status_code=404, detail="Document not found.")
    db.delete(doc)
    db.commit()
    return {"message": "Document deleted."}