import os
import io
import re
import base64
import hashlib
import streamlit as st
import requests
from PIL import Image, ImageOps, ImageFilter
from pdf2image import convert_from_bytes
from openai import OpenAI
from rapidfuzz.distance import Levenshtein

st.set_page_config(page_title="Gov Document OCR", page_icon="📄", layout="wide")

# ─────────────────────────── Config ───────────────────────────
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
if not OPENROUTER_API_KEY:
    st.error("❌ OPENROUTER_API_KEY not found. Please add it to your Space secrets.")
    st.stop()

OPENROUTER_BASE = "https://openrouter.ai/api/v1"
EXTRA_HEADERS = {
    "HTTP-Referer": "https://huggingface.co/spaces",  # replace with your Space URL
    "X-Title": "Gov Document OCR",
}
MODELS = {
    "Gemini 2.5 Flash (recommended)": "google/gemini-2.5-flash",
    "GPT-4o mini (cheapest)": "openai/gpt-4o-mini",
    "GPT-4o (highest cost)": "openai/gpt-4o",
}

# RTL + Nastaliq display: each paragraph/cell picks its own direction from its first letter,
# so mixed English/Urdu documents display correctly line by line.
st.markdown("""
<style>
@import url('https://fonts.googleapis.com/css2?family=Noto+Nastaliq+Urdu:wght@400;700&display=swap');
[data-testid="stMarkdownContainer"] p, [data-testid="stMarkdownContainer"] li,
[data-testid="stMarkdownContainer"] td, [data-testid="stMarkdownContainer"] th,
[data-testid="stMarkdownContainer"] h1, [data-testid="stMarkdownContainer"] h2,
[data-testid="stMarkdownContainer"] h3, textarea {
    unicode-bidi: plaintext;
    text-align: start;
    font-family: "Source Sans Pro", "Source Sans 3", "Noto Nastaliq Urdu", sans-serif;
}
textarea { line-height: 1.9 !important; }
</style>
""", unsafe_allow_html=True)

# ─────────────────────────── Sidebar ───────────────────────────
with st.sidebar:
    st.header("⚙️ Settings")
    model_label = st.selectbox("Model", list(MODELS))
    model_name = MODELS[model_label]
    language = st.selectbox(
        "Document language",
        ["Auto-detect", "English", "Urdu", "Bilingual (English + Urdu)"],
    )

    st.subheader("Document features")
    tag_handwriting = st.checkbox("Tag handwritten annotations", True)
    tag_signatures = st.checkbox("Mark signatures, stamps & logos", True)
    tag_ticks = st.checkbox("Mark ticks & strike-throughs", True)
    ignore_watermarks = st.checkbox("Ignore scanner watermarks (CamScanner etc.)", True)

    st.subheader("Image & cost")
    enhance = st.checkbox("Enhance scan (contrast + sharpen)", True)
    dpi = st.select_slider("PDF render DPI", [150, 200, 250, 300], value=200)
    max_tokens = st.slider("Max output tokens per page", 1000, 4000, 2000, 250)
    st.caption("Lower tokens and Gemini Flash keep OpenRouter costs down.")

# ─────────────────────────── Prompt ───────────────────────────
LANGUAGE_RULES = {
    "Auto-detect": "The document may be in English, Urdu, Pashto or a mix. Detect the language of each block and transcribe it in its original script.",
    "English": "The document is in English. Transcribe it exactly.",
    "Urdu": "The document is in Urdu, usually in Nastaliq script. Transcribe it in Unicode Urdu script in natural right-to-left reading order.",
    "Bilingual (English + Urdu)": "The document mixes English and Urdu. Keep every segment in its original language and script.",
}

def build_prompt():
    rules = [
        "You are an OCR engine for official government documents (notifications, office orders, letters).",
        LANGUAGE_RULES[language],
        "Never translate, transliterate, summarise, or correct spelling. Transcribe exactly what is printed.",
        "Transcribe ALL printed text, including reference numbers, dates, names, designations and pay scales (e.g. BS-04).",
        "Reproduce tables as Markdown tables with every row and column. Merge text that wraps inside a cell into one cell. For right-to-left tables, list columns in reading order.",
        "Use Markdown headings only for letterhead/title lines. Keep numbered lists as numbered lists.",
        "Where handwriting, signatures or stamps overlap printed text, still transcribe the printed text underneath in full.",
    ]
    if tag_handwriting:
        rules.append("Transcribe handwritten additions (dates, numbers, notes) separately as [HANDWRITTEN: text] where they appear, or [HANDWRITTEN: illegible]. Never merge handwriting into printed text.")
    if tag_signatures:
        rules.append("Do not try to read signatures: write [SIGNATURE] instead. Write stamps/seals as [STAMP: any readable text] and emblems/logos as [LOGO].")
    if tag_ticks:
        rules.append("If an item has a tick/checkmark next to it, start that item with [✓]. If it is struck through, start it with [STRUCK].")
    if ignore_watermarks:
        rules.append("Ignore scanner-app watermarks and branding such as 'CamScanner' or 'Scanned with ...', and ignore page-edge noise.")
    rules.append("Return only the transcription in Markdown. No commentary and no code fences.")
    return "\n".join(f"- {r}" for r in rules)

