"""
Surya OCR - Universal REST API
One endpoint: POST /ocr — any file in, extracted text out.
Supports: Images, PDF, Word (.docx/.doc), PowerPoint (.pptx/.ppt), and more.
Secret key authentication via X-API-Key header.
Hugging Face Spaces (Free Tier) compatible.

`document_type` is REQUIRED — must be one of the supported types.
`parameters`    is REQUIRED — comma-separated field names the caller wants extracted.
Both drive structured JSON extraction via the Sarvam-M LLM.
"""

import io
import gc
import os
import json
import logging
import subprocess
import tempfile
from contextlib import asynccontextmanager
from pathlib import Path

import httpx
import uvicorn
from fastapi import FastAPI, File, Form, UploadFile, HTTPException, Security
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from fastapi.security import APIKeyHeader
from PIL import Image

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ── Auth ──────────────────────────────────────────────────────────────────────
API_SECRET_KEY = os.environ.get("API_SECRET_KEY", "")

api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)

def verify_key(api_key: str = Security(api_key_header)):
    if not API_SECRET_KEY:
        raise RuntimeError("API_SECRET_KEY environment variable is not set on the server.")
    if api_key != API_SECRET_KEY:
        raise HTTPException(
            status_code=401,
            detail="Invalid or missing API key. Pass it as X-API-Key header.",
        )
    return api_key


# ── Global model holders ──────────────────────────────────────────────────────
recognition_predictor = None
detection_predictor   = None

MAX_FILE_SIZE_MB = 30
MAX_PAGES        = 10

SUPPORTED_EXTENSIONS = {
    ".jpg", ".jpeg", ".png", ".webp", ".tiff", ".tif", ".bmp",
    ".pdf",
    ".docx", ".doc",
    ".pptx", ".ppt",
    ".odt", ".odp",
    ".rtf",
}


# ── Document Type Registry ────────────────────────────────────────────────────
# Only descriptions live here — fields are supplied by the caller via `parameters`.
#
# To add a new document type:
#   1. Add a new entry with a "description" key below.
#   2. No other code changes needed.

DOCUMENT_TYPES: dict[str, dict] = {
    "income_certificate": {
        "description": (
            "A government-issued income certificate that certifies the annual income "
            "of an individual or family for official and welfare purposes."
        ),
    },
    "caste_certificate": {
        "description": (
            "A government-issued caste certificate that officially recognises the caste "
            "category (SC / ST / OBC / etc.) of an individual for reservation and welfare purposes."
        ),
    },
    "domicile_certificate": {
        "description": (
            "A government-issued domicile certificate that proves an individual's permanent "
            "residence in a particular state or union territory."
        ),
    },
}

SUPPORTED_DOCUMENT_TYPES = sorted(DOCUMENT_TYPES.keys())


# ── Sarvam LLM Integration ────────────────────────────────────────────────────
SARVAM_API_KEY = os.environ.get("SARVAM_API_KEY", "")
SARVAM_API_URL = "https://api.sarvam.ai/v1/chat/completions"
SARVAM_MODEL   = "sarvam-m"


def parse_parameters(raw: str) -> list[str]:
    """
    Parse the comma-separated `parameters` string into a clean list of field names.
    e.g.  "name, income , income_in_words"  →  ["name", "income", "income_in_words"]
    """
    fields = [f.strip() for f in raw.split(",") if f.strip()]
    if not fields:
        raise HTTPException(
            status_code=400,
            detail=(
                "`parameters` must be a non-empty comma-separated list of field names. "
                "Example: parameters=name,income,income_in_words,validity"
            ),
        )
    return fields


def build_extraction_prompt(doc_type: str, fields: list[str], extracted_text: str) -> str:
    """Build a prompt so the LLM extracts exactly the fields the caller requested."""
    description = DOCUMENT_TYPES[doc_type]["description"]
    fields_list = "\n".join(f"- {f}" for f in fields)
    json_keys   = ", ".join(f'"{f}"' for f in fields)

    return f"""You are a precise document parser. You will be given raw OCR text extracted from a document.

Document Type: {doc_type.replace("_", " ").title()}
Description: {description}

Your task is to extract the following fields from the text:
{fields_list}

Rules:
- Return ONLY a valid JSON object with these exact keys: {json_keys}
- If a field cannot be found in the text, set its value to null.
- If any two related parameters do not match (possibly due to an OCR error), smartly return null.
- Do NOT include any explanation, markdown, or extra text — pure JSON only.
- Normalise values where sensible (e.g. trim whitespace, fix obvious OCR typos in names/numbers).

OCR Text:
\"\"\"
{extracted_text}
\"\"\"

JSON Output:"""


