import os
os.environ["OMP_THREAD_LIMIT"] = "1"
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["OPENBLAS_NUM_THREADS"] = "1"
import re
import json
import concurrent.futures
import csv

import gradio as gr
import cloudinary
import cloudinary.uploader
import pytesseract
from PIL import Image, ImageFilter
from openai import OpenAI
from dotenv import load_dotenv
from pymongo import MongoClient
import pytz
from datetime import datetime

# ──────────────────────────────────────────────
# A.  Setup & Configuration
# ──────────────────────────────────────────────

load_dotenv()

cloudinary.config(
    cloud_name=os.environ.get("CLOUDINARY_CLOUD_NAME"),
    api_key=os.environ.get("CLOUDINARY_API_KEY"),
    api_secret=os.environ.get("CLOUDINARY_API_SECRET"),
)

client = OpenAI(
    base_url="https://integrate.api.nvidia.com/v1",
    api_key=os.environ.get("NVIDIA_API_KEY"),
    timeout=15.0,
)

if os.name == "nt":
    _tess_path = r"C:\Program Files\Tesseract-OCR\tesseract.exe"
    if os.path.exists(_tess_path):
        pytesseract.pytesseract.tesseract_cmd = _tess_path

mongo_uri = os.environ.get("MONGODB_URI")
mongo_client = None
db = None
collection = None

if mongo_uri:
    try:
        mongo_client = MongoClient(mongo_uri, serverSelectionTimeoutMS=5000)
        mongo_client.admin.command("ping")
        db = mongo_client["police_db"]
        collection = db["warrant"]
        print("✅ Connected successfully to MongoDB!")
    except Exception as exc:
        print(f"❌ MongoDB connection failed: {exc}")
        collection = None

officers_db: dict = {}
_officers_csv_error: str | None = None
try:
    with open("officers.csv", "r", encoding="utf-8") as _f:
        for row in csv.DictReader(_f):
            officers_db[row["Officer_Name"]] = row["Phone_Number"]
except Exception as _e:
    _officers_csv_error = str(_e)
    print(f"⚠️ Could not load officers.csv: {_e}")

# ──────────────────────────────────────────────
# B.  Core Processing Logic
# ──────────────────────────────────────────────

SYSTEM_PROMPT = """You are an extremely precise and strict Indian legal document parser.
Your task is to extract information from raw OCR text of a Punjab court warrant or summons.

CRITICAL RULES TO PREVENT HALLUCINATION & FABRICATION:
1. NEVER assume, guess, or fabricate any field. If a field is not explicitly and clearly
   mentioned in the provided text, you MUST return null for that field.
2. DO NOT use placeholder values unless they are literally printed in the text.
3. Case_FIR_Number: Extract the court case number (usually in a format like 'CHI/339/2021' or similar, found near the top-right or case details) AND/OR the FIR number (usually labeled as 'POLICE STATION/FIR NO' or similar).
   - Use " | " separator only if two DISTINCT numbers exist (e.g., 'CHI/339/2021 | PS WOMEN/39/2020').
   - NEVER duplicate the same number on both sides.
   - Ignore or filter out barcode metadata, serial numbers, or form numbers (like 'PBBT03-003502-2021' or 'Form No. 50') if there is a more specific case number (like 'CHI/339/2021') and FIR number present in the text.
4. Act_and_Sections: Extract only explicitly mentioned sections (e.g. "IPC 302"). null if absent.
5. Person_Name_To_Serve: The person to be served/arrested — found after
   "Whereas [NAME] has been duly served ... has failed to attend".
   NEVER use the accused from "Vs [name]" headers.
6. Hearing_Date: The NEXT hearing date only. Format DD-MM-YYYY (e.g. "05-05-2026").
   - You MUST extract the full day, month, AND year. A year-only value like "2026" is WRONG.
   - Look for labels like "NEXT DATE", "Next Date", "Next Hearing", "Date of Hearing".
   - If only a year can be found and no full date exists in the text, return null.
   - NOT the signing/dated date (e.g. ignore "Dated, this day of ...").
7. Court_Name: Extract the specific court designation/level and location (e.g., 'Judicial Magistrate Ist Class, Bathinda' or 'Judicial Magistrate 1st Class' or 'Chief Judicial Magistrate'), NOT the building/header location (such as 'Criminal Courts, Bathinda' or 'District Courts'). The actual court name is typically written directly below 'IN THE COURT OF Sh. [Name]'.
8. Ground every value in the OCR text. Prefer null over a guess.

Return ONLY valid JSON, no markdown fences, no explanation:
{
  "Case_FIR_Number": "...",
  "Act_and_Sections": null,
  "Type_of_Document": "...",
  "Target_Police_Station": "...",
  "IO_Name_and_Belt_No": "...",
  "IO_Mobile_Number": null,
  "Person_Name_To_Serve": "...",
  "Person_Address": "...",
  "Court_Name": "...",
  "Hearing_Date": "..."
}
"""

REQUIRED_KEYS = [
    "Case_FIR_Number", "Act_and_Sections", "Type_of_Document",
    "Target_Police_Station", "IO_Name_and_Belt_No", "IO_Mobile_Number",
    "Person_Name_To_Serve", "Person_Address", "Court_Name", "Hearing_Date",
]

_MONTH_MAP = {
    "january": "01", "february": "02", "march": "03", "april": "04",
    "may": "05", "june": "06", "july": "07", "august": "08",
    "september": "09", "october": "10", "november": "11", "december": "12",
    "jan": "01", "feb": "02", "mar": "03", "apr": "04",
    "jun": "06", "jul": "07", "aug": "08", "sep": "09",
    "oct": "10", "nov": "11", "dec": "12",
}