# ─────────────────────────── Image handling ───────────────────────────
def load_pages(file_bytes, filename):
    if filename.lower().endswith(".pdf"):
        return convert_from_bytes(file_bytes, dpi=dpi)
    return [ImageOps.exif_transpose(Image.open(io.BytesIO(file_bytes)))]

def prepare_image(img):
    img = img.convert("RGB")  # keep colour: blue/red ink helps the model separate handwriting from print
    if enhance:
        img = ImageOps.autocontrast(img, cutoff=1)
        img = img.filter(ImageFilter.UnsharpMask(radius=2, percent=120, threshold=3))
    img.thumbnail((2000, 2000))  # caps upload size and token cost, keeps aspect ratio
    return img

def to_jpeg_bytes(img):
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=90)
    return buf.getvalue()

# ─────────────────────────── OCR calls ───────────────────────────
def build_messages(prompt, jpeg_bytes):
    b64 = base64.b64encode(jpeg_bytes).decode("utf-8")
    return [{
        "role": "user",
        "content": [
            {"type": "text", "text": prompt},
            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}},
        ],
    }]

def call_sdk(client, messages):
    r = client.chat.completions.create(
        model=model_name, extra_headers=EXTRA_HEADERS, messages=messages,
        max_tokens=max_tokens, temperature=0,
    )
    return r.choices[0].message.content

def call_http(messages):
    r = requests.post(
        f"{OPENROUTER_BASE}/chat/completions",
        headers={"Authorization": f"Bearer {OPENROUTER_API_KEY}",
                 "Content-Type": "application/json", **EXTRA_HEADERS},
        json={"model": model_name, "messages": messages,
              "max_tokens": max_tokens, "temperature": 0},
        timeout=120,
    )
    if r.status_code != 200:
        raise RuntimeError(f"{r.status_code}: {r.text[:300]}")
    return r.json()["choices"][0]["message"]["content"]

def clean_output(text):
    text = (text or "").strip()
    return re.sub(r"^```(?:markdown|md)?\s*|\s*```$", "", text).strip()

def ocr_page(client, prompt, jpeg_bytes):
    """Try the SDK first; if it fails, fall back to raw HTTP for this page."""
    messages = build_messages(prompt, jpeg_bytes)
    try:
        return clean_output(call_sdk(client, messages)), "SDK"
    except Exception as e1:
        try:
            return clean_output(call_http(messages)), "HTTP fallback"
        except Exception as e2:
            raise RuntimeError(f"SDK error: {e1}\nHTTP error: {e2}")

# ─────────────────────────── Evaluation (CER / WER) ───────────────────────────
TAG_RE = re.compile(r"\[(?:HANDWRITTEN|SIGNATURE|STAMP|LOGO|STRUCK|✓)[^\]]*\]")
ARABIC_TO_URDU = str.maketrans({"ي": "ی", "ى": "ی", "ك": "ک", "ه": "ہ"})
DIACRITICS_RE = re.compile(r"[\u064B-\u0652\u0670]")

def markdown_to_plain(md, strip_tags):
    text = re.sub(r"^--- Page \d+ ---$", " ", md, flags=re.M)
    if strip_tags:
        text = TAG_RE.sub(" ", text)
    text = re.sub(r"^\s*\|?\s*:?-{3,}.*$", " ", text, flags=re.M)  # table separator rows
    text = text.replace("|", " ")
    text = re.sub(r"^#{1,6}\s*", "", text, flags=re.M)
    text = text.replace("**", "").replace("__", "")
    return text

def normalize(text, ignore_case, urdu_norm):
    if urdu_norm:
        text = DIACRITICS_RE.sub("", text.translate(ARABIC_TO_URDU))
    if ignore_case:
        text = text.lower()
    return " ".join(text.split())

def cer(ref, hyp):
    return Levenshtein.distance(ref, hyp) / max(len(ref), 1)

def wer(ref, hyp):
    r, h = ref.split(), hyp.split()
    return Levenshtein.distance(r, h) / max(len(r), 1)

# ─────────────────────────── App ───────────────────────────
st.title("📄 Government Document OCR")
st.markdown("Multilingual OCR (English / Urdu) for notifications and office orders. "
            "Handles tables, handwriting, signatures, ticks and scanner watermarks, with built-in CER/WER evaluation.")