async def extract_structured_data(doc_type: str, fields: list[str], extracted_text: str) -> dict:
    """Call Sarvam-M to extract the caller-specified fields from OCR text."""
    if not SARVAM_API_KEY:
        raise HTTPException(
            status_code=503,
            detail="SARVAM_API_KEY environment variable is not set on the server.",
        )

    prompt = build_extraction_prompt(doc_type, fields, extracted_text)

    payload = {
        "model": SARVAM_MODEL,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0.0,  # deterministic for structured extraction
        "max_tokens": 512,
    }
    headers = {
        "Authorization": f"Bearer {SARVAM_API_KEY}",
        "Content-Type": "application/json",
    }

    try:
        async with httpx.AsyncClient(timeout=60) as client:
            response = await client.post(SARVAM_API_URL, json=payload, headers=headers)
            response.raise_for_status()
    except httpx.HTTPStatusError as e:
        raise HTTPException(
            status_code=502,
            detail=f"Sarvam API returned an error: {e.response.status_code} — {e.response.text}",
        )
    except httpx.RequestError as e:
        raise HTTPException(status_code=502, detail=f"Failed to reach Sarvam API: {e}")

    resp_json = response.json()

    try:
        raw_content = resp_json["choices"][0]["message"]["content"].strip()
    except (KeyError, IndexError) as e:
        raise HTTPException(
            status_code=502,
            detail=f"Unexpected Sarvam API response structure: {e} — {resp_json}",
        )

    # Strip <think>...</think> chain-of-thought block if present (Sarvam-M reasoning output)
    if "<think>" in raw_content:
        import re
        raw_content = re.sub(r"<think>.*?</think>", "", raw_content, flags=re.DOTALL).strip()

    # Strip markdown code fences if the LLM added them despite instructions
    if raw_content.startswith("```"):
        raw_content = raw_content.strip("`").lstrip("json").strip()

    try:
        structured = json.loads(raw_content)
    except json.JSONDecodeError:
        logger.warning(f"LLM returned non-JSON content: {raw_content!r}")
        raise HTTPException(
            status_code=502,
            detail=f"LLM did not return valid JSON. Raw response: {raw_content}",
        )

    return structured


# ── Model loading ─────────────────────────────────────────────────────────────

def load_models():
    global recognition_predictor, detection_predictor
    logger.info("Loading Surya OCR models...")
    from surya.foundation import FoundationPredictor
    from surya.recognition import RecognitionPredictor
    from surya.detection import DetectionPredictor

    foundation            = FoundationPredictor()
    recognition_predictor = RecognitionPredictor(foundation)
    detection_predictor   = DetectionPredictor()
    logger.info("Models loaded.")


@asynccontextmanager
async def lifespan(app: FastAPI):
    load_models()
    yield
    global recognition_predictor, detection_predictor
    recognition_predictor = None
    detection_predictor   = None
    gc.collect()


# ── App ───────────────────────────────────────────────────────────────────────

app = FastAPI(
    title="Surya OCR API",
    description=(
        "Upload any document — image, PDF, Word, PowerPoint — and receive raw OCR text "
        "plus structured JSON fields you define. "
        "`document_type` and `parameters` are both required. "
        "Requires X-API-Key header."
    ),
    version="4.0.0",
    lifespan=lifespan,
)

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


# ── File → PIL Images converters ──────────────────────────────────────────────

def convert_image(data: bytes) -> list:
    return [Image.open(io.BytesIO(data)).convert("RGB")]


def convert_pdf(data: bytes) -> list:
    from pdf2image import convert_from_bytes
    return convert_from_bytes(data, dpi=150, first_page=1, last_page=MAX_PAGES, fmt="RGB")


def convert_office(data: bytes, suffix: str) -> list:
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = Path(tmpdir) / f"input{suffix}"
        input_path.write_bytes(data)

        result = subprocess.run(
            ["libreoffice", "--headless", "--convert-to", "pdf",
             "--outdir", tmpdir, str(input_path)],
            capture_output=True, timeout=60,
        )
        if result.returncode != 0:
            raise HTTPException(
                status_code=422,
                detail=f"LibreOffice conversion failed: {result.stderr.decode()}",
            )

        pdf_path = Path(tmpdir) / (input_path.stem + ".pdf")
        if not pdf_path.exists():
            raise HTTPException(
                status_code=422,
                detail="Office-to-PDF conversion produced no output.",
            )

        from pdf2image import convert_from_path
        return convert_from_path(str(pdf_path), dpi=150, first_page=1, last_page=MAX_PAGES, fmt="RGB")


def file_to_images(data: bytes, suffix: str) -> list:
    image_exts  = {".jpg", ".jpeg", ".png", ".webp", ".tiff", ".tif", ".bmp"}
    office_exts = {".docx", ".doc", ".pptx", ".ppt", ".odt", ".odp", ".rtf"}

    if suffix in image_exts:
        return convert_image(data)
    elif suffix == ".pdf":
        return convert_pdf(data)
    elif suffix in office_exts:
        return convert_office(data, suffix)
    else:
        raise HTTPException(status_code=415, detail=f"Unsupported file type: '{suffix}'")


# ── OCR runner ────────────────────────────────────────────────────────────────

