import os
import io
import re
import json
import zipfile
import asyncio
import subprocess
import tempfile
import threading
from datetime import datetime, timezone
from pathlib import Path

import openpyxl
from openpyxl.styles import Font, PatternFill, Alignment
from huggingface_hub import HfApi, hf_hub_download

from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Security
from fastapi.responses import JSONResponse
from fastapi.security.api_key import APIKeyHeader

import httpx
from sarvamai import SarvamAI

app = FastAPI(
    title="Document Extraction API",
    description="Extract structured data from documents using Sarvam Vision",
    version="3.0.0"
)

# ---------------------------------------------------------------------------
# Supported document types
# ---------------------------------------------------------------------------

SUPPORTED_TYPES = {
    "income_certificate":   "Income certificate issued by a government authority",
    "caste_certificate":    "Caste certificate issued by a government authority",
    "domicile_certificate": "Domicile / residence certificate issued by a government authority",
    "ssc_certificate":      "SSC (Secondary School Certificate) / Class 10 marksheet or passing certificate issued by an education board",
    "hsc_certificate":      "HSC (Higher Secondary Certificate) / Class 12 marksheet or passing certificate issued by an education board",
}

SARVAM_API_KEY  = os.environ.get("SARVAM_API_KEY", "")
API_SECRET_KEY  = os.environ.get("API_SECRET_KEY", "")
HF_TOKEN        = os.environ.get("HF_TOKEN", "")

# ---------------------------------------------------------------------------
# Logging config — edit these two constants to match your HF dataset repo
# ---------------------------------------------------------------------------

LOG_REPO_ID  = "PrathameshRaut/VisionModelLogs"
LOG_FILENAME = "extraction_logs.xlsx"

LOG_HEADERS = [
    "timestamp_utc",
    "filename",
    "file_size_bytes",
    "file_type",
    "document_type",
    "parameters_requested",
    "model_raw_thinking",
    "extracted_json",
    "status",
    "error_detail",
]

# A threading lock so concurrent requests don't corrupt the xlsx
_log_lock = threading.Lock()

# ---------------------------------------------------------------------------
# Auth — X-API-Key header guard
# ---------------------------------------------------------------------------

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

async def verify_api_key(key: str = Security(api_key_header)):
    if not API_SECRET_KEY:
        raise HTTPException(status_code=500, detail="API_SECRET_KEY is not configured on the server")
    if not key or key != API_SECRET_KEY:
        raise HTTPException(
            status_code=401,
            detail="Invalid or missing API key. Send it as request header: X-API-Key: <your-key>"
        )

# ---------------------------------------------------------------------------
# Excel log helpers
# ---------------------------------------------------------------------------

def _style_header_row(ws):
    header_fill = PatternFill("solid", start_color="1F4E79")
    header_font = Font(bold=True, color="FFFFFF", name="Arial", size=10)
    for cell in ws[1]:
        cell.fill = header_fill
        cell.font = header_font
        cell.alignment = Alignment(horizontal="center", vertical="center", wrap_text=True)


def _download_or_create_workbook() -> openpyxl.Workbook:
    """Download the existing log workbook from HF, or create a fresh one."""
    if not HF_TOKEN:
        raise RuntimeError("HF_TOKEN secret is not set — cannot write logs")

    try:
        local_path = hf_hub_download(
            repo_id=LOG_REPO_ID,
            filename=LOG_FILENAME,
            repo_type="dataset",
            token=HF_TOKEN,
            force_download=True,
        )
        wb = openpyxl.load_workbook(local_path)
    except Exception:
        # File doesn't exist yet — start fresh
        wb = openpyxl.Workbook()
        ws = wb.active
        ws.title = "Logs"
        ws.append(LOG_HEADERS)
        _style_header_row(ws)
        col_widths = {
            "A": 22,
            "B": 30,
            "C": 16,
            "D": 14,
            "E": 22,
            "F": 35,
            "G": 60,
            "H": 60,
            "I": 12,
            "J": 40,
        }
        for col, width in col_widths.items():
            ws.column_dimensions[col].width = width
        ws.row_dimensions[1].height = 28
    return wb


