backend/app/documents/service.py

"""
Document service — text extraction, AI analysis, multilingual chat.
Supports PDF, DOCX, TXT, CSV.
"""
import io
import os
import re
import subprocess
from typing import Optional


# ─── Text extraction ──────────────────────────────────────────────────────────

def extract_text_from_pdf(file_bytes: bytes) -> str:
    def clean(text: str) -> str:
        return re.sub(r"\n{3,}", "\n\n", text).strip()

    def ocr_page_with_tesseract(page) -> str:
        pix = page.get_pixmap(dpi=200, alpha=False)
        proc = subprocess.run(
            ["tesseract", "stdin", "stdout", "-l", "eng", "--psm", "6"],
            input=pix.tobytes("png"),
            capture_output=True,
            check=False,
            timeout=45,
        )
        if proc.returncode != 0:
            err = proc.stderr.decode("utf-8", errors="replace").strip()
            print(f"[documents] Tesseract OCR error: {err}")
            return ""
        return proc.stdout.decode("utf-8", errors="replace")

    try:
        import pypdf
        reader = pypdf.PdfReader(io.BytesIO(file_bytes))
        pages = []
        for page in reader.pages:
            text = page.extract_text() or ""
            if not text.strip():
                try:
                    text = page.extract_text(extraction_mode="layout") or ""
                except TypeError:
                    text = ""
            pages.append(text)
        extracted = clean("\n".join(pages))
        if extracted:
            return extracted
    except Exception as e:
        print(f"[documents] PDF extraction error: {e}")

    try:
        import fitz  # PyMuPDF
        with fitz.open(stream=file_bytes, filetype="pdf") as doc:
            extracted = "\n".join(page.get_text("text") for page in doc)
            extracted = clean(extracted)
            if extracted:
                return extracted

            ocr_pages = []
            for page in doc:
                try:
                    textpage = page.get_textpage_ocr(language="eng", dpi=200, full=True)
                    ocr_pages.append(page.get_text("text", textpage=textpage))
                except Exception as e:
                    print(f"[documents] PyMuPDF OCR error: {e}")

                if not clean(ocr_pages[-1] if ocr_pages else ""):
                    ocr_pages.append(ocr_page_with_tesseract(page))

        return clean("\n".join(ocr_pages))
    except Exception as e:
        print(f"[documents] PDF fallback extraction error: {e}")
        return ""


def extract_text_from_docx(file_bytes: bytes) -> str:
    try:
        import docx
        doc = docx.Document(io.BytesIO(file_bytes))
        paragraphs = [p.text for p in doc.paragraphs if p.text.strip()]
        return "\n".join(paragraphs).strip()
    except Exception as e:
        print(f"[documents] DOCX extraction error: {e}")
        return ""


def extract_text_from_txt(file_bytes: bytes) -> str:
    try:
        return file_bytes.decode("utf-8", errors="replace").strip()
    except Exception:
        return ""


def extract_text_from_csv(file_bytes: bytes) -> str:
    try:
        import csv
        text = file_bytes.decode("utf-8", errors="replace")
        reader = csv.reader(io.StringIO(text))
        rows = [", ".join(row) for row in reader]
        return "\n".join(rows[:200])  # cap at 200 rows
    except Exception as e:
        print(f"[documents] CSV extraction error: {e}")
        return ""


def extract_text(file_bytes: bytes, file_type: str) -> str:
    ft = file_type.lower()
    if ft == "pdf":
        return extract_text_from_pdf(file_bytes)
    elif ft == "docx":
        return extract_text_from_docx(file_bytes)
    elif ft == "txt":
        return extract_text_from_txt(file_bytes)
    elif ft == "csv":
        return extract_text_from_csv(file_bytes)
    return ""


# ─── LLM caller ───────────────────────────────────────────────────────────────

def _call_llm(messages: list, max_tokens: int = 800) -> Optional[str]:
    openai_key = os.environ.get("OPENAI_API_KEY", "")
    groq_key = os.environ.get("GROQ_API_KEY", "") or os.environ.get("GROQ_KEY", "")

    if openai_key:
        try:
            from openai import OpenAI
            client = OpenAI(api_key=openai_key)
            res = client.chat.completions.create(
                model="gpt-4o-mini",
                messages=messages,
                temperature=0.2,
                max_tokens=max_tokens,
            )
            return res.choices[0].message.content.strip()
        except Exception as e:
            print(f"[documents] OpenAI error: {e}")

    if groq_key:
        try:
            from groq import Groq
            client = Groq(api_key=groq_key)
            res = client.chat.completions.create(
                model="llama-3.3-70b-versatile",
                messages=messages,
                temperature=0.2,
                max_tokens=max_tokens,
            )
            return res.choices[0].message.content.strip()
        except Exception as e:
            print(f"[documents] Groq error: {e}")

    return None


