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",
}

def _get_available_tess_langs() -> str:
    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_MAX_DIM   = 1500
_OCR_MIN_DIM   = 1000
_OCR_TIMEOUT_S = 60


def _preprocess_for_ocr(img: Image.Image) -> Image.Image:
    img = img.convert("L")
    max_dim = max(img.size)
    if max_dim < _OCR_MIN_DIM:
        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:
        scale = _OCR_MAX_DIM / max_dim
        img = img.resize((int(img.width * scale), int(img.height * scale)), Image.LANCZOS)
    img = img.filter(ImageFilter.SHARPEN)
    return img


def _ocr_image(image_path: str) -> str:
    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}"

    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"]

    try:
        kwargs = {}
        if sys.platform != "win32":
            kwargs["start_new_session"] = True
        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
    except Exception:
        pass
    finally:
        try:
            os.unlink(tmp_path)
        except Exception:
            pass

    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}"


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:
    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())
        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:
    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

    hdate = data.get("Hearing_Date") or ""
    if hdate and re.fullmatch(r"\d{4}", str(hdate).strip()):
        recovered = _extract_next_date_from_ocr(raw_ocr)
        data["Hearing_Date"] = recovered
        hdate = recovered or ""

    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:
    val_str = str(val).strip()
    if re.fullmatch(r"\d{4}", val_str):
        return False
    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:
    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()
        if val_str.lower() in raw_ocr_lower:
            return True
        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
        words = [w.lower() for w in re.split(r"[^a-zA-Z0-9]", val_str) if len(w) > 2]
        if not words:
            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 = 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

    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.",
        }


_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.")

    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}")

    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]"

    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}")

    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 helpers
# ──────────────────────────────────────────────

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(){
  // ── 1. Hide the Hugging Face floating repo/files button ─────────────────
  // HF injects a fixed-position pill (LitigationBrach / ALDDS) into the
  // parent page. We can't remove it from Python, so we hide it with CSS
  // injected into <head> and also poll+hide any matching DOM nodes.
  var style = document.createElement('style');
  style.textContent = [
    /* The HF "Files" / repo pill — fixed position, top-right corner */
    'div[class*="SpaceIframeBadge"],',
    'div[class*="space-badge"],',
    'div[class*="iframe-badge"],',
    'a[class*="SpaceIframeBadge"],',
    'a[class*="space-badge"],',
    /* Catch-all: any fixed/absolute element in top-right with HF branding */
    'body > div[style*="position: fixed"][style*="top"][style*="right"],',
    'body > div[style*="position:fixed"][style*="top"][style*="right"]',
    '{ display:none!important; visibility:hidden!important; pointer-events:none!important; }'
  ].join('');
  document.head.appendChild(style);

  // Also poll the DOM and forcibly hide any element that looks like the badge
  function hideBadge(){
    document.querySelectorAll('body > div, body > a').forEach(function(el){
      var cs = window.getComputedStyle(el);
      if((cs.position==='fixed'||cs.position==='absolute') &&
         parseInt(cs.top)<80 && parseInt(cs.right)<200){
        var txt = el.innerText || el.textContent || '';
        // Only hide if it contains HF-style content (heart/like icon or repo name)
        if(txt.includes('♥')||txt.includes('❤')||el.querySelector('svg')||
           el.innerHTML.includes('hf.co')||el.innerHTML.includes('huggingface')||
           el.innerHTML.includes('LitigationBrach')||el.innerHTML.includes('ALDDS')){
          el.style.cssText += ';display:none!important;visibility:hidden!important;';
        }
      }
    });
  }
  hideBadge();
  setTimeout(hideBadge,300);
  setTimeout(hideBadge,800);
  setTimeout(hideBadge,2000);
  setTimeout(hideBadge,5000);
  if(window.MutationObserver){
    new MutationObserver(hideBadge).observe(document.body,{childList:true,subtree:false});
  }

  // ── 2. Fix tab bar horizontal layout ────────────────────────────────────
  function fixTabs(){
    ['[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;';
    });
  }
  fixTabs();
  setTimeout(fixTabs,300);
  setTimeout(fixTabs,800);
  setTimeout(fixTabs,2000);
  if(window.MutationObserver){
    new MutationObserver(fixTabs).observe(document.body,{childList:true,subtree:true});
  }
})();
</script>"""

CUSTOM_CSS = """
/* ── Reset & base ─────────────────────────────────────────────────────── */
*,*::before,*::after { box-sizing:border-box!important; }
html,body { overflow-x:hidden!important; max-width:100vw!important; margin:0; padding:0; }