def _upload_workbook(wb: openpyxl.Workbook):
    """Save workbook to a buffer and push to HF dataset repo."""
    buf = io.BytesIO()
    wb.save(buf)
    buf.seek(0)

    api = HfApi(token=HF_TOKEN)
    api.upload_file(
        path_or_fileobj=buf,
        path_in_repo=LOG_FILENAME,
        repo_id=LOG_REPO_ID,
        repo_type="dataset",
        commit_message=f"Add log entry {datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S')} UTC",
    )


def append_log(
    filename: str,
    file_size: int,
    file_type: str,
    document_type: str,
    parameters: list[str],
    raw_thinking: str,
    extracted_json: dict | None,
    status: str,
    error_detail: str = "",
):
    """Thread-safe: download → append row → upload."""
    if not HF_TOKEN:
        return

    row = [
        datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S"),
        filename,
        file_size,
        file_type,
        document_type,
        ", ".join(parameters),
        (raw_thinking or "")[:5000],
        json.dumps(extracted_json, ensure_ascii=False) if extracted_json else "",
        status,
        error_detail[:1000],
    ]

    with _log_lock:
        try:
            wb = _download_or_create_workbook()
            ws = wb.active

            ws.append(row)

            data_font = Font(name="Arial", size=10)
            wrap_align = Alignment(vertical="top", wrap_text=True)
            row_idx = ws.max_row
            for cell in ws[row_idx]:
                cell.font = data_font
                cell.alignment = wrap_align

            if row_idx % 2 == 0:
                fill = PatternFill("solid", start_color="DCE6F1")
                for cell in ws[row_idx]:
                    cell.fill = fill

            _upload_workbook(wb)
        except Exception as e:
            print(f"[LOG WARNING] Failed to write log entry: {e}")


# ---------------------------------------------------------------------------
# File conversion — DOCX/PPTX → PDF via LibreOffice; images/PDFs pass through
# ---------------------------------------------------------------------------

NATIVE_EXTENSIONS = {".pdf", ".png", ".jpg", ".jpeg"}
NATIVE_MIME       = {"application/pdf", "image/png", "image/jpeg", "image/jpg"}

def prepare_file(file_bytes: bytes, content_type: str, filename: str) -> tuple[bytes, str]:
    suffix = Path(filename).suffix.lower()
    if suffix in NATIVE_EXTENSIONS or content_type in NATIVE_MIME:
        return file_bytes, filename

    with tempfile.TemporaryDirectory() as tmpdir:
        src = os.path.join(tmpdir, "input" + suffix)
        with open(src, "wb") as f:
            f.write(file_bytes)

        result = subprocess.run(
            ["libreoffice", "--headless", "--convert-to", "pdf", "--outdir", tmpdir, src],
            capture_output=True, timeout=180
        )
        if result.returncode != 0:
            raise HTTPException(
                status_code=500,
                detail=f"File conversion to PDF failed: {result.stderr.decode()}"
            )
        pdfs = [f for f in os.listdir(tmpdir) if f.endswith(".pdf")]
        if not pdfs:
            raise HTTPException(status_code=500, detail="PDF conversion produced no output")
        out_path = os.path.join(tmpdir, pdfs[0])
        with open(out_path, "rb") as f:
            return f.read(), Path(filename).stem + ".pdf"


# ---------------------------------------------------------------------------
# Text cleanup — strip base64-encoded image blobs from OCR output
# ---------------------------------------------------------------------------

def clean_raw_text(text: str) -> str:
    # Pattern 1 — markdown image embed
    text = re.sub(
        r'!\[[^\]]*\]\(data:[a-zA-Z]+/[a-zA-Z+\-]+;base64,[A-Za-z0-9+/=\s]+\)',
        '',
        text,
        flags=re.DOTALL,
    )
    # Pattern 2 — bare data URI
    text = re.sub(
        r'data:[a-zA-Z]+/[a-zA-Z+\-]+;base64,[A-Za-z0-9+/=\s]{100,}',
        '',
        text,
        flags=re.DOTALL,
    )
    return text.strip()