def run_ocr(images: list) -> str:
    if recognition_predictor is None:
        raise HTTPException(status_code=503, detail="Models not ready. Please retry in a moment.")

    try:
        predictions = recognition_predictor(images, det_predictor=detection_predictor)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"OCR failed: {e}")

    pages_text = []
    for page in predictions:
        page_lines = [line.text for line in page.text_lines]
        pages_text.append("\n".join(page_lines))

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


# ── Routes ────────────────────────────────────────────────────────────────────

@app.get("/", summary="Health / Info")
def root():
    """Public endpoint — no key required."""
    return {
        "service": "Surya OCR API",
        "status": "ready" if recognition_predictor else "loading",
        "usage": (
            "POST /ocr  |  form fields: file, document_type, parameters  |  header: X-API-Key. "
            "Example: -F 'document_type=income_certificate' -F 'parameters=name,income,validity'"
        ),
        "supported_formats": sorted(SUPPORTED_EXTENSIONS),
        "supported_document_types": SUPPORTED_DOCUMENT_TYPES,
    }


@app.get("/health", summary="Health check")
def health():
    """Public endpoint — no key required."""
    return {"status": "ready" if recognition_predictor else "loading"}


@app.get("/types", summary="List supported document types")
def list_types():
    """Public endpoint — returns all registered document types and their descriptions."""
    return {
        doc_type: {"description": cfg["description"]}
        for doc_type, cfg in DOCUMENT_TYPES.items()
    }


@app.post("/ocr", summary="Extract text and structured data from any document")
async def ocr(
    file: UploadFile = File(
        ...,
        description="Any image, PDF, Word, or PowerPoint file.",
    ),
    document_type: str = Form(
        ...,
        description=(
            "Type of document being uploaded. "
            f"Must be one of: {', '.join(SUPPORTED_DOCUMENT_TYPES)}."
        ),
    ),
    parameters: str = Form(
        ...,
        description=(
            "Comma-separated list of field names to extract from the document. "
            "You decide what to extract — any field name is accepted. "
            "Example: name,income,income_in_words,validity"
        ),
    ),
    _: str = Security(verify_key),
):
    """
    Upload **any** document and receive:

    - **`text`** — raw OCR output from the document.
    - **`data`** — structured JSON containing exactly the fields you requested, extracted by Sarvam-M LLM.

    **All three form fields are required:**

    | Field | Required | Description |
    |---|---|---|
    | `file` | ✅ | Document file (image / PDF / Word / PowerPoint) |
    | `document_type` | ✅ | One of: `income_certificate`, `caste_certificate`, `domicile_certificate` |
    | `parameters` | ✅ | Comma-separated fields to extract, e.g. `name,income,validity` |

    **Requires** `X-API-Key` header.

    **Example:**
    ```bash
    curl -X POST 'https://your-api/ocr' \\
      -H 'X-API-Key: your-secret-key' \\
      -F 'file=@income_cert.pdf;type=application/pdf' \\
      -F 'document_type=income_certificate' \\
      -F 'parameters=name,income,income_in_words,validity'
    ```
    """

    # ── Validate document_type ────────────────────────────────────────────────
    document_type = document_type.strip().lower()
    if document_type not in DOCUMENT_TYPES:
        raise HTTPException(
            status_code=400,
            detail=(
                f"Invalid `document_type`: '{document_type}'. "
                f"Must be one of: {', '.join(SUPPORTED_DOCUMENT_TYPES)}."
            ),
        )

    # ── Validate & parse parameters ───────────────────────────────────────────
    fields = parse_parameters(parameters)
    logger.info(f"Requested fields: {fields}")

    # ── Read & validate file ──────────────────────────────────────────────────
    data = await file.read()

    size_mb = len(data) / (1024 * 1024)
    if size_mb > MAX_FILE_SIZE_MB:
        raise HTTPException(
            status_code=413,
            detail=f"File too large ({size_mb:.1f} MB). Max: {MAX_FILE_SIZE_MB} MB.",
        )

    suffix = Path(file.filename or "").suffix.lower()
    if not suffix:
        raise HTTPException(
            status_code=400,
            detail="Filename must have an extension so the format can be detected.",
        )
    if suffix not in SUPPORTED_EXTENSIONS:
        raise HTTPException(
            status_code=415,
            detail=f"'{suffix}' is not supported. Supported: {sorted(SUPPORTED_EXTENSIONS)}",
        )

    logger.info(
        f"Processing '{file.filename}' ({size_mb:.1f} MB) | "
        f"document_type={document_type!r} | fields={fields}"
    )

    # ── OCR ───────────────────────────────────────────────────────────────────
    images = file_to_images(data, suffix)
    extracted_text = run_ocr(images)

    del images
    gc.collect()

    # ── LLM structured extraction ─────────────────────────────────────────────
    logger.info(f"Calling Sarvam-M for structured extraction (document_type={document_type!r})")
    structured_data = await extract_structured_data(document_type, fields, extracted_text)

    return JSONResponse(content={
        "text": extracted_text,
        "data": structured_data,
    })


if __name__ == "__main__":
    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=False)