/* ── Container ────────────────────────────────────────────────────────── */
.gradio-container {
    max-width:1280px!important;
    margin:auto!important;
    padding:0 12px!important;
}

/* ── Hide HF badge via CSS (belt-and-suspenders with the JS above) ─────── */
div[class*="SpaceIframeBadge"],
div[class*="space-badge"],
div[class*="iframe-badge"],
a[class*="SpaceIframeBadge"],
a[class*="space-badge"] {
    display:none!important;
    visibility:hidden!important;
    pointer-events:none!important;
}

/* ── Tab bar ──────────────────────────────────────────────────────────── */
.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 button ───────────────────────────────────────────────────── */
#process-btn { font-size:1rem; padding:12px 24px; width:100%; margin-top:8px; }

/* ── Status row hint ─────────────────────────────────────────────────── */
#status-row { background:#f0fdf4; border-radius:8px; padding:8px 14px; font-size:0.85rem; }
#status-row, #status-row * { color:#166534!important; }

/* ── Column sizing (desktop) ─────────────────────────────────────────── */
.upload-col  { min-width:0!important; flex:1 1 280px!important; }
.outputs-col { min-width:0!important; flex:2 1 380px!important; }

/* ── Dashboard table / cards ─────────────────────────────────────────── */
.warrant-desktop { display:block; }
.warrant-mobile  { display:none;  }

/* ══════════════════════════════════════════════════════════════════════
   MOBILE  ≤ 768 px
   ══════════════════════════════════════════════════════════════════════ */
@media screen and (max-width:768px){

    /* Prevent any horizontal scroll on the whole page */
    .gradio-container,.main,.wrap,.tabitem,footer {
        overflow-x:hidden!important;
        max-width:100%!important;
    }

    /* Stack ALL flex rows into columns */
    .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;
    }

    /* Full-width every child column/block */
    .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;
    }

    /* Image uploader height */
    .gradio-container [data-testid="image"],
    .gradio-container .upload-container {
        width:100%!important;
        height:220px!important;
    }

    /* Inputs full width */
    .gradio-container textarea,
    .gradio-container input[type="text"] { width:100%!important; }

    /* Search + refresh row */
    #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; }

    /* WhatsApp row — stack dropdowns/inputs vertically */
    .wa-notify-row,
    .wa-notify-row>* {
        flex-direction:column!important;
        width:100%!important;
        min-width:0!important;
        flex:none!important;
    }
    .wa-notify-row .gradio-dropdown,
    .wa-notify-row .gradio-textbox,
    .wa-notify-row button {
        width:100%!important;
        min-width:0!important;
    }

    /* Switch dashboard to card layout on mobile */
    .warrant-desktop { display:none!important; }
    .warrant-mobile  { display:block!important; }

    /* Shrink heading text slightly so it doesn't overflow */
    .gradio-container h1 { font-size:1.35rem!important; }
    .gradio-container h2 { font-size:1.1rem!important; }

    /* Make JSON output scrollable instead of overflowing */
    .gradio-container .json-holder,
    .gradio-container [data-testid="json"] {
        overflow-x:auto!important;
        max-width:100%!important;
    }

    /* Process button — comfortable tap target */
    #process-btn { font-size:1rem!important; padding:14px 20px!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** |
"""


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):
        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] || ''; }}"

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

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:

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

    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():
                # ── Left column: upload + controls ─────────────────────
                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}`)_"
                        )

                # ── Right column: outputs ──────────────────────────────
                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)")
                    # elem_classes used so our mobile CSS can target this row
                    with gr.Row(elem_id="wa-notify-row", elem_classes=["wa-notify-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
            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,
    )