# ---------------------------------------------------------------------------
# Step 1: Extract raw text via Sarvam Document Intelligence SDK
# ---------------------------------------------------------------------------

def extract_text_with_sdk(file_bytes: bytes, filename: str) -> str:
    if not SARVAM_API_KEY:
        raise HTTPException(status_code=500, detail="SARVAM_API_KEY is not configured on the server")

    # FIX: increased timeout so the polling loop doesn't give up too early
    client = SarvamAI(
        api_subscription_key=SARVAM_API_KEY,
        httpx_client=httpx.Client(timeout=180),
    )

    suffix = Path(filename).suffix.lower()
    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
        tmp.write(file_bytes)
        tmp_path = tmp.name

    zip_path = tmp_path + "_output.zip"

    try:
        job = client.document_intelligence.create_job(language="en-IN", output_format="md")
        job.upload_file(tmp_path)
        job.start()
        status = job.wait_until_complete()

        if status.job_state not in ("Completed", "PartiallyCompleted"):
            raise HTTPException(
                status_code=500,
                detail=f"Sarvam document processing ended with state: {status.job_state}"
            )

        job.download_output(zip_path)

        with zipfile.ZipFile(zip_path, "r") as z:
            md_files = sorted(n for n in z.namelist() if n.endswith(".md"))
            if not md_files:
                raise HTTPException(status_code=500, detail="No markdown output found in Sarvam result")
            full_text = "\n\n".join(
                z.read(name).decode("utf-8", errors="replace") for name in md_files
            )

        return full_text

    finally:
        for path in (tmp_path, zip_path):
            try:
                os.unlink(path)
            except FileNotFoundError:
                pass


# ---------------------------------------------------------------------------
# Step 2: Extract structured JSON via SarvamAI SDK chat completions
# ---------------------------------------------------------------------------

def strip_think_tags(text: str) -> str:
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
    text = re.sub(r"<think>.*$", "", text, flags=re.DOTALL)
    return text.strip()


async def extract_json_via_chat(doc_type: str, raw_text: str, parameters: list[str]) -> dict:
    if not SARVAM_API_KEY:
        raise HTTPException(status_code=500, detail="SARVAM_API_KEY is not configured on the server")

    description = SUPPORTED_TYPES[doc_type]
    param_list  = "\n".join(f"  - {p}" for p in parameters)

    # Truncate to avoid blowing context window
    raw_text = raw_text[:6000] if len(raw_text) > 6000 else raw_text

    prompt = f"""You are a document data extraction assistant. Below is the full text extracted from a {description}.
--- DOCUMENT TEXT START ---
{raw_text}
--- DOCUMENT TEXT END ---
Extract the following fields from the document text above:
{param_list}
Rules:
1. Return ONLY a valid JSON object. No explanations, no markdown, no code fences.
2. Use exactly the field names listed above as JSON keys.
3. If a field is not found or not legible, set its value to null.
4. Do not invent or infer values not explicitly present in the document.
5. Respond everything in english language only. Also convert date into proper DD/MM/YYYY format.
Respond with the JSON object only."""

    def _call_sdk() -> str:
        # FIX: increased timeout — sarvam-105b reasoning can take >60s
        client = SarvamAI(
            api_subscription_key=SARVAM_API_KEY,
            httpx_client=httpx.Client(timeout=180),
        )

        # FIX: retry up to 3 times — 105b occasionally returns null content
        for attempt in range(3):
            response = client.chat.completions(
                model="sarvam-105b",
                messages=[{"role": "user", "content": prompt}],
                reasoning_effort=None,
                temperature=0,
                top_p=1,
                max_tokens=4096,
            )

            # FIX: guard against null content before calling .strip()
            content = response.choices[0].message.content
            if content:
                return content.strip()

            print(
                f"[WARN] Attempt {attempt + 1}/3: Sarvam returned null content "
                f"(finish_reason={response.choices[0].finish_reason!r}), retrying..."
            )

        raise ValueError("Sarvam returned null content after 3 attempts")

    try:
        raw = await asyncio.to_thread(_call_sdk)
    except ValueError as exc:
        raise HTTPException(status_code=502, detail=str(exc)) from exc
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"Sarvam chat error: {exc}") from exc

    # Remove <think> blocks
    raw = strip_think_tags(raw)

    # Strip markdown code fences if present
    if raw.startswith("```"):
        raw = raw.split("\n", 1)[-1]
        if raw.endswith("```"):
            raw = raw[: raw.rfind("```")]

    raw = raw.strip()

    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        return {"raw_response": raw, "parse_error": "Model did not return valid JSON"}


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------