# ── Detect available Tesseract languages ──────

def _get_available_tess_langs() -> str:
    """
    Query Tesseract for installed language packs and build the best
    available lang string for Punjab court documents (eng + hin + pan).
    Falls back gracefully if hin/pan are not installed.
    """
    try:
        import subprocess
        result = subprocess.run(
            ["tesseract", "--list-langs"],
            capture_output=True, text=True, timeout=10
        )
        installed = set(result.stdout.lower().split() + result.stderr.lower().split())
    except Exception:
        installed = {"eng"}

    langs = ["eng"]
    if "hin" in installed:
        langs.append("hin")
    if "pan" in installed:
        langs.append("pan")

    lang_str = "+".join(langs)
    print(f"🔤 Tesseract languages available: {lang_str}")
    return lang_str


_TESS_LANG = _get_available_tess_langs()

# ── OCR ───────────────────────────────────────

_OCR_MAX_DIM   = 1500   # px — cap before handing to Tesseract (reduced to prevent HF timeout)
_OCR_MIN_DIM   = 1000   # px — upscale target for very small images (optimized for speed)
_OCR_TIMEOUT_S = 60     # seconds — hard kill on the Tesseract subprocess


def _preprocess_for_ocr(img: Image.Image) -> Image.Image:
    """
    Prepare image for Tesseract:
      1. Convert to greyscale (L mode).
      2. Upscale tiny images so Tesseract has enough pixel density (~300 DPI).
      3. **Downscale large images** — phone camera shots can be 4000×3000+ px,
         which causes Tesseract to hang for minutes.  We cap the longest edge at
         _OCR_MAX_DIM (2000 px) which is more than enough for court-document text.
      4. Sharpen after any resize to compensate for interpolation softness.
    """
    img = img.convert("L")
    max_dim = max(img.size)

    if max_dim < _OCR_MIN_DIM:
        # Too small — upscale
        scale = _OCR_MIN_DIM / max_dim
        img = img.resize(
            (int(img.width * scale), int(img.height * scale)), Image.LANCZOS
        )
    elif max_dim > _OCR_MAX_DIM:
        # Too large — downscale to prevent Tesseract timeout
        scale = _OCR_MAX_DIM / max_dim
        img = img.resize(
            (int(img.width * scale), int(img.height * scale)), Image.LANCZOS
        )

    # Sharpen after any resize
    img = img.filter(ImageFilter.SHARPEN)
    return img


def _ocr_image(image_path: str) -> str:
    """
    Run Tesseract via subprocess with a hard OS-level kill on timeout.

    Key fixes vs original:
      • Images are now CAPPED at _OCR_MAX_DIM (2000 px longest edge).
        Phone camera images (4K+) were the primary cause of the hang.
      • Uses subprocess + os.killpg so the Tesseract process is truly killed
        on timeout (ThreadPoolExecutor.cancel() only abandons the Python
        thread — the Tesseract child process kept running and blocked the next
        request).
      • PSM 6 (uniform block) instead of PSM 3 (full auto-layout) — warrants
        are single-column documents; PSM 3's expensive layout analysis adds
        time without improving accuracy here.
      • Falls back to pytesseract if subprocess approach fails (e.g. Windows).
    """
    import subprocess, tempfile, os, signal, sys

    try:
        img = Image.open(image_path)
        processed = _preprocess_for_ocr(img)
    except Exception as exc:
        return f"[OCR Error] Could not open/preprocess image: {exc}"

    # Write the preprocessed image to a temp file for the subprocess approach
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
        tmp_path = tmp.name
    try:
        processed.save(tmp_path, format="PNG")
    except Exception as exc:
        os.unlink(tmp_path)
        return f"[OCR Error] Could not save temp image: {exc}"

    tess_cmd = pytesseract.pytesseract.tesseract_cmd or "tesseract"
    out_base  = tmp_path.replace(".png", "_out")
    cmd = [
        tess_cmd, tmp_path, out_base,
        "-l", _TESS_LANG,
        "--oem", "1",
        "--psm", "4",   # single column of variable sizes — matches Punjab warrants layout best
    ]

    try:
        # Use subprocess with a real timeout so we can kill the process tree
        kwargs = {}
        if sys.platform != "win32":
            kwargs["start_new_session"] = True   # lets us killpg the whole group

        proc = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            **kwargs,
        )
        try:
            proc.communicate(timeout=_OCR_TIMEOUT_S)
        except subprocess.TimeoutExpired:
            if sys.platform != "win32":
                try:
                    os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
                except Exception:
                    proc.kill()
            else:
                proc.kill()
            proc.communicate()
            return "[OCR Timeout] Tesseract exceeded 60 s — image may be too complex."

        out_txt = out_base + ".txt"
        if os.path.exists(out_txt):
            with open(out_txt, "r", encoding="utf-8", errors="replace") as f:
                result = f.read()
            os.unlink(out_txt)
            return result
        # Subprocess produced no output file — fall through to pytesseract fallback
    except Exception:
        pass
    finally:
        try:
            os.unlink(tmp_path)
        except Exception:
            pass

    # ── Fallback: pytesseract (Windows or if subprocess path fails) ────────
    def _run() -> str:
        return pytesseract.image_to_string(
            processed, lang=_TESS_LANG, config="--oem 1 --psm 4"
        )

    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
        future = executor.submit(_run)
        try:
            return future.result(timeout=_OCR_TIMEOUT_S)
        except concurrent.futures.TimeoutError:
            future.cancel()
            return "[OCR Timeout] Tesseract exceeded 60 s — try a smaller or clearer image."
        except Exception as exc:
            return f"[OCR Error] {exc}"