uploaded = st.file_uploader("📎 Upload a PDF or image", type=["pdf", "png", "jpg", "jpeg"])

if uploaded:
    file_bytes = uploaded.getvalue()
    prompt = build_prompt()
    run_key = hashlib.md5(
        file_bytes + f"{model_name}|{prompt}|{enhance}|{dpi}|{max_tokens}".encode()
    ).hexdigest()

    # OCR runs only on click and is cached, so using the evaluation tab doesn't re-bill you
    if st.button("🚀 Run OCR", type="primary"):
        try:
            raw_pages = load_pages(file_bytes, uploaded.name)
        except Exception as e:
            st.error(f"❌ Could not read file: {e}")
            st.stop()

        client = OpenAI(base_url=OPENROUTER_BASE, api_key=OPENROUTER_API_KEY)
        progress = st.progress(0.0)
        results = []
        for i, page in enumerate(raw_pages, 1):
            progress.progress((i - 1) / len(raw_pages), text=f"🔍 OCR page {i}/{len(raw_pages)}...")
            jpeg = to_jpeg_bytes(prepare_image(page))
            try:
                text, method = ocr_page(client, prompt, jpeg)
                results.append({"page": i, "image": jpeg, "text": text, "method": method, "error": None})
            except Exception as e:
                results.append({"page": i, "image": jpeg, "text": "[ERROR]", "method": None, "error": str(e)})
        progress.progress(1.0, text="✅ Done")
        st.session_state.ocr = {"key": run_key, "name": uploaded.name, "pages": results}

    result = st.session_state.get("ocr")
    if result and result["name"] == uploaded.name:
        if result["key"] != run_key:
            st.info("ℹ️ Settings changed since the last run. Click **Run OCR** again to apply them.")

        pages = result["pages"]
        failed = [p for p in pages if p["error"]]
        if len(failed) == len(pages):
            st.error("❌ OCR failed on all pages.")
        elif failed:
            st.warning(f"⚠️ {len(failed)} of {len(pages)} page(s) failed.")
        else:
            st.success(f"✅ OCR completed on {len(pages)} page(s).")

        full_md = "\n\n".join(f"--- Page {p['page']} ---\n{p['text']}" for p in pages)
        base = os.path.splitext(uploaded.name)[0]

        tab_view, tab_raw, tab_eval = st.tabs(["📖 Side-by-side", "📝 Raw Markdown", "📏 Evaluate (CER/WER)"])

        with tab_view:
            for p in pages:
                st.markdown(f"#### Page {p['page']}" + (f"  ·  _{p['method']}_" if p["method"] else ""))
                col_img, col_txt = st.columns(2)
                col_img.image(p["image"], use_container_width=True)
                with col_txt:
                    if p["error"]:
                        st.error(p["error"])
                    else:
                        st.markdown(p["text"])
                st.divider()

        with tab_raw:
            st.text_area("Extracted Markdown", full_md, height=400)
            c1, c2 = st.columns(2)
            c1.download_button("💾 Download .md", full_md, file_name=f"{base}.md", mime="text/markdown")
            c2.download_button("💾 Download .txt", markdown_to_plain(full_md, strip_tags=False),
                               file_name=f"{base}.txt", mime="text/plain")

        with tab_eval:
            st.markdown("Paste or upload the **correct text** (ground truth) to measure accuracy.")
            gt_file = st.file_uploader("Ground truth .txt", type=["txt"], key="gt")
            gt_text = gt_file.getvalue().decode("utf-8") if gt_file else st.text_area("…or paste it here", height=200)

            e1, e2, e3 = st.columns(3)
            strip_tags = e1.checkbox("Ignore annotation tags", True,
                                     help="Removes [HANDWRITTEN], [SIGNATURE], [✓] etc. before scoring")
            ignore_case = e2.checkbox("Ignore case", False)
            urdu_norm = e3.checkbox("Normalize Urdu characters", True,
                                    help="Unifies Arabic/Urdu letter variants and removes diacritics")

            if st.button("📏 Compute metrics") and gt_text.strip():
                ref = normalize(markdown_to_plain(gt_text, strip_tags), ignore_case, urdu_norm)
                hyp = normalize(markdown_to_plain(full_md, strip_tags), ignore_case, urdu_norm)
                c, w = cer(ref, hyp), wer(ref, hyp)
                m1, m2, m3 = st.columns(3)
                m1.metric("CER", f"{c:.2%}")
                m2.metric("WER", f"{w:.2%}")
                m3.metric("Character accuracy", f"{max(0.0, 1 - c):.2%}")
                st.caption(f"Reference: {len(ref)} chars / {len(ref.split())} words · "
                           f"OCR output: {len(hyp)} chars / {len(hyp.split())} words")