app.py

import os
import re
import traceback
import requests
from fastapi import FastAPI, File, Form, UploadFile
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from io import BytesIO
from docx import Document
from docx_builder import create_docx_with_layout

app = FastAPI()

# ──────────────────────────────── CORS ─────────────────────────────────
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # restrict to frontend domain in production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

# ──────────────────────────────── ROOT ─────────────────────────────────
@app.get("/")
async def root():
    return {
        "status": "ok",
        "message": "Document Processing API (Gemini-2.0-Flash only)",
        "endpoints": {
            "POST /process-document": "Processes a document using Gemini-2.0-Flash model"
        },
    }

# ─────────────────────────────── HELPERS ───────────────────────────────
def clean_ai_response(text: str) -> str:
    """Removes conversational phrases and keeps only the document content."""
    lines = text.strip().split("\n")
    cleaned_lines = []
    for i, line in enumerate(lines):
        line_stripped = line.strip().lower()
        if i < 3 and len(line_stripped) < 100:
            if any(
                kw in line_stripped
                for kw in ["sure", "okay", "certainly", "here is", "here's", "enhanced", "revised", "version", "below is"]
            ):
                continue
        cleaned_lines.append(line)
    return "\n".join(cleaned_lines).strip()

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

def extract_text_from_pdf(content: bytes) -> str:
    try:
        import PyPDF2
        pdf_file = BytesIO(content)
        pdf_reader = PyPDF2.PdfReader(pdf_file)
        return "\n\n".join([page.extract_text() for page in pdf_reader.pages])
    except ImportError:
        raise ValueError("PDF processing not available. Install PyPDF2 or upload .docx/.txt files.")

# ───────────────────────────── GEMINI CALL ─────────────────────────────
def call_gemini_api(text: str, user_prompt: str) -> str:
    """Calls Gemini-2.0-Flash model with both a system instruction and user prompt."""
    if not GEMINI_API_KEY:
        raise ValueError("GEMINI_API_KEY not set")

    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key={GEMINI_API_KEY}"
    headers = {"Content-Type": "application/json"}

    system_instruction = (
        "You are Verolabz, a document enhancement AI. "
        "Your goal is to enhance and refine the document text while keeping its meaning "
        "and layout intact. Do not add explanations, introductions, or summaries. "
        "Return only the improved text with preserved structure and formatting cues."
    )

    payload = {
        "system_instruction": {"parts": [{"text": system_instruction}]},
        "contents": [
            {
                "role": "user",
                "parts": [{"text": f"User instructions: {user_prompt}\n\nDocument text:\n{text}"}],
            }
        ],
    }

    res = requests.post(url, headers=headers, json=payload)
    if res.status_code != 200:
        raise Exception(f"Gemini API error: {res.text}")

    data = res.json()
    try:
        return data["candidates"][0]["content"]["parts"][0]["text"]
    except (KeyError, IndexError):
        raise Exception(f"Unexpected Gemini API response: {data}")

# ────────────────────────────── MAIN ROUTE ─────────────────────────────
@app.post("/process-document")
async def process_document(file: UploadFile = File(...), user_prompt: str = Form(...)):
    try:
        content = await file.read()
        filename = file.filename.lower()

        # Extract text
        if filename.endswith(".docx"):
            text = extract_text_from_docx(content)
        elif filename.endswith(".pdf"):
            text = extract_text_from_pdf(content)
        elif filename.endswith(".txt"):
            text = content.decode("utf-8", errors="ignore")
        else:
            return JSONResponse({"error": "Unsupported file type. Use .docx, .pdf, or .txt"}, status_code=400)

        if not text.strip():
            return JSONResponse({"error": "Document is empty"}, status_code=400)

        # Call Gemini
        result_text = call_gemini_api(text, user_prompt)
        cleaned_text = clean_ai_response(result_text)

        # Rebuild output DOCX
        output = create_docx_with_layout(cleaned_text)

        return StreamingResponse(
            output,
            media_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
            headers={
                "Content-Disposition": f"attachment; filename=enhanced_{file.filename.replace('.pdf','.docx').replace('.txt','.docx')}"
            },
        )
    except Exception as e:
        traceback.print_exc()
        return JSONResponse({"error": str(e)}, status_code=500)