# ── Post-processing helpers ───────────────────

def _extract_person_from_whereas(raw_ocr: str) -> str | None:
    m = re.search(
        r"[Ww]hereas\s+([A-Za-z\s]+(?:,\s*no\.?\s*\d+)?)[,\s]+"
        r"(?:\([^)]*\)[,\s]+)?(?:R/O[^,]*,)?\s*has been duly served",
        raw_ocr,
    )
    return m.group(1).strip() if m else None


def _extract_next_date_from_ocr(raw_ocr: str) -> str | None:
    """
    Fallback: scan OCR text for a date near 'NEXT DATE' / 'Next Date' / 'Next Hearing' labels.
    Returns DD-MM-YYYY string if found, else None.
    """
    # Look for "NEXT DATE" or similar label followed by a date within ~60 characters
    m = re.search(
        r"(?:NEXT\s*DATE|Next\s*Date|Next\s*Hearing|Hearing\s*Date)\s*[:\-]?\s*"
        r"(\d{1,2}[-/.\s]\d{1,2}[-/.\s]\d{4}|\d{4}[-/.\s]\d{1,2}[-/.\s]\d{1,2})",
        raw_ocr, re.IGNORECASE
    )
    if m:
        raw_date = re.sub(r"[\s/.]", "-", m.group(1).strip())
        # Normalise YYYY-MM-DD → DD-MM-YYYY
        yyyy_first = re.fullmatch(r"(\d{4})-(\d{1,2})-(\d{1,2})", raw_date)
        if yyyy_first:
            raw_date = f"{yyyy_first.group(3).zfill(2)}-{yyyy_first.group(2).zfill(2)}-{yyyy_first.group(1)}"
        return raw_date
    return None


def _post_validate(data: dict, raw_ocr: str) -> dict:
    """Correct Person_Name_To_Serve if LLM picked the accused instead of the witness.
    Also repair Hearing_Date if LLM returned only a year (e.g. '2026')."""
    # ── Person fix ──────────────────────────────────────────────────────────
    person = data.get("Person_Name_To_Serve") or ""
    if person and re.search(r"\bVs\s+" + re.escape(person), raw_ocr, re.IGNORECASE):
        fallback = _extract_person_from_whereas(raw_ocr)
        if fallback:
            data["Person_Name_To_Serve"] = fallback
    if not person:
        fallback = _extract_person_from_whereas(raw_ocr)
        if fallback:
            data["Person_Name_To_Serve"] = fallback

    # ── Hearing_Date repair ─────────────────────────────────────────────────
    hdate = data.get("Hearing_Date") or ""
    if hdate and re.fullmatch(r"\d{4}", str(hdate).strip()):
        # LLM returned only a year — try to recover the full date from OCR
        recovered = _extract_next_date_from_ocr(raw_ocr)
        data["Hearing_Date"] = recovered  # may be None if not found
        hdate = recovered or ""

    # Correct OCR character-substitution errors in years (e.g. misreading '2026' as '2028')
    if hdate:
        m_date = re.fullmatch(r"(\d{2})-(\d{2})-(\d{4})", str(hdate).strip())
        if m_date:
            day, month, year = m_date.groups()
            ocr_years = [y for y in re.findall(r"\b(202\d|203\d)\b", raw_ocr)]
            if ocr_years:
                from collections import Counter
                year_counts = Counter(ocr_years)
                most_common_year, count = year_counts.most_common(1)[0]
                if year != most_common_year and year_counts[most_common_year] >= 2 and year_counts[year] <= 1:
                    data["Hearing_Date"] = f"{day}-{month}-{most_common_year}"

    return data


def _is_date_grounded(val: str, raw_ocr_lower: str) -> bool:
    """
    Verify a date string against OCR output.
    Handles numeric dates (DD-MM-YYYY), written months, ordinal suffixes.

    FIX: Reject partial/year-only values like "2026".
    A valid Hearing_Date must contain at least a day, a month, and a year
    (i.e. at least 3 non-empty parts when split on separators).
    Pure 4-digit year strings are always rejected so that a year-only LLM
    output (e.g. "2026") is nulled out rather than passing the grounding check.
    """
    val_str = str(val).strip()

    # Reject bare year (4-digit number only, e.g. "2026")
    if re.fullmatch(r"\d{4}", val_str):
        return False

    # Split into parts and require at least 3 (day / month / year)
    parts = [p for p in re.split(r"[-/.\s]+", val_str) if p]
    if len(parts) < 3:
        return False

    if val_str.lower() in raw_ocr_lower:
        return True

    for part in parts:
        clean = re.sub(r"(?<=\d)(st|nd|rd|th)$", "", part.lower())
        alias = _MONTH_MAP.get(clean)
        if clean not in raw_ocr_lower and (alias is None or alias not in raw_ocr_lower):
            return False
    return True