@app.get("/")
def root():
    return {
        "message": "Document Extraction API powered by Sarvam Vision",
        "supported_document_types": list(SUPPORTED_TYPES.keys()),
        "docs": "/docs"
    }

@app.get("/types")
def get_supported_types():
    return {"supported_types": {k: {"description": v} for k, v in SUPPORTED_TYPES.items()}}


@app.post("/extract", dependencies=[Security(verify_api_key)])
async def extract_document(
    file: UploadFile = File(..., description="Document file (image, PDF, DOCX, PPTX, etc.)"),
    document_type: str = Form(..., description="One of: income_certificate, caste_certificate, domicile_certificate"),
    parameters: str = Form(..., description="Comma-separated field names to extract. E.g: name,income,date_of_issue")
):
    """
    Extract structured JSON data from an uploaded document.
    Requires header: **X-API-Key: your-secret-key**
    """

    if document_type not in SUPPORTED_TYPES:
        raise HTTPException(
            status_code=400,
            detail={
                "error": f"Unsupported document_type: '{document_type}'",
                "supported_types": list(SUPPORTED_TYPES.keys())
            }
        )

    param_list = [p.strip() for p in parameters.split(",") if p.strip()]
    if not param_list:
        raise HTTPException(
            status_code=400,
            detail="'parameters' must contain at least one field name (e.g. 'name,income,date_of_issue')"
        )

    file_bytes   = await file.read()
    if not file_bytes:
        raise HTTPException(status_code=400, detail="Uploaded file is empty")

    content_type = file.content_type or ""
    filename     = file.filename or "document"
    file_size    = len(file_bytes)

    raw_text  = ""
    extracted = {}

    try:
        processed_bytes, processed_name = prepare_file(file_bytes, content_type, filename)
    except HTTPException as exc:
        asyncio.get_event_loop().run_in_executor(
            None, append_log,
            filename, file_size, content_type, document_type, param_list,
            "", None, "error", str(exc.detail)
        )
        raise

    try:
        raw_text = await asyncio.to_thread(extract_text_with_sdk, processed_bytes, processed_name)
    except HTTPException as exc:
        asyncio.get_event_loop().run_in_executor(
            None, append_log,
            filename, file_size, content_type, document_type, param_list,
            "", None, "error", str(exc.detail)
        )
        raise

    # Strip base64 image blobs before sending to chat model
    raw_text = clean_raw_text(raw_text)

    try:
        extracted = await extract_json_via_chat(document_type, raw_text, param_list)
    except HTTPException as exc:
        asyncio.get_event_loop().run_in_executor(
            None, append_log,
            filename, file_size, content_type, document_type, param_list,
            raw_text, None, "error", str(exc.detail)
        )
        raise

    # Fire-and-forget logging (don't block the response)
    asyncio.get_event_loop().run_in_executor(
        None, append_log,
        filename, file_size, content_type, document_type, param_list,
        raw_text, extracted, "success", ""
    )

    return JSONResponse(content={
        "document_type": document_type,
        "filename": filename,
        "parameters_requested": param_list,
        "extracted_data": extracted
    })


@app.get("/health")
def health():
    return {"status": "ok"}