"""
ocr_engine.py - the Colab pipeline (CELLS 2b, 2c, 2d, 3, 4b) as an importable module.

    eng = OCREngine.load()                    # once per process: loads OCR + Real-ESRGAN on the GPU
    result = eng.process(Path("doc.pdf"), lang="kn-IN", work_dir=Path("/tmp/job1"))

Set OCR_MOCK=1 to run without a GPU/model (returns fake blocks) - used to develop the UI.
Everything else is configured with env vars (see CONFIG below); defaults match the notebook.
"""
from __future__ import annotations

import base64, io, json, os, re, sys, threading, time
from collections import Counter
from pathlib import Path

from PIL import Image, ImageDraw, ImageFont, ImageOps

# ---------------------------------------------------------------- config (notebook defaults)
CONFIG = dict(
    REPO=os.getenv("OCR_REPO"),                         # local snapshot of bodhan-ai/indic-ocr (None = download)
    HF_REPO_ID=os.getenv("OCR_HF_REPO_ID", "bodhan-ai/indic-ocr"),
    ESRGAN_PATH=os.getenv("OCR_ESRGAN_PATH", "/models/RealESRGAN_x4plus.pth"),
    UPSCALER=os.getenv("OCR_UPSCALER", "realesrgan"),   # "realesrgan" | "lanczos"
    DEBLUR=os.getenv("OCR_DEBLUR", "1") == "1",
    PDF_DPI=int(os.getenv("OCR_PDF_DPI", "200")),
    MIN_LONG_SIDE=int(os.getenv("OCR_MIN_LONG_SIDE", "1600")),
    MAX_NEW_TOKENS=int(os.getenv("OCR_MAX_NEW_TOKENS", "1024")),
    MAX_PAGES=int(os.getenv("OCR_MAX_PAGES", "10")),     # guard for the public demo
    MASK_AADHAAR=os.getenv("OCR_MASK_AADHAAR", "1") == "1",
    MOCK=os.getenv("OCR_MOCK", "0") == "1",
)

IMAGE_EXT = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp", ".webp"}
LANG_NAMES = {"hi-IN": "Hindi", "kn-IN": "Kannada", "ml-IN": "Malayalam", "ta-IN": "Tamil", "te-IN": "Telugu"}

TYPE_COLORS = {
    "Text": (33, 150, 243), "Title": (142, 36, 170), "SectionHeader": (142, 36, 170),
    "Table": (0, 137, 123), "Picture": (245, 124, 0), "Figure": (245, 124, 0),
    "Equation": (109, 76, 65), "PageHeader": (96, 125, 139), "PageFooter": (96, 125, 139),
    "PageNumber": (96, 125, 139),
}
DEFAULT_COLOR, REVIEW_COLOR = (84, 110, 122), (229, 57, 53)


# ---------------------------------------------------------------- text clean-up (CELL 3, unchanged)
def fix_spaces(text):
    if not text:
        return text
    text = re.sub(r"[ ▁]+", " ", text.replace("▁", " ")).strip()
    return re.sub(r"(.)\1{19,}", lambda m: m.group(1) * 10, text)


def cut_loops(text, window=40, max_unique=4, min_len=800):
    if not text or len(text) < min_len:
        return text, False
    words = text.split(" ")
    norm = [re.sub(r"\d+", "#", w) for w in words]
    for i in range(len(words) - window + 1):
        win = norm[i:i + window]
        if len(set(win)) <= max_unique and any(any(ch.isalpha() for ch in w) for w in win):
            return " ".join(words[:i + 6]) + " [REPETITION_TRUNCATED]", True
    return text, False


def mask_ids(text):
    if not CONFIG["MASK_AADHAAR"] or not text:
        return text
    return re.sub(r"\b(\d{4})[\s-]?(\d{4})[\s-]?(\d{4})\b", lambda m: "XXXX XXXX " + m.group(3), text)


def clean(text):
    text, looped = cut_loops(fix_spaces(text))
    return mask_ids(text), looped