def _strict_grounding_filter(data: dict, raw_ocr: str) -> dict:
    """
    Programmatic hallucination firewall.
    is_grounded() is nested here so it closes over raw_ocr_lower correctly.
    """
    if not isinstance(data, dict):
        return data

    raw_ocr_lower = raw_ocr.lower()

    def is_grounded(val) -> bool:
        if not val or str(val).strip().lower() in ("null", "none", "—", ""):
            return False
        val_str = str(val).strip()

        # 1. Exact substring
        if val_str.lower() in raw_ocr_lower:
            return True

        # 2. Code-token check: "CHI/339/2021" → any segment present is enough
        code_tokens = [t.lower() for t in re.split(r"[/\-]", val_str) if len(t) > 2]
        if code_tokens and any(t in raw_ocr_lower for t in code_tokens):
            return True

        # 3. Prose token check: require 75% of meaningful words present
        words = [w.lower() for w in re.split(r"[^a-zA-Z0-9]", val_str) if len(w) > 2]
        if not words:
            # Pure numeric fallback (belt numbers, years)
            nums = [n for n in re.split(r"\D+", val_str) if n]
            return any(n in raw_ocr_lower for n in nums) if nums else False

        matched = sum(1 for w in words if w in raw_ocr_lower)
        return matched >= max(1, round(len(words) * 0.75))

    # — Case/FIR deduplication & grounding —
    case_fir = data.get("Case_FIR_Number")
    if case_fir:
        case_fir = str(case_fir).strip()
        if " | " in case_fir:
            parts = list(dict.fromkeys(p.strip() for p in case_fir.split("|") if p.strip()))
            grounded = [p for p in parts if is_grounded(p)]
            data["Case_FIR_Number"] = " | ".join(grounded) if grounded else None
        elif not is_grounded(case_fir):
            data["Case_FIR_Number"] = None

    # — All other fields —
    for field in [
        "Act_and_Sections", "Type_of_Document", "Target_Police_Station",
        "IO_Name_and_Belt_No", "IO_Mobile_Number", "Person_Name_To_Serve",
        "Person_Address", "Court_Name", "Hearing_Date",
    ]:
        val = data.get(field)
        if val:
            if field == "Hearing_Date":
                if not _is_date_grounded(val, raw_ocr_lower):
                    data[field] = None
            elif not is_grounded(val):
                data[field] = None

    return data


def _clean_and_parse_json(raw_response: str, raw_ocr: str = "") -> dict:
    cleaned = re.sub(r"^```(?:json)?\s*", "", raw_response.strip())
    cleaned = re.sub(r"\s*```$", "", cleaned).strip()
    try:
        data = json.loads(cleaned)
        if isinstance(data, dict):
            normalized = {k.replace("__", "_").strip(): v for k, v in data.items()}
            final_data = {
                req: next((v for k, v in normalized.items() if k.lower() == req.lower()), None)
                for req in REQUIRED_KEYS
            }
            if raw_ocr:
                final_data = _post_validate(final_data, raw_ocr)
                final_data = _strict_grounding_filter(final_data, raw_ocr)
            return final_data
        return data
    except json.JSONDecodeError:
        return {
            "_parse_error": True,
            "_raw_llm_response": raw_response,
            "_message": "Could not parse LLM response as JSON.",
        }

# ── Main pipeline ─────────────────────────────

_NVIDIA_MODELS = [
    "qwen/qwen3-coder-480b-a35b-instruct",
    "meta/llama-3.3-70b-instruct",
    "nvidia/llama-3.1-nemotron-70b-instruct",
    "qwen/qwen3.5-122b-a10b",
    "qwen/qwen3-coder-480b-a35b-instruct",
]


def process_document(image_path: str, progress=gr.Progress(track_tqdm=False)):
    if image_path is None:
        raise gr.Error("Please upload an image first.")

    # Step 1 — Upload
    progress(0, desc="☁️ Uploading to Cloudinary…")
    yield "⏳ Uploading to Cloudinary…", "", {}

    try:
        upload_result = cloudinary.uploader.upload(
            image_path, folder="warrants", resource_type="image"
        )
        cloudinary_url = upload_result.get("secure_url", "")
    except Exception as exc:
        raise gr.Error(f"Cloudinary upload failed: {exc}")

    # Step 2 — OCR
    progress(0.25, desc="🔍 Running OCR…")
    yield cloudinary_url, "⏳ Extracting text via OCR…", {}

    try:
        raw_text = _ocr_image(image_path)
    except Exception as exc:
        raw_text = f"[OCR Error] {exc}"

    if not raw_text or not raw_text.strip():
        raw_text = "[OCR returned empty text — image may be blank or unreadable]"

    # Step 3 — LLM
    progress(0.5, desc="🤖 Calling AI model…")
    yield cloudinary_url, raw_text, {"status": "⏳ Calling AI model…"}

    prompt = f"{SYSTEM_PROMPT}\n\n--- RAW OCR TEXT ---\n{raw_text}\n--- END ---"
    llm_response: str | None = None
    last_exception: Exception | None = None

    for model_name in _NVIDIA_MODELS:
        try:
            completion = client.chat.completions.create(
                model=model_name,
                messages=[{"role": "user", "content": prompt}],
                temperature=0.0,
                top_p=0.1,
                max_tokens=1024,
                stream=True,
            )
            chunks = []
            for chunk in completion:
                delta = chunk.choices[0].delta.content if chunk.choices else None
                if delta is not None:
                    chunks.append(delta)
                    yield cloudinary_url, raw_text, {"streaming_raw_response": "".join(chunks)}
            llm_response = "".join(chunks)
            break
        except Exception as exc:
            last_exception = exc
            continue

    if llm_response is None:
        raise gr.Error(f"All NVIDIA models failed. Last error: {last_exception}")

    # Step 4 — Parse + Store
    parsed_json = _clean_and_parse_json(llm_response, raw_text)

    if collection is not None and "_parse_error" not in parsed_json:
        try:
            collection.insert_one({
                **parsed_json,
                "cloudinary_url": cloudinary_url,
                "raw_ocr_text": raw_text,
                "uploaded_at": datetime.now(pytz.timezone("Asia/Kolkata")),
            })
        except Exception as exc:
            print(f"❌ MongoDB insert failed: {exc}")

    progress(1.0, desc="✅ Done!")
    yield cloudinary_url, raw_text, parsed_json


