"""
GUVI Hackathon Track 2 - AI-Powered Document Analysis & Extraction
==================================================================
Endpoint : POST /api/document-analyze
Auth     : x-api-key header (401 if missing/invalid)
Formats  : PDF, DOCX, Image (OCR)
AI Stack : Gemini 2.5 Flash (primary) → spaCy + DistilBERT + sumy (fallback)
"""

# ─────────────────────────────────────────────────────────────────────────────
# SECTION 1 — Imports & Configuration
# ─────────────────────────────────────────────────────────────────────────────
import base64
import io
import os
import re
import sys
import logging
import tempfile
from contextlib import asynccontextmanager
from typing import Any

import nltk
from dotenv import load_dotenv
from fastapi import Depends, FastAPI, Header, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel

from pathlib import Path
# Load .env — try multiple locations to be robust regardless of cwd
for _candidate in [
    Path(__file__).resolve().parent.parent / ".env",   # hcl-docu/.env (when running from src/)
    Path(__file__).resolve().parent / ".env",           # src/.env (if copied there)
    Path.cwd() / ".env",                               # cwd
    Path.cwd().parent / ".env",                        # cwd parent
]:
    if _candidate.exists():
        load_dotenv(_candidate, override=True)
        break

logging.basicConfig(level=logging.INFO, format="%(levelname)s | %(message)s")
logger = logging.getLogger(__name__)

# ── API Keys ──────────────────────────────────────────────────────────────────
API_KEY: str = os.getenv("API_KEY", "")
GEMINI_API_KEY: str = os.getenv("GEMINI_API_KEY", "")

# Startup debug — confirm key is loaded (prints only first 10 chars)
print(f"DEBUG | GEMINI_API_KEY loaded: {GEMINI_API_KEY[:10]}... (len={len(GEMINI_API_KEY)})" if GEMINI_API_KEY else "DEBUG | GEMINI_API_KEY is EMPTY!")

# ── Tesseract path: .env override → auto-detect common install locations ──────
import pytesseract
TESSERACT_PATH: str = os.getenv("TESSERACT_PATH", "")
if not TESSERACT_PATH:
    # Auto-detect on Windows / Linux / macOS
    _candidates = [
        r"C:\Program Files\Tesseract-OCR\tesseract.exe",
        r"C:\Program Files (x86)\Tesseract-OCR\tesseract.exe",
        os.path.expanduser("~/tesseract/tesseract.exe"),
        "/usr/bin/tesseract",
        "/usr/local/bin/tesseract",
        "/opt/homebrew/bin/tesseract",
    ]
    for _c in _candidates:
        if os.path.isfile(_c):
            TESSERACT_PATH = _c
            break
if TESSERACT_PATH:
    pytesseract.pytesseract.tesseract_cmd = TESSERACT_PATH
    print(f"DEBUG | Tesseract found: {TESSERACT_PATH}")
else:
    print("DEBUG | Tesseract not found — OCR for images will fail. Set TESSERACT_PATH in .env")

# ── Global model containers (loaded once at startup) ──────────────────────────
_nlp = None            # spaCy model
_sentiment_pipe = None # HuggingFace DistilBERT pipeline
_gemini_model = None   # Gemini generative model


# ─────────────────────────────────────────────────────────────────────────────
# SECTION 2 — Pydantic Request / Response Models
# ─────────────────────────────────────────────────────────────────────────────
class DocumentRequest(BaseModel):
    """Incoming request body: base64-encoded document."""
    fileName: str
    fileType: str   # pdf | docx | image
    fileBase64: str


class EntitiesModel(BaseModel):
    names: list[str] = []
    dates: list[str] = []
    organizations: list[str] = []
    amounts: list[str] = []


class DocumentResponse(BaseModel):
    """Exact response shape required by the problem statement."""
    status: str                   # "success" | "error"
    fileName: str
    summary: str = ""
    entities: EntitiesModel = EntitiesModel()
    sentiment: str = "Neutral"   # Positive | Neutral | Negative