def deblur(im):
    import cv2, numpy as np
    g = cv2.cvtColor(np.asarray(im.convert("RGB")), cv2.COLOR_RGB2GRAY)
    g = cv2.fastNlMeansDenoising(g, h=4, templateWindowSize=7, searchWindowSize=21)
    g = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)).apply(g)
    for sigma, amount in ((2.0, 1.0), (0.7, 0.7)):
        g = cv2.addWeighted(g, 1 + amount, cv2.GaussianBlur(g, (0, 0), sigma), -amount, 0)
    return Image.fromarray(cv2.cvtColor(g, cv2.COLOR_GRAY2RGB))


# ---------------------------------------------------------------- line boxes inside a layout block
# The layout model returns one box per paragraph / list, which on handwritten pages covers 10-25
# lines plus margins. These helpers find the actual written lines from the ink so the highlight
# sits on the text. Tables, pictures and equations keep their block box.
NO_LINE_TYPES = {"Table", "Picture", "Figure", "Equation"}


def _ink_mask(gray):
    import cv2, numpy as np
    h, w = gray.shape
    k = max(15, (min(h, w) // 20) | 1)
    bg = cv2.morphologyEx(gray, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (k, k)))
    norm = cv2.divide(gray, bg, scale=255)                 # flatten shadows / paper tint
    _, m = cv2.threshold(norm, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    for ks in [(max(40, w // 6), 1), (1, max(40, h // 6))]:  # ruled notebook lines, margin line
        m = cv2.subtract(m, cv2.morphologyEx(m, cv2.MORPH_OPEN, cv2.getStructuringElement(cv2.MORPH_RECT, ks)))
    n, lab, st, _ = cv2.connectedComponentsWithStats(m, 8)
    keep = np.zeros(n, bool)
    keep[1:] = st[1:, cv2.CC_STAT_AREA] >= 12               # specks
    return keep[lab]


def _line_pitch(prof, lo, hi):
    # distance between text lines = first autocorrelation peak of the row-ink profile
    import numpy as np
    x = prof - prof.mean()
    ac = np.correlate(x, x, "full")[len(x) - 1:]
    if ac[0] <= 0 or hi <= lo:
        return None
    ac = ac / ac[0]
    lag = lo + int(np.argmax(ac[lo:hi]))
    return lag if ac[lag] > 0.1 else None


def find_lines(gray, box):
    """Tight boxes around each written line inside `box` (pixel coords of `gray`)."""
    import numpy as np
    x0, y0, x1, y1 = [int(round(v)) for v in box]
    x0, y0 = max(x0, 0), max(y0, 0)
    crop = gray[y0:y1, x0:x1]
    if crop.size == 0:
        return []
    m = _ink_mask(crop)
    raw = m.sum(1).astype(float)
    if raw.max() == 0:
        return []
    h = len(raw)
    pitch = (_line_pitch(raw, max(8, h // 40), h // 2 + 1) if h > 40 else None) or h
    k = max(3, pitch // 4)
    prof = np.convolve(raw, np.ones(k) / k, "same")
    cuts = [0]                      # cut at the emptiest row between consecutive lines
    while cuts[-1] < h:
        lo, hi = cuts[-1] + int(0.6 * pitch), min(h, cuts[-1] + int(1.4 * pitch))
        if lo >= h or hi <= lo:
            break
        c = lo + int(np.argmin(prof[lo:hi]))
        if prof[c] > 0.5 * prof[cuts[-1]:hi].max():
            c = hi                  # no real gap in this window
        cuts.append(c)
    cuts = sorted(set(cuts + [h]))
    floor = 0.08 * np.percentile(raw[raw > 0], 90)
    out = []
    for a, b in zip(cuts, cuts[1:]):
        rows = np.where(raw[a:b] > floor)[0]
        if len(rows) == 0:
            continue
        a, b = a + rows[0], a + rows[-1] + 1
        if b - a < 0.3 * pitch:
            continue                # ruled-line residue, stray dots
        cols = np.where(m[a:b].any(0))[0]
        out.append([x0 + int(cols[0]), y0 + int(a), x0 + int(cols[-1]) + 1, y0 + int(b)])
    return out


def add_lines(page, im):
    """Adds b["lines"] = [{"bbox_xyxy", "text"}] (page coords) to text blocks; `im` is the page image.
    Line text is filled in only when the OCR text has as many lines as were found on the page."""
    import numpy as np
    sx = im.width / (page.get("width") or im.width)
    sy = im.height / (page.get("height") or im.height)
    gray = np.asarray(im.convert("L"))
    pad = max(2, round(max(im.size) / 700))
    for b in page.get("blocks", []):
        text = b.get("text") or ""
        if not text.strip() or b.get("type") in NO_LINE_TYPES or "XXXX XXXX" in text or not b.get("bbox_xyxy"):
            continue
        x0, y0, x1, y1 = b["bbox_xyxy"]
        found = find_lines(gray, [x0 * sx + pad, y0 * sy + pad, x1 * sx - pad, y1 * sy - pad])
        texts = [t for t in text.split("\n") if t.strip()]
        b["lines"] = [{"bbox_xyxy": [round(l[0] / sx, 1), round(l[1] / sy, 1), round(l[2] / sx, 1), round(l[3] / sy, 1)],
                       "text": texts[i] if len(texts) == len(found) else None} for i, l in enumerate(found)]
    return page


def highlight_page(page, im, fill_alpha=50, redact=True):
    """CELL 4b drawing, shared by the notebook and the app: one filled box per written line
    (or per block for tables / pictures), a "3 Text" tag on each block, masked IDs blacked out.
    Blocks without text (empty page headers / footers) are skipped."""
    im = im.convert("RGBA")
    sx = im.width / (page.get("width") or im.width)
    sy = im.height / (page.get("height") or im.height)
    overlay = Image.new("RGBA", im.size, (0, 0, 0, 0))
    d = ImageDraw.Draw(overlay)
    lw = max(2, round(max(im.size) / 700))
    font = get_font(max(14, round(max(im.size) / 80)))
    for b in page.get("blocks", []):
        text = b.get("text") or ""
        if not text.strip() or not b.get("bbox_xyxy"):
            continue
        scale = lambda bb: [bb[0] * sx, bb[1] * sy, bb[2] * sx, bb[3] * sy]
        if redact and "XXXX XXXX" in text:
            d.rectangle(scale(b["bbox_xyxy"]), fill=(0, 0, 0, 255))
            continue
        review = b.get("needs_review")
        c = REVIEW_COLOR if review else TYPE_COLORS.get(b.get("type"), DEFAULT_COLOR)
        boxes = [scale(l["bbox_xyxy"]) for l in b.get("lines") or []] or [scale(b["bbox_xyxy"])]
        for box in boxes:
            d.rectangle(box, fill=c + (fill_alpha,), outline=c + (255,), width=lw)
        label = f"{b.get('order')} {b.get('type', '')}" + (" REVIEW" if review else "")
        l, t, r, btm = d.textbbox((0, 0), label, font=font)
        w, h = r - l, btm - t
        x, y = boxes[0][0], boxes[0][1]
        ly = y - h - 6 if y - h - 6 >= 0 else y
        d.rectangle([x, ly, x + w + 8, ly + h + 6], fill=c + (235,))
        d.text((x + 4, ly + 3 - t), label, fill=(255, 255, 255, 255), font=font)
    return Image.alpha_composite(im, overlay).convert("RGB")


def get_font(size):
    for p in ["/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
              "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf"]:
        if Path(p).exists():
            return ImageFont.truetype(p, size)
    return ImageFont.load_default()


def to_data_url(im, max_side=1800, quality=85):
    im = im.copy()
    im.thumbnail((max_side, max_side))
    buf = io.BytesIO()
    im.convert("RGB").save(buf, "JPEG", quality=quality)
    return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()


# ---------------------------------------------------------------- engine
class OCREngine:
    """One instance per process. process() is serialised with a lock: one GPU, one job at a time."""

    def __init__(self):
        self.lock = threading.Lock()
        self.doc_lang = None          # read by the recognizer patches (was DOC_LANG in the notebook)
        self.parser = None
        self.ready = False
        self.nb = None                # notebook globals when attached with from_notebook()

    # -------------------------------------------------------- loading
    @classmethod
    def load(cls) -> "OCREngine":
        eng = cls()
        if CONFIG["MOCK"]:
            eng.ready = True
            return eng
        eng._load_models()
        eng.ready = True
        return eng

    @classmethod
    def from_notebook(cls, ns: dict) -> "OCREngine":
        """Reuse the Colab kernel's already-loaded pipeline (CELLS 2b-3): no second model copy,
        no session restart, and results identical to the notebook. Call with globals()."""
        missing = [k for k in ("parser", "ocr_image", "DOC_LANG") if k not in ns]
        if missing:
            raise RuntimeError(f"run CELLS 2b, 2c, 2d and 3 first (missing: {', '.join(missing)})")
        eng = cls()
        eng.nb = ns
        eng.ready = True
        return eng

    def _load_models(self):
        import numpy as np, torch
        repo = CONFIG["REPO"]
        if not repo:
            from huggingface_hub import snapshot_download
            repo = snapshot_download(CONFIG["HF_REPO_ID"])   # uses HF_TOKEN from the env
        self.repo = repo
        sys.path.insert(0, repo)
        sys.path.insert(0, str(Path(__file__).parent))       # validate.py next to this file

        # --- CELL 2b: Real-ESRGAN
        self.sr_model = None
        if CONFIG["UPSCALER"] == "realesrgan":
            from spandrel import ModelLoader
            self.sr_model = ModelLoader().load_from_file(CONFIG["ESRGAN_PATH"]).cuda().eval()
            self.sr_half = self.sr_model.supports_half
            if self.sr_half:
                self.sr_model.half()

        # --- CELL 2c: shipped tokenizer   + CELL 2d: language hint & script remap
        import idp_recognizer
        from transformers import PreTrainedTokenizerFast
        try:
            from validate import to_script
        except ImportError:
            print("[warn] validate.py not found next to ocr_engine.py - script remap disabled")
            to_script = lambda text, lang: text
        eng = self
        Rec = idp_recognizer.HfRecognizer
        orig_init, orig_prompt, orig_transcribe = Rec.__init__, Rec._prompt, Rec.transcribe

        def init(self_, *a, **k):
            orig_init(self_, *a, **k)
            tok = PreTrainedTokenizerFast.from_pretrained(self_.ckpt)
            tok.padding_side = "left"
            assert tok.eos_token_id == self_.eos_token_id
            self_.processor.tokenizer = tok

        def prompt(self_, text):
            if eng.doc_lang and text.startswith("Transcribe the text in this image"):
                name = LANG_NAMES[eng.doc_lang]
                text = f"Transcribe the {name} text in this image. Write it in {name} script only."
            return orig_prompt(self_, text)

        def transcribe(self_, requests):
            out = orig_transcribe(self_, requests)
            return [to_script(t, eng.doc_lang) for t in out] if eng.doc_lang else out

        Rec.__init__, Rec._prompt, Rec.transcribe = init, prompt, transcribe

        # --- CELL 3: parser
        from indic_ocr import IndicOCR
        self.parser = IndicOCR.from_pretrained(repo, table_format="markdown",
                                               max_new_tokens=CONFIG["MAX_NEW_TOKENS"])
        self._torch, self._np = torch, np

    # -------------------------------------------------------- pre-processing (CELL 3)
    def _sr_upscale(self, im):
        torch, np = self._torch, self._np
        with torch.inference_mode():
            x = torch.from_numpy(np.array(im.convert("RGB"))).permute(2, 0, 1)[None].cuda()
            x = (x.half() if self.sr_half else x.float()) / 255
            y = self.sr_model(x).clamp(0, 1)[0].permute(1, 2, 0).float().cpu().numpy()
        return Image.fromarray((y * 255).round().astype(np.uint8))

    def _prepare_image(self, img_path, work_dir):
        im = ImageOps.exif_transpose(Image.open(img_path)).convert("RGB")
        if max(im.size) >= CONFIG["MIN_LONG_SIDE"] and not CONFIG["DEBLUR"]:
            return img_path
        if CONFIG["DEBLUR"]:
            im = deblur(im)
        s = max(1.0, CONFIG["MIN_LONG_SIDE"] / max(im.size))
        size = (round(im.width * s), round(im.height * s))
        if self.sr_model is not None and s > 1:
            try:
                im = self._sr_upscale(im)
            except Exception as e:
                print(f"SR failed, using LANCZOS: {e}")
                self._torch.cuda.empty_cache()
        im = im.resize(size, Image.LANCZOS)
        out = work_dir / "_pages" / f"{Path(img_path).stem}_prep_x{s:.1f}.png"
        out.parent.mkdir(parents=True, exist_ok=True)
        im.save(out)
        return out

    def _pdf_to_images(self, pdf_path, work_dir):
        import pypdfium2 as pdfium
        pdf = pdfium.PdfDocument(str(pdf_path))
        if len(pdf) > CONFIG["MAX_PAGES"]:
            raise ValueError(f"PDF has {len(pdf)} pages; this demo allows {CONFIG['MAX_PAGES']}")
        out_dir = work_dir / "_pages"
        out_dir.mkdir(parents=True, exist_ok=True)
        paths = []
        for i in range(len(pdf)):
            p = out_dir / f"page_{i + 1:03d}.png"
            pdf[i].render(scale=CONFIG["PDF_DPI"] / 72).to_pil().save(p)
            paths.append(p)
        return paths

    def _ocr_image(self, img_path, work_dir):
        if CONFIG["MOCK"]:
            return _mock_page(img_path)
        if self.nb is not None:
            self.nb["DOC_LANG"] = self.doc_lang        # what the notebook's recognizer patches read
            try:
                return self.nb["ocr_image"](Path(img_path))
            finally:
                self.nb["DOC_LANG"] = None
        page = self.parser.parse(str(self._prepare_image(img_path, work_dir)))
        n_loops = 0
        for b in page.get("blocks", []):
            b["text"], looped = clean(b.get("text", ""))
            b["needs_review"] = looped
            n_loops += looped
        page["markdown"] = "\n".join(clean(l)[0] for l in page.get("markdown", "").splitlines())
        page["loops_truncated"] = n_loops
        return page

    # -------------------------------------------------------- rendering (CELL 4b)
    def _parsed_image(self, page, src, work_dir):
        """The image the model parsed (deblurred / upscaled); page["image"] is only its file name."""
        name = Path(page.get("image") or "").name
        dirs = [work_dir / "_pages"] + ([Path(self.nb["OUTPUT_DIR"]) / "_pages"] if self.nb and "OUTPUT_DIR" in self.nb else [])
        return next((d / name for d in dirs if name and (d / name).is_file()), src)

    @staticmethod
    def _redacted_page(page, src):
        """Page with masked-ID blocks blacked out (so the UI never shows the raw number)."""
        im = ImageOps.exif_transpose(Image.open(src)).convert("RGB")
        sx = im.width / (page.get("width") or im.width)
        sy = im.height / (page.get("height") or im.height)
        d = ImageDraw.Draw(im)
        for b in page.get("blocks", []):
            if "XXXX XXXX" in (b.get("text") or ""):
                x0, y0, x1, y1 = b["bbox_xyxy"]
                d.rectangle([x0 * sx, y0 * sy, x1 * sx, y1 * sy], fill=(0, 0, 0))
        return im, sx, sy

    # -------------------------------------------------------- public entry point
    def process(self, file_path: Path, lang: str | None, work_dir: Path) -> dict:
        if lang and lang not in LANG_NAMES:
            raise ValueError(f"unknown language {lang!r}")
        work_dir.mkdir(parents=True, exist_ok=True)
        suffix = file_path.suffix.lower()
        if suffix not in IMAGE_EXT | {".pdf"}:
            raise ValueError(f"unsupported file type {suffix}")

        t0 = time.time()
        with self.lock:
            self.doc_lang = lang
            try:
                sources = self._pdf_to_images(file_path, work_dir) if suffix == ".pdf" else [file_path]
                raw_pages = [self._ocr_image(p, work_dir) for p in sources]
            finally:
                self.doc_lang = None
        elapsed = round(time.time() - t0, 1)

        pages, md_parts, highlighted = [], [], []
        for i, (page, src) in enumerate(zip(raw_pages, sources), 1):
            base, sx, sy = self._redacted_page(page, self._parsed_image(page, src, work_dir))
            add_lines(page, base)
            highlighted.append(highlight_page(page, base))
            md_parts.append(f"<!-- page {i} -->\n{page.get('markdown', '')}")
            pages.append({
                "index": i,
                "width": page.get("width") or base.width,     # bbox coordinate space
                "height": page.get("height") or base.height,
                "image": to_data_url(base),
                "markdown": page.get("markdown", ""),
                "loops_truncated": page.get("loops_truncated", 0),
                "blocks": [{
                    "order": b.get("order", k),
                    "type": b.get("type", "Text"),
                    "text": b.get("text", ""),
                    "bbox_xyxy": b.get("bbox_xyxy"),
                    "lines": b.get("lines") or [],
                    "needs_review": bool(b.get("needs_review")),
                    "redacted": "XXXX XXXX" in (b.get("text") or ""),
                } for k, b in enumerate(page.get("blocks", []), 1)
                    if b.get("bbox_xyxy") and (b.get("text") or "").strip()],   # empty header/footer boxes
            })

        # same three artefacts the notebook writes
        stem = file_path.stem
        (work_dir / f"{stem}.json").write_text(json.dumps(raw_pages, ensure_ascii=False, indent=2,
                                                          default=str), encoding="utf-8")
        (work_dir / f"{stem}.md").write_text("\n\n".join(md_parts), encoding="utf-8")
        highlighted[0].save(work_dir / f"{stem}_highlighted.pdf", "PDF", resolution=150,
                            save_all=True, append_images=highlighted[1:])

        blocks = [b for p in pages for b in p["blocks"]]
        return {
            "file": file_path.name,
            "lang": lang,
            "lang_name": LANG_NAMES.get(lang, "Auto"),
            "stats": {
                "pages": len(pages),
                "blocks": len(blocks),
                "chars": sum(len(b["text"]) for b in blocks),
                "needs_review": sum(b["needs_review"] for b in blocks),
                "redacted": sum(b["redacted"] for b in blocks),
                "seconds": elapsed,
            },
            "pages": pages,
            "markdown": "\n\n".join(md_parts),
            "artifacts": {"json": f"{stem}.json", "md": f"{stem}.md", "pdf": f"{stem}_highlighted.pdf"},
        }


# ---------------------------------------------------------------- mock output for UI work
def _mock_page(img_path):
    time.sleep(1.2)
    im = ImageOps.exif_transpose(Image.open(img_path))
    W, H = im.size
    samples = [("Title", "ಕರ್ನಾಟಕ ಸರ್ಕಾರ — ಅರ್ಜಿ ನಮೂನೆ"),
               ("Text", "ಹೆಸರು: ರಮೇಶ್ ಕುಮಾರ್   ತಂದೆಯ ಹೆಸರು: ಶಿವಪ್ಪ"),
               ("Text", "ವಿಳಾಸ: ನಂ. 12, 3ನೇ ಅಡ್ಡರಸ್ತೆ, ಜಯನಗರ, ಬೆಂಗಳೂರು - 560041"),
               ("Text", "Aadhaar: " + mask_ids("1234 5678 9012")),
               ("Table", "| ಕ್ರಮ | ವಿವರ | ಮೊತ್ತ |\n|---|---|---|\n| 1 | ಶುಲ್ಕ | 500 |\n| 2 | ಠೇವಣಿ | 1000 |"),
               ("Text", "ದಿನಾಂಕ ದಿನಾಂಕ ದಿನಾಂಕ ... [REPETITION_TRUNCATED]"),
               ("PageFooter", "ಸಹಿ / Signature")]
    blocks, y = [], 0.06
    heights = [0.07, 0.06, 0.08, 0.05, 0.22, 0.06, 0.05]
    for k, ((t, txt), h) in enumerate(zip(samples, heights), 1):
        blocks.append({"order": k, "type": t, "text": txt, "needs_review": "TRUNCATED" in txt,
                       "bbox_xyxy": [round(W * 0.08), round(H * y), round(W * (0.92 if t != "PageFooter" else 0.4)),
                                     round(H * (y + h))]})
        y += h + 0.04
    md = "\n\n".join(("# " + b["text"]) if b["type"] == "Title" else b["text"] for b in blocks)
    return {"width": W, "height": H, "blocks": blocks, "markdown": md, "loops_truncated": 1}