# ──────────────────────────────────────────────
# C.  Dashboard — inline styles so dark theme can't override
# ──────────────────────────────────────────────

C_LABEL       = "#4f46e5"
C_VALUE       = "#1a1a2e"
C_CARD_BG     = "#ffffff"
C_CARD_BG2    = "#f8f7ff"
C_CARD_BORDER = "#e0e7ff"
C_PILL_BG     = "#ede9fe"
C_PILL_TEXT   = "#3730a3"
C_SEPARATOR   = "#ede9fe"
C_MUTED       = "#6b7280"
C_BTN_BG      = "#4f46e5"
C_BTN_TEXT    = "#ffffff"
C_TH_BG1      = "#4f46e5"
C_TH_BG2      = "#7c3aed"
C_TH_TEXT     = "#ffffff"


def _card_row(label: str, value: str, is_pill: bool = False, is_btn: bool = False) -> str:
    label_html = (
        f'<span style="flex:0 0 44%;max-width:44%;font-weight:700;font-size:0.73rem;'
        f'color:{C_LABEL};white-space:nowrap;overflow:hidden;text-overflow:ellipsis;">'
        f'{label}</span>'
    )
    if is_btn:
        value_html = (
            f'<a href="{value}" target="_blank" rel="noopener" '
            f'style="flex:1;display:flex;align-items:center;justify-content:center;'
            f'padding:7px 0;background:{C_BTN_BG};color:{C_BTN_TEXT};border-radius:8px;'
            f'font-weight:600;font-size:0.82rem;text-decoration:none;">🖼 View</a>'
        )
    elif is_pill:
        value_html = (
            f'<span style="flex:1;text-align:right;">'
            f'<span style="display:inline-block;padding:3px 10px;border-radius:999px;'
            f'background:{C_PILL_BG};color:{C_PILL_TEXT};font-size:0.75rem;font-weight:700;">'
            f'{value}</span></span>'
        )
    elif value and value != "—":
        value_html = (
            f'<span style="flex:1;text-align:right;color:{C_VALUE};'
            f'font-size:0.82rem;word-break:break-word;">{value}</span>'
        )
    else:
        value_html = (
            f'<span style="flex:1;text-align:right;color:{C_MUTED};font-size:0.82rem;">—</span>'
        )

    return (
        f'<div style="display:flex;align-items:flex-start;justify-content:space-between;'
        f'gap:8px;padding:8px 14px;border-bottom:1px solid {C_SEPARATOR};">'
        f'{label_html}{value_html}</div>'
    )


def _build_html_table(rows: list) -> str:
    if not rows:
        return (
            f'<div style="text-align:center;padding:32px;color:{C_MUTED};'
            f'font-size:0.95rem;font-family:Segoe UI,sans-serif;">📭 No records found.</div>'
        )

    headers = [
        "Uploaded At", "Case / FIR No.", "Type", "Target Station",
        "IO Name & Belt No.", "IO Mobile", "Person to Serve",
        "Address", "Court", "Hearing Date",
    ]
    th = f'padding:10px 12px;text-align:left;white-space:nowrap;font-weight:600;color:{C_TH_TEXT};font-size:0.82rem;'
    header_html = "".join(f'<th style="{th}">{h}</th>' for h in headers)
    header_html += f'<th style="{th}">Document</th>'

    desktop_rows = ""
    for i, row in enumerate(rows):
        url = row[-1] if row[-1] else ""
        bg = C_CARD_BG2 if i % 2 else C_CARD_BG
        cells = ""
        for j, cell in enumerate(row[:-1]):
            val = str(cell) if cell else "—"
            if j == 1:
                cells += (
                    f'<td style="padding:8px 12px;vertical-align:top;max-width:180px;word-break:break-word;">'
                    f'<span style="display:inline-block;padding:2px 8px;border-radius:999px;'
                    f'background:{C_PILL_BG};color:{C_PILL_TEXT};font-size:0.75rem;font-weight:700;">'
                    f'{val}</span></td>'
                )
            else:
                cells += (
                    f'<td style="padding:8px 12px;vertical-align:top;color:{C_VALUE};'
                    f'max-width:180px;word-break:break-word;font-size:0.82rem;">{val}</td>'
                )
        link_cell = (
            f'<td style="padding:8px 12px;vertical-align:top;">'
            f'<a href="{url}" target="_blank" rel="noopener" '
            f'style="display:inline-flex;align-items:center;gap:4px;padding:4px 10px;'
            f'border-radius:6px;background:{C_BTN_BG};color:{C_BTN_TEXT};'
            f'font-size:0.75rem;font-weight:600;text-decoration:none;white-space:nowrap;">🖼 View</a></td>'
            if url else f'<td style="padding:8px 12px;color:{C_MUTED};">—</td>'
        )
        desktop_rows += (
            f'<tr style="background:{bg};border-bottom:1px solid {C_CARD_BORDER};">'
            f'{cells}{link_cell}</tr>'
        )

    desktop_table = (
        f'<div class="warrant-desktop" style="overflow-x:auto;border-radius:10px;'
        f'box-shadow:0 2px 12px rgba(0,0,0,0.10);margin-top:12px;">'
        f'<table style="width:100%;border-collapse:collapse;font-family:Segoe UI,sans-serif;">'
        f'<thead><tr style="background:linear-gradient(90deg,{C_TH_BG1},{C_TH_BG2});">'
        f'{header_html}</tr></thead><tbody>{desktop_rows}</tbody></table></div>'
    )

    mobile_cards = '<div class="warrant-mobile">'
    for i, row in enumerate(rows):
        url = row[-1] if row[-1] else ""
        bg = C_CARD_BG2 if i % 2 else C_CARD_BG
        values = [str(c) if c else "—" for c in row[:-1]]
        card_rows = ""
        for j, (label, val) in enumerate(zip(headers, values)):
            row_html = _card_row(label, val, is_pill=(j == 1))
            if j == len(headers) - 1:
                row_html = row_html.replace(f"border-bottom:1px solid {C_SEPARATOR};", "border-bottom:none;")
            card_rows += row_html
        if url:
            card_rows += (
                '<div style="padding:10px 14px;">'
                + _card_row("Document", url, is_btn=True)
                  .replace(f"border-bottom:1px solid {C_SEPARATOR};", "border-bottom:none;")
                + "</div>"
            )
        mobile_cards += (
            f'<div style="background:{bg};border:1px solid {C_CARD_BORDER};'
            f'border-radius:12px;margin-bottom:14px;overflow:hidden;'
            f'box-shadow:0 2px 8px rgba(79,70,229,0.08);">{card_rows}</div>'
        )
    mobile_cards += "</div>"
    return desktop_table + mobile_cards


