main.py

import os
import io
import re
import json
import tempfile
from typing import Optional

import uvicorn
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel

import google.generativeai as genai

# ── Document parsers ────────────────────────────────────────────────────────
import PyPDF2
import docx

# ── App setup ───────────────────────────────────────────────────────────────
app = FastAPI(
    title="DocAI – Document Analyzer",
    description="Upload PDF, DOCX, or TXT files and get AI-generated title, summary, and keywords via Gemini 2.5 Flash.",
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── Gemini client ────────────────────────────────────────────────────────────
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
if GEMINI_API_KEY:
    genai.configure(api_key=GEMINI_API_KEY)

GEMINI_MODEL = "gemini-2.5-flash"

# ── Pydantic models ──────────────────────────────────────────────────────────
class AnalysisResult(BaseModel):
    title: str
    summary: str
    keywords: list[str]
    word_count: int
    char_count: int
    file_name: str
    file_type: str


# ── Helpers ──────────────────────────────────────────────────────────────────
ALLOWED_TYPES = {
    "application/pdf": "pdf",
    "application/vnd.openxmlformats-officedocument.wordprocessingml.document": "docx",
    "application/msword": "doc",
    "text/plain": "txt",
}

MAX_FILE_SIZE_MB = 10
MAX_FILE_SIZE_BYTES = MAX_FILE_SIZE_MB * 1024 * 1024


def extract_text_from_pdf(file_bytes: bytes) -> str:
    reader = PyPDF2.PdfReader(io.BytesIO(file_bytes))
    pages = []
    for page in reader.pages:
        text = page.extract_text()
        if text:
            pages.append(text.strip())
    return "\n\n".join(pages)


def extract_text_from_docx(file_bytes: bytes) -> str:
    doc = docx.Document(io.BytesIO(file_bytes))
    paragraphs = [p.text for p in doc.paragraphs if p.text.strip()]
    return "\n\n".join(paragraphs)


def extract_text_from_txt(file_bytes: bytes) -> str:
    for enc in ("utf-8", "latin-1", "cp1252"):
        try:
            return file_bytes.decode(enc)
        except UnicodeDecodeError:
            continue
    raise ValueError("Cannot decode text file with supported encodings.")


def extract_text(file_bytes: bytes, file_type: str) -> str:
    if file_type == "pdf":
        return extract_text_from_pdf(file_bytes)
    elif file_type in ("docx", "doc"):
        return extract_text_from_docx(file_bytes)
    elif file_type == "txt":
        return extract_text_from_txt(file_bytes)
    raise ValueError(f"Unsupported file type: {file_type}")


def truncate_text(text: str, max_chars: int = 30_000) -> str:
    """Gemini has a large context window; we still cap to keep costs low."""
    if len(text) <= max_chars:
        return text
    return text[:max_chars] + "\n\n[... document truncated for analysis ...]"


def build_prompt(text: str) -> str:
    return f"""You are an expert document analyst. Analyze the document text below and respond ONLY with valid JSON — no markdown, no code fences, no extra text.

Return exactly this structure:
{{
  "title": "A professional, descriptive title for this document (max 15 words)",
  "summary": "A concise, coherent summary preserving the key information (150-250 words)",
  "keywords": ["keyword1", "keyword2", "keyword3", "keyword4", "keyword5", "keyword6", "keyword7", "keyword8", "keyword9", "keyword10"]
}}

Rules:
- title: Meaningful, professional, specific to the content. Do NOT use generic titles.
- summary: Preserve facts, figures, names, and key arguments. Write in third-person prose.
- keywords: Exactly 10 keywords. Order by importance (most important first). Single words or short phrases only.

DOCUMENT TEXT:
\"\"\"
{text}
\"\"\"
"""


def call_gemini(prompt: str) -> dict:
    if not GEMINI_API_KEY:
        raise HTTPException(
            status_code=500,
            detail="GEMINI_API_KEY environment variable is not set.",
        )
    model = genai.GenerativeModel(
        model_name=GEMINI_MODEL,
        generation_config=genai.GenerationConfig(
            temperature=0.3,
            max_output_tokens=8192,
        ),
    )
    response = model.generate_content(prompt)
    raw = response.text.strip()

    # Strip markdown fences if model adds them anyway
    raw = re.sub(r"^```(?:json)?\s*", "", raw, flags=re.IGNORECASE | re.MULTILINE)
    raw = re.sub(r"\s*```\s*$", "", raw, flags=re.MULTILINE)
    raw = raw.strip()

    # Extract the first JSON object if there is surrounding text
    match = re.search(r"\{.*\}", raw, re.DOTALL)
    if match:
        raw = match.group(0)

    return json.loads(raw)


# ── Routes ───────────────────────────────────────────────────────────────────
@app.get("/", tags=["Health"])
def root():
    return {
        "service": "DocAI – Document Analyzer",
        "model": GEMINI_MODEL,
        "status": "running",
        "endpoints": {
            "analyze": "POST /analyze",
            "health": "GET /health",
            "docs": "GET /docs",
        },
    }


@app.get("/health", tags=["Health"])
def health():
    return {"status": "ok", "model": GEMINI_MODEL}


@app.post("/analyze", response_model=AnalysisResult, tags=["Analysis"])
async def analyze_document(file: UploadFile = File(...)):
    """
    Upload a PDF, DOCX, or TXT file.
    Returns an AI-generated title, summary, and 10 keywords.
    """
    # ── Validate content type ────────────────────────────────────────────────
    content_type = file.content_type or ""
    file_type = ALLOWED_TYPES.get(content_type)

    # Fallback: guess from extension
    if not file_type and file.filename:
        ext = file.filename.rsplit(".", 1)[-1].lower()
        if ext in ("pdf", "docx", "doc", "txt"):
            file_type = ext

    if not file_type:
        raise HTTPException(
            status_code=415,
            detail=f"Unsupported file type '{content_type}'. Allowed: PDF, DOCX, TXT.",
        )

    # ── Read & size-check ────────────────────────────────────────────────────
    file_bytes = await file.read()
    if len(file_bytes) > MAX_FILE_SIZE_BYTES:
        raise HTTPException(
            status_code=413,
            detail=f"File too large. Maximum size is {MAX_FILE_SIZE_MB} MB.",
        )
    if len(file_bytes) == 0:
        raise HTTPException(status_code=400, detail="Uploaded file is empty.")

    # ── Extract text ─────────────────────────────────────────────────────────
    try:
        raw_text = extract_text(file_bytes, file_type)
    except Exception as exc:
        raise HTTPException(
            status_code=422,
            detail=f"Could not extract text from document: {exc}",
        )

    raw_text = raw_text.strip()
    if not raw_text:
        raise HTTPException(
            status_code=422,
            detail="No readable text found in the document.",
        )

    # ── Analyze with Gemini ──────────────────────────────────────────────────
    truncated = truncate_text(raw_text)
    prompt = build_prompt(truncated)

    try:
        result = call_gemini(prompt)
    except json.JSONDecodeError as exc:
        raise HTTPException(
            status_code=502,
            detail=f"Gemini returned non-JSON response: {exc}",
        )
    except Exception as exc:
        raise HTTPException(
            status_code=502,
            detail=f"Gemini API error: {exc}",
        )

    # ── Validate and normalize result ────────────────────────────────────────
    title = str(result.get("title", "Untitled Document")).strip()
    summary = str(result.get("summary", "")).strip()
    keywords = result.get("keywords", [])
    if not isinstance(keywords, list):
        keywords = []
    keywords = [str(k).strip() for k in keywords if k][:10]

    return AnalysisResult(
        title=title,
        summary=summary,
        keywords=keywords,
        word_count=len(raw_text.split()),
        char_count=len(raw_text),
        file_name=file.filename or "unknown",
        file_type=file_type,
    )


# ── Entry point (local dev) ──────────────────────────────────────────────────
if __name__ == "__main__":
    uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True)