# ─────────────────────────────────────────────────────────────────────────────
# SECTION 3 — Authentication
# ─────────────────────────────────────────────────────────────────────────────
def verify_api_key(x_api_key: str | None = Header(default=None, alias="x-api-key")) -> str:
    """
    FastAPI dependency that validates the x-api-key header.
    Returns 401 for missing OR invalid key (PS requirement).
    """
    if not API_KEY:
        raise HTTPException(status_code=500, detail="Server API_KEY not configured.")
    if not x_api_key or x_api_key != API_KEY:
        raise HTTPException(status_code=401, detail="Unauthorized: invalid or missing API key.")
    return x_api_key


# ─────────────────────────────────────────────────────────────────────────────
# SECTION 4 — Text Extraction Layer
# ─────────────────────────────────────────────────────────────────────────────

def _ocr_bytes(image_bytes: bytes) -> str:
    """Run Tesseract OCR on raw image bytes with preprocessing for better accuracy."""
    import pytesseract
    from PIL import Image, ImageFilter, ImageOps

    img = Image.open(io.BytesIO(image_bytes))

    # Upscale small images (Tesseract works best at 300 DPI / large text)
    w, h = img.size
    if w < 1000 or h < 1000:
        scale = max(2, 1500 // min(w, h))
        img = img.resize((w * scale, h * scale), Image.LANCZOS)

    # Preprocessing: grayscale → sharpen → contrast → binarize
    img = img.convert("L")
    img = img.filter(ImageFilter.SHARPEN)
    img = ImageOps.autocontrast(img)
    img = img.point(lambda x: 0 if x < 140 else 255)  # binarize

    text = pytesseract.image_to_string(img, config="--psm 6")
    return text.strip()


def extract_text_from_pdf(file_bytes: bytes) -> str:
    """
    Extract text from PDF bytes using pdfplumber.
    Falls back to OCR (via pdf2image + pytesseract) for scanned/image-only pages.
    """
    import pdfplumber

    full_text: list[str] = []
    scanned_page_indices: list[int] = []

    with pdfplumber.open(io.BytesIO(file_bytes)) as pdf:
        for i, page in enumerate(pdf.pages):
            page_text = page.extract_text()
            if page_text and page_text.strip():
                full_text.append(page_text.strip())
            else:
                scanned_page_indices.append(i)  # mark for OCR

    # OCR fallback for pages that yielded no text
    if scanned_page_indices:
        try:
            from pdf2image import convert_from_bytes
            images = convert_from_bytes(file_bytes, dpi=200)
            for idx in scanned_page_indices:
                if idx < len(images):
                    ocr_text = _ocr_bytes(
                        _pil_to_bytes(images[idx])
                    )
                    if ocr_text:
                        full_text.append(ocr_text)
        except Exception as e:
            logger.warning(f"PDF OCR fallback failed: {e}")

    return "\n\n".join(full_text)


def extract_text_from_docx(file_bytes: bytes) -> str:
    """
    Extract text from DOCX bytes using python-docx.
    Includes paragraphs, headings, and table cell content.
    """
    from docx import Document

    doc = Document(io.BytesIO(file_bytes))
    parts: list[str] = []

    # Paragraphs & headings
    for para in doc.paragraphs:
        if para.text.strip():
            parts.append(para.text.strip())

    # Table content
    for table in doc.tables:
        for row in table.rows:
            row_texts = [cell.text.strip() for cell in row.cells if cell.text.strip()]
            if row_texts:
                parts.append(" | ".join(row_texts))

    return "\n".join(parts)


def extract_text_from_image(file_bytes: bytes) -> str:
    """Extract text from image bytes using Tesseract OCR."""
    return _ocr_bytes(file_bytes)


def _pil_to_bytes(pil_image) -> bytes:
    """Convert a PIL Image to PNG bytes."""
    buf = io.BytesIO()
    pil_image.save(buf, format="PNG")
    return buf.getvalue()


def extract_text(file_bytes: bytes, file_type: str) -> str:
    """
    Router: dispatch to the correct extraction function based on file type.
    Returns the extracted plain-text string.
    """
    ft = file_type.lower().strip()
    if ft == "pdf":
        return extract_text_from_pdf(file_bytes)
    elif ft == "docx":
        return extract_text_from_docx(file_bytes)
    elif ft in ("image", "img", "png", "jpg", "jpeg", "tiff", "bmp", "gif"):
        return extract_text_from_image(file_bytes)
    else:
        raise ValueError(f"Unsupported fileType: '{file_type}'. Use pdf, docx, or image.")


# ─────────────────────────────────────────────────────────────────────────────
# SECTION 5 — AI Analysis: Gemini 2.5 Flash (Primary)
# ─────────────────────────────────────────────────────────────────────────────

_GEMINI_PROMPT_TEMPLATE = """You are a precise document analysis expert. Analyze the document text below and return ONLY valid JSON — no markdown fences, no explanation, nothing else.

Document filename: {file_name}

Document text:
\"\"\"
{text}
\"\"\"

Return this exact JSON structure:
{{
  "summary": "<concise 1–2 sentence summary capturing the main purpose and key facts of the document>",
  "entities": {{
    "names": ["<person names found in the document>"],
    "dates": ["<dates found, preserve the original format exactly as it appears>"],
    "organizations": ["<company or organization names found>"],
    "amounts": ["<monetary amounts found, include currency symbol>"]
  }},
  "sentiment": "<exactly one of: Positive, Neutral, Negative>"
}}

Rules:
- summary: capture the who/what/when/how-much of the document in 1-2 sentences
- names: only real person names (not job titles, not places)
- dates: preserve original format (e.g. "10 March 2026", "2026-03-10", "10/03/2026")
- organizations: company names, institutions, brands
- amounts: include currency symbol and formatting (e.g. "₹10,000", "$500.00", "Rs. 2,50,000")
- sentiment: Positive for praise/good news/approvals, Negative for complaints/rejections/problems, Neutral for factual/informational content
- Return empty arrays [] for entity types not found
- Return ONLY the JSON object, nothing else"""


def analyze_with_gemini(text: str, file_name: str) -> dict[str, Any]:
    """
    Call Gemini 2.5 Flash via direct HTTP API (no SDK dependency issues).
    Returns a dict with keys: summary, entities (dict), sentiment.
    Raises RuntimeError if Gemini is unavailable or returns invalid JSON.
    """
    import json
    import requests as _requests

    if not GEMINI_API_KEY:
        raise RuntimeError("GEMINI_API_KEY not set.")

    # Truncate to avoid token limits (keep first 6000 chars — plenty for context)
    truncated_text = text[:6000] if len(text) > 6000 else text

    prompt = _GEMINI_PROMPT_TEMPLATE.format(
        file_name=file_name,
        text=truncated_text,
    )

    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent?key={GEMINI_API_KEY}"
    payload = {
        "contents": [{"parts": [{"text": prompt}]}],
        "generationConfig": {
            "temperature": 0.1,
            "maxOutputTokens": 1024,
        },
    }

    resp = _requests.post(url, json=payload, timeout=60)
    if resp.status_code != 200:
        raise RuntimeError(f"Gemini HTTP {resp.status_code}: {resp.text[:300]}")

    data = resp.json()
    raw = data["candidates"][0]["content"]["parts"][0]["text"].strip()

    # Strip any accidental markdown fences
    if raw.startswith("```"):
        raw = re.sub(r"^```[a-z]*\n?", "", raw)
        raw = re.sub(r"\n?```$", "", raw)
        raw = raw.strip()

    result = json.loads(raw)

    # Normalise entity keys to match response schema
    entities = result.get("entities", {})
    return {
        "summary": str(result.get("summary", "")).strip(),
        "entities": {
            "names":         _dedup(entities.get("names", [])),
            "dates":         _dedup(entities.get("dates", [])),
            "organizations": _dedup(entities.get("organizations", [])),
            "amounts":       _dedup(entities.get("amounts", [])),
        },
        "sentiment": _normalize_sentiment(result.get("sentiment", "Neutral")),
    }


# ─────────────────────────────────────────────────────────────────────────────
# SECTION 6 — AI Analysis: Offline Fallback
#             (spaCy NER + regex + DistilBERT + sumy TextRank)
# ─────────────────────────────────────────────────────────────────────────────

# ── Regex patterns ────────────────────────────────────────────────────────────
_DATE_PATTERNS = [
    # "10 March 2026", "10th March 2026"
    r"\b\d{1,2}(?:st|nd|rd|th)?\s+(?:January|February|March|April|May|June|"
    r"July|August|September|October|November|December)\s+\d{4}\b",
    # "March 10, 2026" / "March 10 2026"
    r"\b(?:January|February|March|April|May|June|July|August|September|"
    r"October|November|December)\s+\d{1,2},?\s+\d{4}\b",
    # "Jan 10, 2026" abbreviated
    r"\b(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)\.?\s+\d{1,2},?\s+\d{4}\b",
    # "10/03/2026", "2026-03-10", "10-03-2026"
    r"\b\d{1,2}[/\-]\d{1,2}[/\-]\d{2,4}\b",
    r"\b\d{4}[/\-]\d{1,2}[/\-]\d{1,2}\b",
]

_AMOUNT_PATTERNS = [
    # ₹10,000 / ₹10000 / ₹ 10,000.50
    r"₹\s?[\d,]+(?:\.\d{1,2})?",
    # Rs. 10,000 / Rs 10000
    r"Rs\.?\s?[\d,]+(?:\.\d{1,2})?",
    # INR 10,000
    r"INR\s?[\d,]+(?:\.\d{1,2})?",
    # $500 / $500.00
    r"\$\s?[\d,]+(?:\.\d{1,2})?",
    # USD 500
    r"USD\s?[\d,]+(?:\.\d{1,2})?",
    # € 100
    r"€\s?[\d,]+(?:\.\d{1,2})?",
    # £ 100
    r"£\s?[\d,]+(?:\.\d{1,2})?",
]


def _regex_extract_dates(text: str) -> list[str]:
    found = []
    for pattern in _DATE_PATTERNS:
        found.extend(re.findall(pattern, text, re.IGNORECASE))
    return _dedup(found)


def _regex_extract_amounts(text: str) -> list[str]:
    found = []
    for pattern in _AMOUNT_PATTERNS:
        found.extend(re.findall(pattern, text))
    return _dedup(found)


def generate_summary_fallback(text: str) -> str:
    """Summarise text using sumy TextRank algorithm (100% offline)."""
    try:
        from sumy.parsers.plaintext import PlaintextParser
        from sumy.nlp.tokenizers import Tokenizer
        from sumy.summarizers.text_rank import TextRankSummarizer

        # Use first 4000 chars to keep sumy fast
        snippet = text[:4000]
        parser = PlaintextParser.from_string(snippet, Tokenizer("english"))
        summarizer = TextRankSummarizer()
        sentences = summarizer(parser.document, sentences_count=2)
        summary = " ".join(str(s) for s in sentences).strip()
        return summary if summary else text[:250].strip() + "..."
    except Exception as e:
        logger.warning(f"sumy summarisation failed: {e}")
        # Last-resort: first 250 characters as stub summary
        return text[:250].strip() + ("..." if len(text) > 250 else "")


def extract_entities_fallback(text: str) -> dict[str, list[str]]:
    """
    Extract named entities using spaCy NER + regex.
    spaCy handles PERSON / ORG; regex handles dates and monetary amounts.
    """
    global _nlp
    names: list[str] = []
    organizations: list[str] = []

    if _nlp is not None:
        try:
            # Truncate for speed (spaCy can be slow on very long docs)
            doc = _nlp(text[:5000])
            for ent in doc.ents:
                if ent.label_ == "PERSON":
                    names.append(ent.text.strip())
                elif ent.label_ in ("ORG", "GPE"):
                    organizations.append(ent.text.strip())
        except Exception as e:
            logger.warning(f"spaCy NER failed: {e}")

    dates = _regex_extract_dates(text)
    amounts = _regex_extract_amounts(text)

    return {
        "names":         _dedup(names),
        "dates":         _dedup(dates),
        "organizations": _dedup(organizations),
        "amounts":       _dedup(amounts),
    }


def analyze_sentiment_fallback(text: str) -> str:
    """
    Classify sentiment using DistilBERT SST-2 (offline HuggingFace model).
    Maps POSITIVE → "Positive", NEGATIVE → "Negative".
    Score < 0.65 maps to "Neutral" to avoid overconfident labelling.
    """
    global _sentiment_pipe
    if _sentiment_pipe is None:
        return "Neutral"

    try:
        # Use first 512 tokens worth of text
        snippet = text[:1500]
        result = _sentiment_pipe(snippet, truncation=True, max_length=512)[0]
        label: str = result["label"]   # "POSITIVE" or "NEGATIVE"
        score: float = result["score"]

        if score < 0.65:
            return "Neutral"
        return "Positive" if label == "POSITIVE" else "Negative"
    except Exception as e:
        logger.warning(f"Sentiment analysis failed: {e}")
        return "Neutral"


def analyze_with_fallback(text: str, file_name: str) -> dict[str, Any]:
    """Orchestrate offline analysis: summary + entities + sentiment."""
    summary = generate_summary_fallback(text)
    entities = extract_entities_fallback(text)
    sentiment = analyze_sentiment_fallback(text)
    return {
        "summary": summary,
        "entities": entities,
        "sentiment": sentiment,
    }


# ─────────────────────────────────────────────────────────────────────────────
# SECTION 7 — Main Analysis Dispatcher
# ─────────────────────────────────────────────────────────────────────────────

def analyze_document(text: str, file_name: str) -> dict[str, Any]:
    """
    Try Gemini first; if it fails (network error, quota, bad JSON) fall back
    to the fully-offline pipeline so the API never returns an empty response.
    """
    if GEMINI_API_KEY:
        try:
            result = analyze_with_gemini(text, file_name)
            logger.info("Analysis completed via Gemini.")
            return result
        except Exception as e:
            logger.warning(f"Gemini failed ({e}), switching to offline fallback.")

    result = analyze_with_fallback(text, file_name)
    logger.info("Analysis completed via offline fallback.")
    return result


# ─────────────────────────────────────────────────────────────────────────────
# SECTION 8 — Startup: preload heavy models
# ─────────────────────────────────────────────────────────────────────────────

@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    Preload spaCy and DistilBERT at startup so the first request isn't slow.
    NLTK punkt tokenizer (required by sumy) is also downloaded if missing.
    """
    global _nlp, _sentiment_pipe

    # spaCy
    try:
        import spacy
        _nlp = spacy.load("en_core_web_sm")
        logger.info("spaCy model loaded.")
    except Exception as e:
        logger.warning(f"spaCy load failed (entity extraction will be regex-only): {e}")

    # DistilBERT sentiment
    try:
        from transformers import pipeline as hf_pipeline
        _sentiment_pipe = hf_pipeline(
            "sentiment-analysis",
            model="distilbert-base-uncased-finetuned-sst-2-english",
            device=-1,   # CPU
        )
        logger.info("DistilBERT sentiment model loaded.")
    except Exception as e:
        logger.warning(f"DistilBERT load failed (sentiment will default to Neutral): {e}")

    # NLTK punkt (required by sumy)
    try:
        nltk.download("punkt", quiet=True)
        nltk.download("punkt_tab", quiet=True)
        nltk.download("stopwords", quiet=True)
        logger.info("NLTK data ready.")
    except Exception as e:
        logger.warning(f"NLTK download failed: {e}")

    yield  # application runs

    logger.info("Shutting down.")


# ─────────────────────────────────────────────────────────────────────────────
# SECTION 9 — FastAPI Application
# ─────────────────────────────────────────────────────────────────────────────

app = FastAPI(
    title="GUVI Track 2 — Document Analysis API",
    description="Multi-format document analysis: PDF, DOCX, Image (OCR). AI-powered summary, entity extraction, and sentiment analysis.",
    version="1.0.0",
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["POST", "GET"],
    allow_headers=["*"],
)


# ── Global exception handler: ensure ALL errors return the required JSON shape ─
@app.exception_handler(HTTPException)
async def http_exception_handler(request: Request, exc: HTTPException):
    """Return errors in the same JSON shape the PS requires (except 401 which is standard)."""
    if exc.status_code == 401:
        return JSONResponse(status_code=401, content={"detail": exc.detail})
    return JSONResponse(
        status_code=exc.status_code,
        content={
            "status": "error",
            "fileName": "",
            "summary": str(exc.detail),
            "entities": {"names": [], "dates": [], "organizations": [], "amounts": []},
            "sentiment": "Neutral",
        },
    )


# ─────────────────────────────────────────────────────────────────────────────
# SECTION 10 — Endpoints
# ─────────────────────────────────────────────────────────────────────────────

@app.get("/", tags=["Health"])
async def health_check():
    """Health check endpoint."""
    return {"status": "ok", "message": "Document Analysis API is running."}


@app.post(
    "/api/document-analyze",
    response_model=DocumentResponse,
    tags=["Document Analysis"],
    summary="Analyse a base64-encoded PDF, DOCX, or image document",
)
async def document_analyze(
    body: DocumentRequest,
    _key: str = Depends(verify_api_key),
):
    """
    Main analysis endpoint.

    - **fileName**: original file name (echoed back in response)
    - **fileType**: `pdf` | `docx` | `image`
    - **fileBase64**: base64-encoded file content

    Returns summary, named entities, and sentiment.
    """

    try:
        # ── 1. Decode base64 ──────────────────────────────────────────────
        try:
            file_bytes = base64.b64decode(body.fileBase64)
        except Exception:
            raise HTTPException(status_code=400, detail="Invalid base64 encoding in fileBase64.")

        if not file_bytes:
            raise HTTPException(status_code=400, detail="fileBase64 decoded to empty bytes.")

        # ── 2. Extract text ───────────────────────────────────────────────
        try:
            text = extract_text(file_bytes, body.fileType)
        except ValueError as e:
            raise HTTPException(status_code=400, detail=str(e))
        except Exception as e:
            logger.error(f"Text extraction failed: {e}")
            raise HTTPException(status_code=422, detail=f"Text extraction error: {e}")

        if not text or not text.strip():
            # Return a graceful error rather than crashing
            return DocumentResponse(
                status="error",
                fileName=body.fileName,
                summary="Could not extract any text from the document.",
                entities=EntitiesModel(),
                sentiment="Neutral",
            )

        # ── 3. AI Analysis ────────────────────────────────────────────────
        analysis = analyze_document(text.strip(), body.fileName)

        # ── 4. Build response ─────────────────────────────────────────────
        ents = analysis.get("entities", {})
        return DocumentResponse(
            status="success",
            fileName=body.fileName,
            summary=analysis.get("summary", ""),
            entities=EntitiesModel(
                names=ents.get("names", []),
                dates=ents.get("dates", []),
                organizations=ents.get("organizations", []),
                amounts=ents.get("amounts", []),
            ),
            sentiment=analysis.get("sentiment", "Neutral"),
        )

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Unexpected error: {e}", exc_info=True)
        return DocumentResponse(
            status="error",
            fileName=body.fileName,
            summary=f"Processing failed: {str(e)}",
            entities=EntitiesModel(),
            sentiment="Neutral",
        )


# ─────────────────────────────────────────────────────────────────────────────
# SECTION 11 — Helpers
# ─────────────────────────────────────────────────────────────────────────────

def _dedup(items: list) -> list[str]:
    """Remove duplicates while preserving order; cast all items to str."""
    seen: set = set()
    result: list[str] = []
    for item in items:
        s = str(item).strip()
        if s and s.lower() not in seen:
            seen.add(s.lower())
            result.append(s)
    return result


def _normalize_sentiment(raw: str) -> str:
    """Normalise any sentiment string to exactly Positive / Neutral / Negative."""
    mapping = {
        "positive": "Positive",
        "negative": "Negative",
        "neutral":  "Neutral",
    }
    return mapping.get(raw.strip().lower(), "Neutral")


# ─────────────────────────────────────────────────────────────────────────────
# SECTION 12 — Entry point
# ─────────────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=False)