def fetch_live_warrants(search_query: str = "") -> str:
    if search_query is None:
        search_query = ""
    if collection is None:
        return (
            f'<div style="text-align:center;padding:32px;color:{C_MUTED};'
            f'font-family:Segoe UI,sans-serif;">⚠️ Database connection not available.</div>'
        )
    query: dict = {}
    if search_query.strip():
        rgx = {"$regex": search_query.strip(), "$options": "i"}
        query = {"$or": [
            {"Case_FIR_Number": rgx}, {"Type_of_Document": rgx},
            {"Target_Police_Station": rgx}, {"IO_Name_and_Belt_No": rgx},
            {"Person_Name_To_Serve": rgx}, {"Court_Name": rgx},
        ]}
    try:
        IST = pytz.timezone("Asia/Kolkata")
        rows = []
        for item in collection.find(query).sort("uploaded_at", -1):
            uploaded_str = ""
            if "uploaded_at" in item:
                dt = item["uploaded_at"]
                if dt.tzinfo is None:
                    dt = pytz.utc.localize(dt)
                uploaded_str = dt.astimezone(IST).strftime("%Y-%m-%d %H:%M")
            rows.append([
                uploaded_str,
                item.get("Case_FIR_Number") or "",
                item.get("Type_of_Document") or "",
                item.get("Target_Police_Station") or "",
                item.get("IO_Name_and_Belt_No") or "",
                item.get("IO_Mobile_Number") or "",
                item.get("Person_Name_To_Serve") or "",
                item.get("Person_Address") or "",
                item.get("Court_Name") or "",
                item.get("Hearing_Date") or "",
                item.get("cloudinary_url") or "",
            ])
        return _build_html_table(rows)
    except Exception as exc:
        return (
            f'<div style="text-align:center;padding:32px;color:#dc2626;'
            f'font-family:Segoe UI,sans-serif;">❌ Error fetching data: {exc}</div>'
        )


# ──────────────────────────────────────────────
# D.  CSS + JS
# ──────────────────────────────────────────────

TAB_FIX_JS = """<script>
(function(){
  function fix(){
    ['[role="tablist"]','.tab-nav','.tab-nav > div','.tab-nav > div > div'].forEach(function(s){
      document.querySelectorAll(s).forEach(function(el){
        el.style.cssText+=';display:flex!important;flex-direction:row!important;'+
          'flex-wrap:nowrap!important;overflow-x:auto!important;overflow-y:visible!important;'+
          '-webkit-overflow-scrolling:touch!important;';
      });
    });
    document.querySelectorAll('[role="tab"],.tab-nav button').forEach(function(btn){
      var n=btn.parentElement,d=0;
      while(n&&d<12){
        var s=window.getComputedStyle(n);
        if(s.overflow==='hidden'||s.overflowX==='hidden'){n.style.overflow='visible';n.style.overflowX='auto';}
        n=n.parentElement;d++;
      }
      btn.style.cssText+=';flex-shrink:0!important;white-space:nowrap!important;'+
        'pointer-events:auto!important;touch-action:manipulation!important;';
    });
  }
  fix();setTimeout(fix,300);setTimeout(fix,800);setTimeout(fix,2000);
  if(window.MutationObserver){new MutationObserver(fix).observe(document.body,{childList:true,subtree:true});}
})();
</script>"""