# ─── Language helpers ─────────────────────────────────────────────────────────

LANG_INSTRUCTIONS = {
    "en": "Respond in English.",
    "hi": "हिंदी में जवाब दें। (Respond in Hindi.)",
    "mr": "मराठीत उत्तर द्या. (Respond in Marathi.)",
}


def _lang_instruction(language: str) -> str:
    return LANG_INSTRUCTIONS.get(language, LANG_INSTRUCTIONS["en"])


# ─── AI document analysis ─────────────────────────────────────────────────────

def analyze_document(extracted_text: str, filename: str, language: str = "en") -> dict:
    """
    Generates an AI summary + financial insights from extracted document text.
    Returns {"summary": str, "insights": list[str], "suspicious": list[str]}
    """
    if not extracted_text.strip():
        return {
            "summary": "Could not extract text from this document.",
            "insights": [],
            "suspicious": [],
        }

    # Truncate to ~6000 chars to stay within token limits
    text_chunk = extracted_text[:6000]
    lang_note = _lang_instruction(language)

    system = (
        "You are an expert financial document analyst. "
        "Analyze the provided document and extract key financial information. "
        f"{lang_note}"
    )

    user_prompt = f"""Analyze this financial document: "{filename}"

DOCUMENT CONTENT:
{text_chunk}

Provide:
1. SUMMARY: A 2-3 sentence summary of what this document contains.
2. KEY FINANCIAL INSIGHTS: List 3-5 specific financial facts, amounts, or patterns found.
3. SUSPICIOUS ITEMS: List any transactions or entries that look unusual or suspicious (or say "None detected").

Format your response exactly as:
SUMMARY:
[your summary]

KEY INSIGHTS:
- [insight 1]
- [insight 2]
- [insight 3]

SUSPICIOUS:
- [item 1 or "None detected"]"""

    messages = [
        {"role": "system", "content": system},
        {"role": "user", "content": user_prompt},
    ]

    response = _call_llm(messages, max_tokens=600)

    if not response:
        # Rule-based fallback
        amounts = re.findall(r'\$[\d,]+\.?\d*', text_chunk)
        return {
            "summary": f"Document '{filename}' processed. Contains {len(extracted_text.split())} words.",
            "insights": [f"Found {len(amounts)} monetary amounts"] if amounts else ["No structured financial data detected"],
            "suspicious": [],
        }

    # Parse response
    summary = ""
    insights = []
    suspicious = []

    lines = response.split("\n")
    section = None
    for line in lines:
        line = line.strip()
        if line.upper().startswith("SUMMARY"):
            section = "summary"
        elif "KEY INSIGHT" in line.upper():
            section = "insights"
        elif "SUSPICIOUS" in line.upper():
            section = "suspicious"
        elif line.startswith("- ") and section == "insights":
            insights.append(line[2:])
        elif line.startswith("- ") and section == "suspicious":
            suspicious.append(line[2:])
        elif section == "summary" and line and not line.upper().startswith("SUMMARY"):
            summary += line + " "

    return {
        "summary": summary.strip() or response[:300],
        "insights": insights[:5],
        "suspicious": [s for s in suspicious if s.lower() != "none detected"][:5],
    }


# ─── AI document chat ─────────────────────────────────────────────────────────

def chat_with_document(
    question: str,
    extracted_text: str,
    filename: str,
    history: list,
    language: str = "en",
) -> str:
    """
    Answers a question about the document using only the document's content.
    history: list of {"role": "user"|"assistant", "content": str}
    """
    if not extracted_text.strip():
        return "I couldn't extract text from this document. Please try uploading again."

    text_chunk = extracted_text[:5000]
    lang_note = _lang_instruction(language)

    system = f"""You are a document analysis assistant. You have access to the content of the document "{filename}".

CRITICAL RULES:
1. Answer ONLY based on the document content provided below.
2. If the answer is not in the document, say "This information is not in the document."
3. Never make up information not present in the document.
4. Be specific — quote exact figures, dates, and names from the document.
5. {lang_note}

DOCUMENT CONTENT:
{text_chunk}"""

    messages = [{"role": "system", "content": system}]

    # Add conversation history (last 6 exchanges)
    for msg in history[-12:]:
        messages.append({"role": msg["role"], "content": msg["content"]})

    messages.append({"role": "user", "content": question})

    response = _call_llm(messages, max_tokens=500)

    if not response:
        return f"I found the document '{filename}' but couldn't generate a response. Please try again."

    return response