import os
import sys
import subprocess

# ──────────────────────────────────────────────
# 0. Auto-fix PaddlePaddle 3.0+ PIR Crash
# ──────────────────────────────────────────────
if os.environ.get("_PADDLE_VER_FIXED") != "1":
    try:
        import paddleocr
        if paddleocr.__version__.startswith(("3.", "2.10")):
            print("⏳ Auto-fixing PaddleOCR (downgrading to stable v2.9.1 to fix CPU crash)...")
            subprocess.check_call([sys.executable, "-m", "pip", "install", 
                                   "paddleocr==2.9.1", "paddlepaddle==2.6.2", "-q"])
            os.environ["_PADDLE_VER_FIXED"] = "1"
            print("🔄 Restarting process with fixed versions...")
            os.execv(sys.executable, [sys.executable] + sys.argv)
    except Exception as e:
        print(f"⚠️ PaddleOCR auto-fix skipped: {e}")

# Thread limits for CPU stability
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 csv

import gradio as gr
import cloudinary
import cloudinary.uploader
import numpy as np
from PIL import Image, ImageFilter
from paddleocr import PaddleOCR
from dotenv import load_dotenv
from pymongo import MongoClient
import pytz
from datetime import datetime

# Local LLM dependencies
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

# ──────────────────────────────────────────────
# 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"),
)

# Initialize PaddleOCR (v2.9.1 stable syntax)
paddle_ocr = PaddleOCR(use_angle_cls=True, lang='en', show_log=False)

# Initialize Local LLM (Qwen2.5-1.5B - Ultra-fast for HF Free Tier 2 vCPUs)
print("📥 Loading Ultra-Fast Local LLM (Qwen2.5-1.5B-Instruct)...")
model_path = hf_hub_download(
    repo_id="bartowski/Qwen2.5-1.5B-Instruct-GGUF",
    filename="Qwen2.5-1.5B-Instruct-Q4_K_M.gguf"
)

has_gpu = bool(os.environ.get("CUDA_VISIBLE_DEVICES", "").strip())

local_llm = Llama(
    model_path=model_path,
    n_ctx=2048,                        # Increased to fit prompt + full JSON without truncating
    n_threads=min(4, os.cpu_count() or 2), # Capped specifically for HF Free Tier (2 vCPUs)
    n_gpu_layers=-1 if has_gpu else 0,      
    n_batch=512,                       # Safe batch size for 2 CPUs
    use_mlock=True,                    # Prevents swapping to disk
    verbose=False
)
print("✅ Local LLM Loaded Successfully!")

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 AND/OR the FIR number.
   - Use " | " separator only if two DISTINCT numbers exist.
   - Ignore or filter out barcode metadata, serial numbers, or form numbers.
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".
7. Court_Name: Extract the specific court designation/level and location.
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",
}

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

_OCR_MAX_DIM   = 1500   # px — cap before handing to PaddleOCR
_OCR_MIN_DIM   = 1000   # px — upscale target for very small images


def _preprocess_for_ocr(img: Image.Image) -> Image.Image:
    img = img.convert("RGB")
    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:
    try:
        img = Image.open(image_path)
        processed = _preprocess_for_ocr(img)
        img_array = np.array(processed)
        
        # v2.9.1 stable syntax
        result = paddle_ocr.ocr(img_array, cls=True)
        
        if not result or not result[0]:
            return "[OCR returned empty text — image may be blank or unreadable]"
            
        lines = []
        for line in result[0]:
            text = line[1][0]
            lines.append(text)
            
        return "\n".join(lines)
    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:
    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:
    # Used \x60 to prevent markdown parser truncation bug
    cleaned = re.sub(r"^\x60\x60\x60(?:json)?\s*", "", raw_response.strip())
    cleaned = re.sub(r"\s*\x60\x60\x60$", "", 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 ─────────────────────────────

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 Local AI model…")
    yield cloudinary_url, raw_text, {"status": "⏳ Calling Local AI model…"}

    prompt = f"--- RAW OCR TEXT ---\n{raw_text}\n--- END ---"
    
    try:
        # stream=False runs significantly faster on limited CPUs (HF Free Tier)
        response = local_llm.create_chat_completion(
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": prompt}
            ],
            temperature=0.1,        
            top_p=0.1,
            max_tokens=1024,        # Plenty of room to finish the JSON
            stop=["\n\n", "```"],   
            stream=False, 
        )
        llm_response = response["choices"][0]["message"]["content"]
        yield cloudinary_url, raw_text, {"streaming_raw_response": llm_response}
    except Exception as exc:
        raise gr.Error(f"Local LLM processing failed: {exc}")

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

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 **PaddleOCR** |
| 🤖 3 | Parse structured fields using **Local Qwen2.5-1.5B LLM** |
| 🗄️ 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):
        if not url.startswith("http"):
            s = _status_html("☁️", "Uploading image to Cloudinary…", "#6366f1")
        elif ocr == "⏳ Extracting text via OCR…":
            s = _status_html("🔍", "Running PaddleOCR — extracting text from image…", "#0891b2")
        elif isinstance(data, dict) and "status" in data:
            s = _status_html("🤖", "Local 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"Local AI parsing… ({n} chars generated)", "#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)  

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

    with gr.Tabs():

        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",
                    )
                    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()

            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,
            )

        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,
    )