CUSTOM_CSS = """
*,*::before,*::after{box-sizing:border-box!important;}
body,html{overflow-x:hidden!important;max-width:100vw!important;}
.gradio-container{max-width:1280px!important;margin:auto!important;padding:0 12px!important;}

.tabs,div[class*="tabs"],div[data-testid="tabs"],
.tabitem>.block,.tabitem>div>.block{overflow:visible!important;}

.tab-nav,.tab-nav>div,.tab-nav>div>div,
[role="tablist"],div[class*="tab-nav"],div[data-testid="tab-nav"]{
    display:flex!important;flex-direction:row!important;flex-wrap:nowrap!important;
    overflow-x:auto!important;overflow-y:visible!important;
    -webkit-overflow-scrolling:touch!important;gap:4px!important;
    scrollbar-width:none!important;-ms-overflow-style:none!important;
}
.tab-nav::-webkit-scrollbar,.tab-nav>div::-webkit-scrollbar,
[role="tablist"]::-webkit-scrollbar{display:none!important;}

[role="tab"],.tab-nav button{
    flex-shrink:0!important;white-space:nowrap!important;
    min-width:max-content!important;pointer-events:auto!important;
    touch-action:manipulation!important;cursor:pointer!important;
}

#process-btn{font-size:1rem;padding:12px 24px;width:100%;margin-top:8px;}
#status-row{background:#f0fdf4;border-radius:8px;padding:8px 14px;font-size:0.85rem;}
#status-row,#status-row *{color:#166534!important;}
.upload-col{min-width:0!important;flex:1 1 280px!important;}
.outputs-col{min-width:0!important;flex:2 1 380px!important;}

.warrant-desktop{display:block;}
.warrant-mobile{display:none;}

@media screen and (max-width:768px){
    .gradio-container,.main,.wrap,.tabitem,footer{overflow-x:hidden!important;max-width:100%!important;}

    .gradio-container div.flex,.gradio-container div.gap,
    .gradio-container .gr-row,.gradio-container [class*="flex-row"],
    .gradio-container form>div{flex-direction:column!important;flex-wrap:nowrap!important;}

    .gradio-container div.flex>*,.gradio-container div.gap>*,
    .gradio-container .gr-row>*,.upload-col,.outputs-col,
    .gradio-container .block,.gradio-container .col,
    .gradio-container [data-testid="column"]{
        width:100%!important;max-width:100%!important;min-width:0!important;flex:none!important;
    }
    .gradio-container [data-testid="image"],
    .gradio-container .upload-container{width:100%!important;height:220px!important;}
    .gradio-container textarea,
    .gradio-container input[type="text"]{width:100%!important;}

    #search-refresh-row,#search-refresh-row>*{
        flex-direction:column!important;width:100%!important;min-width:0!important;flex:none!important;
    }
    #refresh-btn{width:100%!important;margin-top:6px;}

    .warrant-desktop{display:none!important;}
    .warrant-mobile{display:block!important;}
}
"""

DESCRIPTION = """
Upload a photo of a **bailable warrant**, **summon**, or similar legal document.

| Step | Action |
|------|--------|
| ☁️ 1 | Host the image on **Cloudinary** |
| 🔍 2 | Extract raw text via **Tesseract OCR** |
| 🤖 3 | Parse structured fields using **NVIDIA LLM Suite** |
| 🗄️ 4 | Store the record securely in **MongoDB** |
"""

# ──────────────────────────────────────────────
# E.  Gradio Interface
# ──────────────────────────────────────────────

def _status_html(icon: str, message: str, color: str, done: bool = False) -> str:
    if done:
        return (
            '<div style="display:flex;align-items:center;gap:10px;padding:10px 14px;'
            'background:#16a34a18;border-left:4px solid #16a34a;border-radius:8px;">'
            '<span style="font-size:1.4rem">✅</span>'
            '<div style="font-weight:600;color:#16a34a">Processing complete!</div></div>'
        )
    return (
        f'<div style="display:flex;align-items:center;gap:10px;padding:10px 14px;'
        f'background:{color}18;border-left:4px solid {color};border-radius:8px;margin-bottom:8px;">'
        f'<span style="font-size:1.4rem;line-height:1">{icon}</span>'
        f'<div><div style="font-weight:600;color:{color};font-size:.9rem">{message}</div>'
        f'<div style="font-size:.75rem;color:#6b7280;margin-top:2px">Please wait, do not close this tab</div></div>'
        f'<div style="margin-left:auto;width:20px;height:20px;border:3px solid {color}40;'
        f'border-top-color:{color};border-radius:50%;animation:spin 1s linear infinite"></div></div>'
        f'<style>@keyframes spin{{to{{transform:rotate(360deg)}}}}</style>'
    )


def _process_and_render(image_path):
    for url, ocr, data in process_document(image_path):
        # Determine status banner
        if not url.startswith("http"):
            s = _status_html("☁️", "Uploading image to Cloudinary…", "#6366f1")
        elif ocr == "⏳ Extracting text via OCR…":
            s = _status_html("🔍", "Running OCR — extracting text from image…", "#0891b2")
        elif isinstance(data, dict) and "status" in data:
            s = _status_html("🤖", "AI model is parsing the document fields…", "#7c3aed")
        elif isinstance(data, dict) and "streaming_raw_response" in data:
            n = len(data["streaming_raw_response"])
            s = _status_html("🤖", f"AI model streaming… ({n} chars)", "#7c3aed")
        elif isinstance(data, dict) and any(k in data for k in ["Case_FIR_Number", "_parse_error"]):
            s = _status_html("", "", "", done=True)
        else:
            s = _status_html("⏳", "Processing…", "#6b7280")

        link_html = ""
        if url and url.startswith("http"):
            link_html = (
                f'<a href="{url}" target="_blank" rel="noopener" '
                f'style="display:inline-flex;align-items:center;gap:6px;padding:6px 14px;'
                f'background:#4f46e5;color:#fff;border-radius:7px;font-weight:600;'
                f'text-decoration:none;font-size:.85rem;">🖼 Open on Cloudinary ↗</a>'
            )
        yield s, url, link_html, ocr, data


_WA_JS = f"""
(phone, url, data) => {{
    if (!phone) return '<span style="color:#dc2626">❌ Please enter a WhatsApp number.</span>';
    if (!url)   return '<span style="color:#dc2626">❌ No document uploaded yet.</span>';
    const caseNo = data?.Case_FIR_Number || "Unknown Case";
    const court  = data?.Court_Name      || "Unknown Court";
    const text   = `🚨 *New Warrant Uploaded*\\n*Case:* ${{caseNo}}\\n*Court:* ${{court}}\\n*Document:* ${{url}}`;
    const clean  = phone.replace(/[^0-9]/g, '');
    if (!clean) return '<span style="color:#dc2626">❌ Invalid phone number.</span>';
    window.open(`https://wa.me/${{clean}}?text=${{encodeURIComponent(text)}}`, '_blank');
    return '<span style="color:#16a34a">✅ WhatsApp opened — click Send in the app.</span>';
}}
"""

_PHONE_JS = f"(name) => {{ const db = {json.dumps(officers_db)}; return db[name] || ''; }}"

with gr.Blocks(
    title="⚖️ Legal Document Digitization",
    theme=gr.themes.Soft(primary_hue="violet", secondary_hue="indigo", neutral_hue="slate"),
    css=CUSTOM_CSS,
) as demo:

    gr.HTML(TAB_FIX_JS)  # must be first child — patches tab bar before render

    gr.Markdown("# ⚖️ Automated Legal Document Digitization System")
    gr.Markdown("*Digitize warrants & summons in seconds — OCR → AI parsing → secure storage*")

    with gr.Tabs():

        # ── Tab 1: Pipeline ────────────────────────────────────
        with gr.Tab("📥 Digitization Pipeline"):
            gr.Markdown(DESCRIPTION)

            with gr.Row():
                with gr.Column(elem_classes=["upload-col"]):
                    image_input = gr.Image(
                        type="filepath",
                        label="📎 Upload Warrant / Summon Photo",
                        height=300,
                    )
                    submit_btn = gr.Button(
                        "🚀 Process Document",
                        variant="primary",
                        elem_id="process-btn",
                        size="lg",
                    )
                    gr.Markdown(
                        "**Tip:** Use a clear, well-lit photo for best OCR accuracy.",
                        elem_id="status-row",
                    )
                    if _officers_csv_error:
                        gr.Markdown(
                            f"⚠️ **Officer CSV not loaded** — WhatsApp dropdown will be empty. "
                            f"Ensure `officers.csv` has `Officer_Name` and `Phone_Number` columns. "
                            f"_(Error: `{_officers_csv_error}`)_"
                        )

                with gr.Column(elem_classes=["outputs-col"]):
                    status_out           = gr.HTML(value="", elem_id="status-display")
                    cloudinary_url_out   = gr.Textbox(label="☁️ Cloudinary URL", interactive=False)
                    cloudinary_link_html = gr.HTML()
                    raw_ocr_out          = gr.Textbox(label="🔍 Raw OCR Text", lines=8, interactive=False)
                    json_out             = gr.JSON(label="📋 Extracted Structured Data (JSON)")

                    gr.Markdown("### 📨 Notify Investigating Officer (WhatsApp)")
                    with gr.Row():
                        io_dropdown     = gr.Dropdown(
                            label="Select Officer (from CSV)",
                            choices=list(officers_db.keys()),
                            scale=2,
                        )
                        manual_phone_in = gr.Textbox(
                            label="WhatsApp Mobile No.",
                            placeholder="e.g. 919876543210",
                            scale=2,
                        )
                        send_wa_btn = gr.Button("💬 Send via WhatsApp", variant="secondary", scale=1)
                    wa_status_out = gr.HTML()

            # Wire events — defined after all components exist
            io_dropdown.change(
                fn=None, inputs=[io_dropdown], outputs=[manual_phone_in], js=_PHONE_JS
            )
            submit_btn.click(
                fn=_process_and_render,
                inputs=[image_input],
                outputs=[status_out, cloudinary_url_out, cloudinary_link_html, raw_ocr_out, json_out],
            )
            send_wa_btn.click(
                fn=None,
                inputs=[manual_phone_in, cloudinary_url_out, json_out],
                outputs=[wa_status_out],
                js=_WA_JS,
            )

        # ── Tab 2: Dashboard ───────────────────────────────────
        with gr.Tab("👮 Live Police Dashboard"):
            gr.Markdown("## 📋 Real-Time Stored Warrants & Summons")
            gr.Markdown(
                "Browse and search all digitized legal documents stored in MongoDB. "
                "Click **View** in the *Document* column to open the original image."
            )

            with gr.Row(elem_id="search-refresh-row"):
                search_box = gr.Textbox(
                    placeholder="🔍 Search by Case No., IO Name, Person, or Station…",
                    show_label=False,
                    scale=4,
                )
                refresh_btn = gr.Button(
                    "🔄 Refresh", variant="secondary", scale=1, elem_id="refresh-btn"
                )

            dashboard_html = gr.HTML(
                value=(
                    "<div style='text-align:center;padding:32px;color:#6b7280;"
                    "font-family:Segoe UI,sans-serif;'>Click 🔄 Refresh to load records.</div>"
                )
            )

            search_box.change(fn=fetch_live_warrants, inputs=[search_box], outputs=[dashboard_html])
            refresh_btn.click(fn=fetch_live_warrants, inputs=[search_box], outputs=[dashboard_html])


if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        ssr_mode=False,
    )