#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
📝 PDF EDITER — Advanced PDF Editor with OCR
=============================================
Upload PDFs → OCR scanned pages → Edit text → Manage pages →
Merge / split → Watermark → Compress → Protect → Download

Built with Gradio + PyMuPDF + Tesseract OCR.
"""

import io
import base64
import math
import os
import tempfile

from PIL import Image, ImageDraw, ImageFont

import gradio as gr
import pymupdf as fitz

# --------------------------------------------------------------------------
#  Constants
# --------------------------------------------------------------------------
RENDER_ZOOM = 1.6          # zoom used for the interactive page preview
THUMB_SCALE = 0.35         # zoom used for gallery thumbnails
MAX_UNDO = 25

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# Noto Sans script fonts, embedded as base64 (keeps the HF repo text-only).
# Latin/Cyrillic/Greek is served by the pip package `pymupdf-fonts` (fontname "notos").
try:
    from fonts_data import FONTS_B64
except Exception:
    FONTS_B64 = {}

_font_cache = {}


def get_font_file(key):
    """Decode an embedded font once, cache its temp-file path."""
    if key in _font_cache:
        return _font_cache[key]
    data = base64.b64decode(FONTS_B64[key])
    fd, path = tempfile.mkstemp(suffix=".ttf", prefix=f"pdfediter_{key}_")
    with os.fdopen(fd, "wb") as f:
        f.write(data)
    _font_cache[key] = path
    return path


def resolve_font(key):
    """(fontname, fontfile) ready for insert_text; fontfile None = built-in/pip font."""
    name, ref = FONTS[key]
    return name, (get_font_file(ref) if ref else None)


# Font choices: display name -> (pymupdf name, embedded-font key or None)
FONTS = {
    "Helvetica (default)": ("helv", None),
    "Times Roman": ("tiro", None),
    "Courier": ("cour", None),
    "Noto Sans (Latin/Cyrillic/Greek)": ("notos", None),
    "Noto Sans Devanagari (हिन्दी)": ("notosdev", "devanagari"),
    "Noto Sans Tamil (தமிழ்)": ("notostam", "tamil"),
    "Noto Sans Bengali (বাংলা)": ("notosben", "bengali"),
    "Noto Sans Arabic (العربية)": ("notosara", "arabic"),
    "Noto Sans Gurmukhi (ਪੰਜਾਬੀ)": ("notosgur", "gurmukhi"),
}

# Tesseract languages: display name -> tesseract code
OCR_LANGS = {
    "English": "eng",
    "Hindi (हिन्दी)": "hin",
    "Tamil (தமிழ்)": "tam",
    "Telugu (తెలుగు)": "tel",
    "Kannada (ಕನ್ನಡ)": "kan",
    "Malayalam (മലയാളം)": "mal",
    "Bengali (বাংলা)": "ben",
    "Marathi (मराठी)": "mar",
    "Gujarati (ગુજરાતી)": "guj",
    "Punjabi (ਪੰਜਾਬੀ)": "pan",
    "Spanish": "spa",
    "French": "fra",
    "German": "deu",
    "Italian": "ita",
    "Portuguese": "por",
    "Russian (русский)": "rus",
    "Arabic (العربية)": "ara",
    "Chinese Simplified (简体中文)": "chi-sim",
    "Japanese (日本語)": "jpn",
}

# Unicode ranges used to auto-pick a font for arbitrary text
R_DEV = (0x0900, 0x097F)
R_TAM = (0x0B80, 0x0BFF)
R_BEN = (0x0980, 0x09FF)
R_ARA = (0x0600, 0x06FF)
R_GUR = (0x0A00, 0x0A7F)


# --------------------------------------------------------------------------
#  Small helpers
# --------------------------------------------------------------------------
def new_state():
    return {
        "bytes": None,
        "name": "untitled.pdf",
        "page": 0,
        "undo": [],
        "redo": [],
        "form_fields": [],
        "pw": "",
        "last": "",
    }


def open_doc(S):
    if not S.get("bytes"):
        raise gr.Error("📄 No PDF loaded yet — upload a file or press “Load demo PDF”.")
    doc = fitz.open(stream=S["bytes"], filetype="pdf")
    if doc.needs_pass:
        pw = S.get("pw") or ""
        if not doc.authenticate(pw):
            doc.close()
            raise gr.Error("🔒 This PDF is password-protected — unlock it in the "
                           "Merge & Advanced tab first.")
    return doc


def snapshot(S):
    S["undo"].append(S["bytes"])
    if len(S["undo"]) > MAX_UNDO:
        S["undo"].pop(0)
    S["redo"] = []


def commit_bytes(S, data, label):
    snapshot(S)
    S["bytes"] = data
    S["last"] = label
    if S.get("page") is not None:
        try:
            d = fitz.open(stream=data, filetype="pdf")
            S["page"] = min(S["page"], d.page_count - 1)
            d.close()
        except Exception:
            S["page"] = 0


def commit(S, doc, label):
    commit_bytes(S, doc.tobytes(garbage=4, deflate=True), label)


def hex_to_rgb(h):
    """'#rrggbb' → (r, g, b) floats in 0..1 (PyMuPDF color convention)."""
    h = h.lstrip("#")
    return tuple(int(h[i:i + 2], 16) / 255 for i in (0, 2, 4))


def has_script(text, ranges):
    for ch in text:
        o = ord(ch)
        for a, b in ranges:
            if a <= o <= b:
                return True
    return False


def pick_font(text):
    """Return (fontname, fontfile) that can render the given text."""
    if has_script(text, [R_DEV]):
        return resolve_font("Noto Sans Devanagari (हिन्दी)")
    if has_script(text, [R_TAM]):
        return resolve_font("Noto Sans Tamil (தமிழ்)")
    if has_script(text, [R_BEN]):
        return resolve_font("Noto Sans Bengali (বাংলা)")
    if has_script(text, [R_ARA]):
        return resolve_font("Noto Sans Arabic (العربية)")
    if has_script(text, [R_GUR]):
        return resolve_font("Noto Sans Gurmukhi (ਪੰਜਾਬੀ)")
    return resolve_font("Noto Sans (Latin/Cyrillic/Greek)")


def fmt_size(n):
    if n < 1024:
        return f"{n} B"
    if n < 1024 * 1024:
        return f"{n / 1024:.1f} KB"
    return f"{n / (1024 * 1024):.2f} MB"


def save_temp(data, suffix, display_name):
    fd, path = tempfile.mkstemp(suffix=suffix, prefix="pdfediter_")
    with os.fdopen(fd, "wb") as f:
        f.write(data)
    pretty = os.path.join(os.path.dirname(path), display_name)
    try:
        os.replace(path, pretty)
        return pretty
    except OSError:
        return path


def parse_page_ranges(text, n):
    """Parse '1,3,5-8' (1-based) into a sorted 0-based unique int list."""
    out = set()
    for part in str(text).replace(" ", "").split(","):
        if not part:
            continue
        if "-" in part:
            a, b = part.split("-")
            lo, hi = int(a), int(b)
            if lo > hi:
                lo, hi = hi, lo
            out.update(range(max(1, lo), min(n, hi) + 1))
        else:
            v = int(part)
            if 1 <= v <= n:
                out.add(v)
    return sorted(v - 1 for v in out)


def pg_slider_update(n_pages, current):
    return gr.update(maximum=n_pages, value=current + 1)


def xy_slider_update(w, h):
    return gr.update(maximum=w, value=0.0), gr.update(maximum=h, value=0.0)


# --------------------------------------------------------------------------
#  Rendering helpers
# --------------------------------------------------------------------------
def render_pil(S, pno=None, zoom=RENDER_ZOOM):
    doc = open_doc(S)
    if pno is None:
        pno = min(S["page"], doc.page_count - 1)
    page = doc[pno]
    pix = page.get_pixmap(matrix=fitz.Matrix(zoom, zoom), alpha=False)
    img = Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
    doc.close()
    return img


def gallery_pils(S, max_pages=250):
    doc = open_doc(S)
    imgs = []
    for pno in range(min(doc.page_count, max_pages)):
        page = doc[pno]
        pix = page.get_pixmap(matrix=fitz.Matrix(THUMB_SCALE, THUMB_SCALE), alpha=False)
        imgs.append(Image.frombytes("RGB", (pix.width, pix.height), pix.samples))
    doc.close()
    return imgs


def session_info(S):
    if not S.get("bytes"):
        return "No document loaded"
    doc = open_doc(S)
    n = doc.page_count
    doc.close()
    return (f"**{S['name']}** — {n} page{'s' if n != 1 else ''} · "
            f"{fmt_size(len(S['bytes']))} · view page **{S['page'] + 1}/{n}** · "
            f"last action: _{S['last'] or '—'}_")


def page_dims(S):
    doc = open_doc(S)
    page = doc[min(S["page"], doc.page_count - 1)]
    w, h = page.rect.width, page.rect.height
    doc.close()
    return float(w), float(h)


# --------------------------------------------------------------------------
#  Demo PDF generator
# --------------------------------------------------------------------------
def make_demo_pdf():
    doc = fitz.open()

    # --- page 1: title + text to test find/replace
    p = doc.new_page()
    p.insert_text((72, 110), "PDF EDITER — Demo Document", fontsize=24,
                  fontname="helv", color=hex_to_rgb("#4f46e5"))
    p.insert_text((72, 150), "This demo was generated automatically so you can try every feature.",
                  fontsize=12, fontname="helv")
    p.insert_text((72, 185), "Replace Me: apples bananas cherries durian",
                  fontsize=13, fontname="helv")
    p.insert_text((72, 215), "Try the Edit Text tab: search for “apples” and replace it.",
                  fontsize=11, fontname="helv")

    # --- page 2: 'scanned style' page (text baked into an image) for OCR demo
    p2 = doc.new_page()
    try:
        fnt = ImageFont.truetype(get_font_file("devanagari"), 44)  # font incl. Latin glyphs
    except Exception:
        fnt = ImageFont.load_default()
    img = Image.new("RGB", (1600, 900), "white")
    d = ImageDraw.Draw(img)
    d.text((120, 120), "SCANNED STYLE PAGE - FOR OCR", fill="black", font=fnt)
    d.text((120, 240), "This page is a picture of text, not real text.", fill="black", font=fnt)
    d.text((120, 320), "Use the OCR tab to extract a searchable text layer.", fill="black", font=fnt)
    d.text((120, 400), "Invoice number 2026-0815 amount 4,999.99", fill="black", font=fnt)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    p2.insert_image(p2.rect, stream=buf.getvalue())

    # --- page 3: a simple form to demo the Forms tab
    p3 = doc.new_page()
    p3.insert_text((72, 100), "Sample Form", fontsize=18, fontname="helv",
                   color=hex_to_rgb("#4f46e5"))
    for lab, y in [("Name", 140), ("Email", 190), ("City", 240), ("Amount", 290)]:
        p3.insert_text((72, y), f"{lab}:", fontsize=12, fontname="helv")
        w = fitz.Widget()
        w.field_type = fitz.PDF_WIDGET_TYPE_TEXT
        w.field_name = lab.lower()
        w.rect = fitz.Rect(150, y - 12, 420, y + 8)
        w.field_value = ""
        p3.add_widget(w)

    data = doc.tobytes(garbage=4, deflate=True)
    doc.close()
    return data


# --------------------------------------------------------------------------
#  Top bar actions
# --------------------------------------------------------------------------
def load_demo(S):
    S["bytes"] = make_demo_pdf()
    S["name"] = "demo.pdf"
    S["page"] = 0
    S["undo"], S["redo"] = [], []
    S["last"] = "Demo loaded"
    w, h = page_dims(S)
    xup, yup = xy_slider_update(w, h)
    return (render_pil(S), gallery_pils(S), session_info(S),
            "✅ Demo PDF loaded — try OCR, replace, add text, forms…", xup, yup,
            pg_slider_update(3, 0), S)


def load_files(files, S):
    if not files:
        raise gr.Error("Select at least one PDF file.")
    paths = [f if isinstance(f, str) else f.name for f in files]
    merged = fitz.open()
    for path in paths:
        try:
            src = fitz.open(path)
        except Exception:
            raise gr.Error(f"Could not open {os.path.basename(path)} — is it a valid PDF?")
        merged.insert_pdf(src)
        src.close()
    n = merged.page_count
    if n == 0:
        merged.close()
        raise gr.Error("The file(s) contain no pages.")
    S["bytes"] = merged.tobytes(garbage=4, deflate=True)
    merged.close()
    S["name"] = os.path.basename(paths[0]) if len(paths) == 1 else f"merged_{len(paths)}_files.pdf"
    S["page"] = 0
    S["undo"], S["redo"] = [], []
    S["last"] = "Loaded"
    w, h = page_dims(S)
    xup, yup = xy_slider_update(w, h)
    return (render_pil(S), gallery_pils(S), session_info(S),
            f"✅ Loaded {len(paths)} file(s) — {n} page(s) in the working document.",
            xup, yup, pg_slider_update(n, 0), S)


def nav_go(page_no, S):
    doc = open_doc(S)
    n = doc.page_count
    doc.close()
    page_no = max(0, min(int(page_no) - 1, n - 1))
    S["page"] = page_no
    w, h = page_dims(S)
    xup, yup = xy_slider_update(w, h)
    return render_pil(S), session_info(S), xup, yup, pg_slider_update(n, page_no), S


def nav_delta(delta, S):
    doc = open_doc(S)
    n = doc.page_count
    doc.close()
    S["page"] = max(0, min(S["page"] + delta, n - 1))
    w, h = page_dims(S)
    xup, yup = xy_slider_update(w, h)
    return render_pil(S), session_info(S), xup, yup, pg_slider_update(n, S["page"]), S


def undo_action(S):
    if not S["undo"]:
        return (gr.update(), gr.update(), gr.update(), "Nothing to undo.",
                gr.update(), gr.update(), gr.update(), S)
    S["redo"].append(S["bytes"])
    S["bytes"] = S["undo"].pop()
    S["last"] = "Undo"
    n, w, h = _count_dims(S)
    xup, yup = xy_slider_update(w, h)
    return (render_pil(S), gallery_pils(S), session_info(S), "↩️ Undo applied.",
            xup, yup, pg_slider_update(n, S["page"]), S)


def redo_action(S):
    if not S["redo"]:
        return (gr.update(), gr.update(), gr.update(), "Nothing to redo.",
                gr.update(), gr.update(), gr.update(), S)
    S["undo"].append(S["bytes"])
    S["bytes"] = S["redo"].pop()
    S["last"] = "Redo"
    n, w, h = _count_dims(S)
    xup, yup = xy_slider_update(w, h)
    return (render_pil(S), gallery_pils(S), session_info(S), "↪️ Redo applied.",
            xup, yup, pg_slider_update(n, S["page"]), S)


def reset_all(S):
    S.update(new_state())
    return (None, [], "No document loaded", "🧹 Session reset.",
            gr.update(maximum=595, value=0.0), gr.update(maximum=842, value=0.0),
            gr.update(maximum=1, value=1), S)


def _count_dims(S):
    doc = open_doc(S)
    n = doc.page_count
    page = doc[min(S["page"], n - 1)]
    w, h = page.rect.width, page.rect.height
    doc.close()
    return n, float(w), float(h)


def download_current(S):
    if not S.get("bytes"):
        raise gr.Error("Nothing to download yet.")
    path = save_temp(S["bytes"], ".pdf", S["name"])
    return path, "⬇️ Download ready — click the file above to save it."


# --------------------------------------------------------------------------
#  Tab 2 — OCR
# --------------------------------------------------------------------------
def run_ocr(lang_labels, dpi, page_sel, S, progress=gr.Progress()):
    if not S.get("bytes"):
        raise gr.Error("Load a PDF first.")
    langs = "+".join(OCR_LANGS[l] for l in lang_labels if l in OCR_LANGS) or "eng"

    doc = open_doc(S)
    pages = list(range(doc.page_count)) if page_sel == "All pages" else [S["page"]]
    out_parts = []
    progress(0, desc="Preparing OCR…")
    for i, pno in enumerate(pages):
        progress((i + 1) / len(pages), desc=f"OCR page {pno + 1}/{doc.page_count} [{langs}] @ {dpi}dpi")
        page = doc[pno]
        tp = page.get_textpage_ocr(language=langs, dpi=int(dpi), full=True)
        text = page.get_text("text", textpage=tp)
        out_parts.append(f"\n\n===== PAGE {pno + 1} =====\n{text}")
        # Embed an invisible, searchable & selectable text layer — but only on
        # pages that have NO real text yet (they are the genuinely scanned ones).
        existing = page.get_text("text").strip()
        if text.strip() and not existing:
            words = page.get_text("words", textpage=tp)
            for w in words:
                x0, y0, x1, y1, word = w[0], w[1], w[2], w[3], w[4]
                h = max(1.0, y1 - y0)
                fn, ff = pick_font(word)
                page.insert_text(fitz.Point(x0, y0 + h * 0.85), word,
                                 fontsize=h * 0.9, fontname=fn, fontfile=ff,
                                 render_mode=3, color=(1, 1, 1), overlay=False)
    commit(S, doc, f"OCR ({langs}) on {len(pages)} page(s)")
    doc.close()

    out_text = "".join(out_parts).strip()
    txt_path = save_temp(out_text.encode("utf-8"), ".txt", "ocr_text.txt")
    pdf_path = save_temp(S["bytes"], ".pdf", "searchable.pdf")
    return (out_text, txt_path, pdf_path, render_pil(S), gallery_pils(S),
            session_info(S),
            f"✅ OCR finished on {len(pages)} page(s). Text below is copyable — "
            f"download the searchable PDF or the .txt file.", S)


def extract_text_txt(S):
    doc = open_doc(S)
    parts = []
    for pno in range(doc.page_count):
        parts.append(f"\n===== PAGE {pno + 1} =====\n{doc[pno].get_text('text')}")
    doc.close()
    txt = "".join(parts).strip()
    path = save_temp(txt.encode("utf-8"), ".txt", "extracted_text.txt")
    return path, f"📄 Extracted {len(txt)} characters from all pages."


# --------------------------------------------------------------------------
#  Tab 3 — Edit text
# --------------------------------------------------------------------------
def _search_hits(page, find, case):
    hits = []
    for r in page.search_for(find):
        # ignore sub-pixel artifacts (e.g. from 1pt invisible text layers)
        if r.width < 2 or r.height < 2:
            continue
        if case:
            if page.get_textbox(r).strip() == find:
                hits.append(r)
        else:
            hits.append(r)
    return hits


def highlight_matches(find, case, S):
    if not find.strip():
        raise gr.Error("Type something to search for.")
    img = render_pil(S)
    doc = open_doc(S)
    page = doc[S["page"]]
    hits = _search_hits(page, find, case)
    doc.close()
    draw = ImageDraw.Draw(img)
    for r in hits:
        draw.rectangle([r.x0 * RENDER_ZOOM, r.y0 * RENDER_ZOOM,
                        r.x1 * RENDER_ZOOM, r.y1 * RENDER_ZOOM],
                       outline=(220, 38, 38), width=4)
    return img, f"🔎 {len(hits)} match(es) on page {S['page'] + 1} (highlighted in red)."


def replace_text(find, repl, case, current_only, S):
    if not find.strip():
        raise gr.Error("Enter the text to find.")
    doc = open_doc(S)
    pages = [S["page"]] if current_only else range(doc.page_count)
    total = 0
    for pno in pages:
        page = doc[pno]
        hits = _search_hits(page, find, case)
        if not hits:
            continue
        for r in hits:
            page.add_redact_annot(r, fill=(1, 1, 1))
        page.apply_redactions()
        fontname, fontfile = pick_font(repl)
        for r in hits:
            fs = max(6.0, r.height * 0.9)
            lines = repl.split("\n")
            y = r.y0 + fs
            for ln in lines:
                page.insert_text(fitz.Point(r.x0, y), ln, fontsize=fs,
                                 fontname=fontname, fontfile=fontfile, color=(0, 0, 0))
                y += fs * 1.25
        total += len(hits)
    commit(S, doc, f"Replaced {total} occurrence(s) of “{find}”")
    doc.close()
    return (render_pil(S), gallery_pils(S), session_info(S),
            f"✏️ Replaced {total} occurrence(s). Undo is available in the top bar.", S)


def add_text(text, font_key, size, color, x, y, S):
    if not text.strip():
        raise gr.Error("Enter the text to add.")
    doc = open_doc(S)
    page = doc[S["page"]]
    fontname, fontfile = resolve_font(font_key)
    try:
        page.insert_text(fitz.Point(float(x), float(y)), text, fontsize=float(size),
                         fontname=fontname, fontfile=fontfile, color=hex_to_rgb(color))
    except Exception as e:
        doc.close()
        raise gr.Error(f"Could not place text: {e}")
    commit(S, doc, f"Added text on page {S['page'] + 1}")
    doc.close()
    return (render_pil(S), gallery_pils(S), session_info(S),
            f"✅ Text added on page {S['page'] + 1} at ({x:.0f}, {y:.0f}).", S)


def on_preview_click(evt: gr.SelectData, S):
    if not S.get("bytes"):
        return 0.0, 0.0
    x, y = evt.index
    px, py = x / RENDER_ZOOM, y / RENDER_ZOOM
    w, h = page_dims(S)
    return round(min(max(px, 0), w), 1), round(min(max(py, 0), h), 1)


# --------------------------------------------------------------------------
#  Tab 4 — Page tools
# --------------------------------------------------------------------------
def gallery_select(evt: gr.SelectData, S):
    idx = int(evt.index)
    S["page"] = idx
    n, w, h = _count_dims(S)
    xup, yup = xy_slider_update(w, h)
    return (render_pil(S), session_info(S), xup, yup,
            pg_slider_update(n, idx), S)


def delete_pages(pages_str, S):
    doc = open_doc(S)
    n = doc.page_count
    idx = parse_page_ranges(pages_str, n)
    if not idx:
        doc.close()
        raise gr.Error("Enter valid page numbers, e.g. 1,3 or 2-5.")
    for i in sorted(idx, reverse=True):
        doc.delete_page(i)
    commit(S, doc, f"Deleted {len(idx)} page(s)")
    doc.close()
    n2, w, h = _count_dims(S)
    xup, yup = xy_slider_update(w, h)
    return (render_pil(S), gallery_pils(S), session_info(S),
            f"🗑️ Deleted page(s): {', '.join(str(i + 1) for i in idx)}",
            xup, yup, pg_slider_update(n2, S["page"]), S)


def rotate_pages(deg, all_pages, S):
    doc = open_doc(S)
    pages = range(doc.page_count) if all_pages else [S["page"]]
    for pno in pages:
        p = doc[pno]
        p.set_rotation((p.rotation + deg) % 360)
    commit(S, doc, f"Rotated {deg}° ({'all pages' if all_pages else 'current page'})")
    doc.close()
    return (render_pil(S), gallery_pils(S), session_info(S),
            f"🔄 Rotated {deg}°.", S)


def duplicate_page(S):
    doc = open_doc(S)
    pno = S["page"]
    doc.copy_page(pno)                          # appends a copy at the end
    doc.move_page(doc.page_count - 1, pno + 1)  # place it right after the original
    commit(S, doc, f"Duplicated page {pno + 1}")
    doc.close()
    n, w, h = _count_dims(S)
    xup, yup = xy_slider_update(w, h)
    return (render_pil(S), gallery_pils(S), session_info(S),
            f"📑 Duplicated page {pno + 1}.",
            xup, yup, pg_slider_update(n, S["page"]), S)


def move_page(dirn, S):
    doc = open_doc(S)
    pno = S["page"]
    target = pno + dirn
    if 0 <= target < doc.page_count:
        doc.move_page(pno, target)
        S["page"] = target
        commit(S, doc, f"Moved page {pno + 1} {'up' if dirn < 0 else 'down'}")
        doc.close()
        return (render_pil(S), gallery_pils(S), session_info(S), "↕️ Page moved.",
                gr.update(), gr.update(), gr.update(), S)
    doc.close()
    return (gr.update(), gr.update(), gr.update(), "Already at the edge.",
            gr.update(), gr.update(), gr.update(), S)


def reorder_pages(order_str, S):
    doc = open_doc(S)
    n = doc.page_count
    parts = [p.strip() for p in str(order_str).split(",") if p.strip()]
    try:
        order = [int(p) for p in parts]
    except ValueError:
        doc.close()
        raise gr.Error("Enter a comma-separated list of page numbers, e.g. 3,1,2,4")
    if sorted(order) != list(range(1, n + 1)):
        doc.close()
        raise gr.Error(f"Order must be a permutation of 1..{n} (all pages, no repeats).")
    nd = fitz.open()
    for p in order:
        nd.insert_pdf(doc, from_page=p - 1, to_page=p - 1)
    doc.close()
    commit_bytes(S, nd.tobytes(garbage=4, deflate=True), "Pages reordered")
    nd.close()
    n, w, h = _count_dims(S)
    xup, yup = xy_slider_update(w, h)
    return (render_pil(S), gallery_pils(S), session_info(S),
            f"🔀 Pages reordered to: {order_str}",
            xup, yup, pg_slider_update(n, S["page"]), S)


def split_pdf(chunk, S):
    doc = open_doc(S)
    n = doc.page_count
    chunk = max(1, int(chunk))
    paths = []
    nd = fitz.open()
    for pno in range(n):
        if pno > 0 and pno % chunk == 0:
            paths.append(save_temp(nd.tobytes(garbage=4, deflate=True), ".pdf",
                                   f"part_{len(paths) + 1}.pdf"))
            nd.close()
            nd = fitz.open()
        nd.insert_pdf(doc, from_page=pno, to_page=pno)
    if nd.page_count:
        paths.append(save_temp(nd.tobytes(garbage=4, deflate=True), ".pdf",
                               f"part_{len(paths) + 1}.pdf"))
    nd.close()
    doc.close()
    return paths, f"✂️ Split into {len(paths)} part(s) of up to {chunk} page(s) each."


def extract_pages(pages_str, S):
    doc = open_doc(S)
    n = doc.page_count
    idx = parse_page_ranges(pages_str, n)
    if not idx:
        doc.close()
        raise gr.Error("Enter valid page numbers, e.g. 1,3 or 2-5.")
    nd = fitz.open()
    for i in idx:
        nd.insert_pdf(doc, from_page=i, to_page=i)
    doc.close()
    path = save_temp(nd.tobytes(garbage=4, deflate=True), ".pdf", "extracted_pages.pdf")
    nd.close()
    return path, f"📦 Extracted {len(idx)} page(s)."


# --------------------------------------------------------------------------
#  Tab 5 — Merge & advanced
# --------------------------------------------------------------------------
def merge_files(files, S):
    if not files:
        raise gr.Error("Select at least two PDFs to merge.")
    paths = [f if isinstance(f, str) else f.name for f in files]
    nd = fitz.open()
    for path in paths:
        try:
            src = fitz.open(path)
        except Exception:
            raise gr.Error(f"Could not open {os.path.basename(path)}")
        nd.insert_pdf(src)
        src.close()
    n = nd.page_count
    if n == 0:
        nd.close()
        raise gr.Error("Nothing to merge.")
    commit_bytes(S, nd.tobytes(garbage=4, deflate=True), "Merged PDFs")
    S["name"] = "merged.pdf"
    path = save_temp(S["bytes"], ".pdf", "merged.pdf")
    nd.close()
    n, w, h = _count_dims(S)
    xup, yup = xy_slider_update(w, h)
    return (path, render_pil(S), gallery_pils(S), session_info(S),
            f"🔗 Merged {len(paths)} file(s) → {n} pages (now your working document).",
            xup, yup, pg_slider_update(n, S["page"]), S)


def add_watermark(text, size, color, opacity, angle, S):
    if not text.strip():
        raise gr.Error("Enter watermark text.")
    doc = open_doc(S)
    fontname, fontfile = pick_font(text)
    f = fitz.Font(fontname=fontname, fontfile=fontfile)
    rgb = hex_to_rgb(color)
    rad = math.radians(angle)
    cos, sin = math.cos(rad), math.sin(rad)
    m = fitz.Matrix(cos, sin, -sin, cos, 0, 0)
    for pno in range(doc.page_count):
        page = doc[pno]
        rect = page.rect
        tw = fitz.TextWriter(rect, color=rgb)
        cx, cy = rect.width / 2, rect.height / 2
        tw.append(fitz.Point(cx, cy), text, font=f, fontsize=float(size))
        tw.write_text(page, morph=(fitz.Point(cx, cy), m),
                      opacity=float(opacity), overlay=True)
    commit(S, doc, "Watermark added to all pages")
    doc.close()
    return (render_pil(S), gallery_pils(S), session_info(S),
            f"💧 Watermark “{text}” applied to all pages.", S)


def compress_pdf(S):
    doc = open_doc(S)
    before = len(S["bytes"])
    try:
        doc.subset_fonts()
    except Exception:
        pass
    data = doc.tobytes(garbage=4, deflate=True)
    doc.close()
    after = len(data)
    pct = (1 - after / before) * 100 if before else 0
    commit_bytes(S, data, "Compressed")
    path = save_temp(data, ".pdf", "compressed.pdf")
    return (path, render_pil(S), gallery_pils(S), session_info(S),
            f"🗜️ Compressed: {fmt_size(before)} → {fmt_size(after)} "
            f"({pct:+.1f}%). Download below.", S)


def encrypt_pdf(pw, S):
    if not pw or len(pw) < 3:
        raise gr.Error("Password must be at least 3 characters.")
    doc = open_doc(S)
    data = doc.tobytes(garbage=4, deflate=True, encryption=fitz.PDF_ENCRYPT_AES_256,
                       user_pw=pw, owner_pw=pw)
    doc.close()
    commit_bytes(S, data, "Encrypted with password")
    S["pw"] = pw
    path = save_temp(data, ".pdf", "protected.pdf")
    return (path, render_pil(S), gallery_pils(S), session_info(S),
            f"🔐 Encrypted (AES-256). Password: “{pw}” — keep it safe!", S)


def decrypt_pdf(pw, S):
    doc = open_doc(S)
    if doc.needs_pass:
        if not doc.authenticate(pw or ""):
            doc.close()
            raise gr.Error("❌ Wrong password.")
    data = doc.tobytes(garbage=4, deflate=True)
    doc.close()
    commit_bytes(S, data, "Password removed")
    S["pw"] = ""
    path = save_temp(data, ".pdf", "unlocked.pdf")
    return (path, render_pil(S), gallery_pils(S), session_info(S),
            "🔓 Password removed — file is now open.", S)


# --------------------------------------------------------------------------
#  Tab 6 — Forms
# --------------------------------------------------------------------------
def list_forms(S):
    doc = open_doc(S)
    fields = []
    for pno in range(doc.page_count):
        for w in doc[pno].widgets():
            name = w.field_name or f"field_{pno + 1}"
            fields.append({"label": f"{name}  ·  p{pno + 1} ({w.field_type_string})",
                           "pno": pno, "name": name})
    doc.close()
    S["form_fields"] = fields
    choices = [f["label"] for f in fields]
    if not fields:
        return gr.update(choices=[], value=None), "📋 No form fields found in this PDF."
    return (gr.update(choices=choices, value=choices[0]),
            f"📋 Found {len(fields)} form field(s). Pick one, type a value, press Apply.")


def set_form(label, value, S):
    fields = S.get("form_fields") or []
    if not fields:
        raise gr.Error("Click “Scan form fields” first.")
    f = next((x for x in fields if x["label"] == label), None)
    if not f:
        raise gr.Error("Pick a field from the list.")
    doc = open_doc(S)
    page = doc[f["pno"]]
    done = False
    for w in page.widgets():
        if w.field_name == f["name"]:
            w.field_value = value
            w.update()
            done = True
            break
    if not done:
        doc.close()
        raise gr.Error("Field not found on the page.")
    commit(S, doc, f"Filled form field “{f['name']}”")
    doc.close()
    return (render_pil(S), gallery_pils(S), session_info(S),
            f"✅ Field “{f['name']}” set to “{value}”.", S)


# ==========================================================================
#  UI
# ==========================================================================
CSS = """
.gradio-container {max-width: 1320px !important; margin: 0 auto;}
#hd {background: linear-gradient(135deg,#4f46e5 0%,#7c3aed 55%,#2563eb 100%);
     border-radius: 20px; padding: 26px 30px; color:#fff; margin: 8px 0 16px;
     box-shadow: 0 10px 30px -12px rgba(79,70,229,.55);}
#hd h1 {margin:0; font-size:32px; font-weight:800; letter-spacing:-.5px;}
#hd p {margin:8px 0 0; opacity:.95; font-size:15px;}
#hd .sub {margin-top:14px;}
.badge {display:inline-block; background:rgba(255,255,255,.16); border:1px solid rgba(255,255,255,.35);
        border-radius:999px; padding:3px 13px; font-size:12.5px; margin:0 6px 6px 0;}
#toolbar {display:flex; align-items:center; gap:10px; flex-wrap:wrap; padding:10px 4px;}
footer {display:none !important;}
"""

theme = gr.themes.Soft(
    primary_hue="indigo",
    secondary_hue="blue",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
)


def build_demo():
    with gr.Blocks(title="PDF EDITER — Advanced PDF Editor with OCR") as demo:
        state = gr.State(value=new_state())

        # ---------------- header ----------------
        with gr.Column(elem_id="hd"):
            gr.HTML("""
            <h1>📝 PDF EDITER</h1>
            <p>Advanced PDF editor with <b>OCR</b> — upload, edit text, manage pages,
            merge, watermark, compress &amp; protect. Powered by PyMuPDF + Tesseract.</p>
            <div class="sub">
                <span class="badge">🔍 OCR · searchable PDF</span>
                <span class="badge">✏️ Find &amp; Replace</span>
                <span class="badge">➕ Add text (multi-language)</span>
                <span class="badge">🖼️ Page tools</span>
                <span class="badge">🔗 Merge / Split</span>
                <span class="badge">💧 Watermark</span>
                <span class="badge">🗜️ Compress</span>
                <span class="badge">🔐 Protect</span>
                <span class="badge">📋 Forms</span>
            </div>
            """)

        # ---------------- top toolbar ----------------
        with gr.Row(elem_id="toolbar"):
            session_info = gr.Markdown("No document loaded", elem_id="status")
            undo_btn = gr.Button("↩️ Undo", scale=0)
            redo_btn = gr.Button("↪️ Redo", scale=0)
            reset_btn = gr.Button("🧹 Reset", scale=0)
            dl_btn = gr.Button("⬇️ Download PDF", variant="primary", scale=0)
            dl_file = gr.File(visible=False)
        status = gr.Markdown("👋 Welcome! Upload a PDF below, or press “Load demo PDF” to explore.")

        # ---------------- tabs ----------------
        with gr.Tabs():
            # ================= TAB 1 : UPLOAD =================
            with gr.Tab("📥 Upload & Preview"):
                with gr.Row():
                    with gr.Column(scale=1):
                        upload = gr.File(file_count="multiple", file_types=[".pdf"],
                                         label="Upload one or more PDFs (multiple = merged)")
                        load_btn = gr.Button("📂 Load PDF(s)", variant="primary")
                        demo_btn = gr.Button("✨ Load demo PDF (try everything instantly)")
                        gr.Markdown("### Navigate")
                        prev_btn = gr.Button("◀ Previous page")
                        page_slider = gr.Slider(1, 100, step=1, value=1, label="Go to page (1-based)")
                        next_btn = gr.Button("Next page ▶")
                    with gr.Column(scale=2):
                        preview = gr.Image(label="Page preview — click to set text position",
                                           interactive=False, height=680)
                gr.Markdown("💡 **Tip:** upload a scanned PDF and go to the **OCR** tab to make it "
                            "searchable; then **Edit Text** works on the extracted text.")

            # ================= TAB 2 : OCR =================
            with gr.Tab("🔍 OCR"):
                with gr.Row():
                    with gr.Column(scale=1):
                        ocr_langs = gr.Dropdown(choices=list(OCR_LANGS.keys()),
                                                value=["English"], multiselect=True,
                                                label="OCR language(s) — hold Ctrl/Cmd to multi-select")
                        ocr_dpi = gr.Slider(150, 400, value=200, step=10,
                                            label="Scan resolution (DPI — higher = slower but more accurate)")
                        ocr_pages = gr.Radio(["All pages", "Current page only"], value="All pages",
                                             label="Pages to OCR")
                        ocr_btn = gr.Button("🧠 Run OCR — make it searchable", variant="primary")
                        extract_btn = gr.Button("📄 Extract text as .txt (no OCR)")
                        ocr_txt_dl = gr.File(label="OCR / extracted text (.txt)")
                        ocr_pdf_dl = gr.File(label="Searchable PDF (download)")
                        gr.Markdown("⏱️ OCR runs on CPU — large documents take a while. "
                                    "Start with 1–2 pages to check quality.")
                    with gr.Column(scale=2):
                        ocr_out = gr.Textbox(label="Recognized text (editable)", lines=24)

            # ================= TAB 3 : EDIT TEXT =================
            with gr.Tab("✏️ Edit Text"):
                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("### 🔎 Find & Replace")
                        find_tb = gr.Textbox(label="Text to find", placeholder="e.g. apples")
                        repl_tb = gr.Textbox(label="Replace with", placeholder="e.g. oranges")
                        case_cb = gr.Checkbox(label="Case-sensitive", value=False)
                        cur_cb = gr.Checkbox(label="Current page only", value=False)
                        with gr.Row():
                            find_btn = gr.Button("🔎 Highlight matches")
                            repl_btn = gr.Button("✏️ Replace", variant="primary")
                        gr.Markdown("---")
                        gr.Markdown("### ➕ Add text (multi-language)")
                        add_tb = gr.Textbox(label="Text", placeholder="नमस्ते / Hello / مرحبا …",
                                            value="Hello")
                        font_dd = gr.Dropdown(choices=list(FONTS.keys()),
                                              value="Noto Sans (Latin/Cyrillic/Greek)",
                                              label="Font")
                        with gr.Row():
                            size_sl = gr.Slider(4, 96, value=14, step=1, label="Font size")
                            color_pk = gr.ColorPicker(value="#111111", label="Color")
                        with gr.Row():
                            x_sl = gr.Slider(0, 595, value=72, step=1, label="X (pt)")
                            y_sl = gr.Slider(0, 842, value=100, step=1, label="Y (pt)")
                        add_btn = gr.Button("➕ Place text on this page", variant="primary")
                        gr.Markdown("💡 Click anywhere on the preview image to auto-fill X/Y.")
                    with gr.Column(scale=2):
                        preview2 = gr.Image(label="Preview", interactive=False, height=680)

            # ================= TAB 4 : PAGE TOOLS =================
            with gr.Tab("🖼️ Page Tools"):
                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("### Selected page (click a thumbnail)")
                        sel_info = gr.Markdown("Click a thumbnail to select the working page.")
                        with gr.Row():
                            rot_cw = gr.Button("⟳ Rotate 90° CW")
                            rot_ccw = gr.Button("⟲ Rotate 90° CCW")
                        rot_all_cb = gr.Checkbox(label="Apply rotation to ALL pages", value=False)
                        with gr.Row():
                            dup_btn = gr.Button("📑 Duplicate page")
                            del_btn = gr.Button("🗑️ Delete page(s)")
                        del_tb = gr.Textbox(label="Pages to delete (1-based, e.g. 2,5 or 3-6)",
                                            placeholder="2,5")
                        with gr.Row():
                            mv_up = gr.Button("⬆️ Move up")
                            mv_dn = gr.Button("⬇️ Move down")
                        gr.Markdown("---")
                        gr.Markdown("### 🔀 Reorder / Split / Extract")
                        order_tb = gr.Textbox(label="New page order (permutation of 1..N)",
                                              placeholder="3,1,2,4")
                        order_btn = gr.Button("🔀 Apply order")
                        split_sl = gr.Slider(1, 50, value=1, step=1, label="Split after every N pages")
                        split_btn = gr.Button("✂️ Split PDF")
                        split_files = gr.File(file_count="multiple", label="Split parts")
                        extract_tb = gr.Textbox(label="Extract pages (e.g. 1,3 or 2-5)",
                                                placeholder="1,3")
                        extract_btn = gr.Button("📦 Extract pages")
                        extract_file = gr.File(label="Extracted PDF")
                    with gr.Column(scale=2):
                        gallery = gr.Gallery(label="All pages — click to select", columns=3,
                                             rows=2, height=520, object_fit="contain")
                        preview3 = gr.Image(label="Selected page preview", interactive=False, height=500)

            # ================= TAB 5 : MERGE & ADVANCED =================
            with gr.Tab("🧩 Merge & Advanced"):
                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("### 🔗 Merge PDFs")
                        merge_files_in = gr.File(file_count="multiple", file_types=[".pdf"],
                                                 label="Pick PDFs to merge (in order)")
                        merge_btn = gr.Button("🔗 Merge into working document", variant="primary")
                        merge_dl = gr.File(label="Merged PDF")
                        gr.Markdown("---")
                        gr.Markdown("### 💧 Watermark")
                        wm_tb = gr.Textbox(label="Watermark text", value="CONFIDENTIAL")
                        wm_size = gr.Slider(10, 120, value=48, step=2, label="Size")
                        wm_color = gr.ColorPicker(value="#999999", label="Color")
                        wm_opa = gr.Slider(0.05, 1.0, value=0.25, step=0.05, label="Opacity")
                        wm_ang = gr.Slider(-90, 90, value=-45, step=5, label="Angle (°)")
                        wm_btn = gr.Button("💧 Apply watermark to all pages")
                    with gr.Column(scale=1):
                        gr.Markdown("### 🗜️ Compress & 🔐 Protect")
                        comp_btn = gr.Button("🗜️ Compress PDF (garbage collect + subset fonts)")
                        comp_dl = gr.File(label="Compressed PDF")
                        pw_tb = gr.Textbox(label="Password", type="password")
                        with gr.Row():
                            enc_btn = gr.Button("🔐 Encrypt (AES-256)")
                            dec_btn = gr.Button("🔓 Unlock / remove password")
                        prot_dl = gr.File(label="Protected / unlocked PDF")
                        gr.Markdown("---")
                        gr.Markdown("### 📋 Fill PDF forms")
                        scan_form_btn = gr.Button("📋 Scan form fields")
                        form_dd = gr.Dropdown(choices=[], label="Field")
                        form_val = gr.Textbox(label="Value")
                        form_set_btn = gr.Button("✅ Set field value")
                    with gr.Column(scale=2):
                        preview4 = gr.Image(label="Preview", interactive=False, height=680)

        # ================= events =================
        # tab1
        load_btn.click(load_files, [upload, state],
                       [preview, gallery, session_info, status, x_sl, y_sl, page_slider, state])
        demo_btn.click(load_demo, [state],
                       [preview, gallery, session_info, status, x_sl, y_sl, page_slider, state])
        prev_btn.click(nav_delta, [gr.State(-1), state],
                       [preview, session_info, x_sl, y_sl, page_slider, state])
        next_btn.click(nav_delta, [gr.State(1), state],
                       [preview, session_info, x_sl, y_sl, page_slider, state])
        page_slider.change(nav_go, [page_slider, state],
                           [preview, session_info, x_sl, y_sl, page_slider, state])

        # top toolbar
        undo_btn.click(undo_action, [state],
                       [preview, gallery, session_info, status, x_sl, y_sl, page_slider, state])
        redo_btn.click(redo_action, [state],
                       [preview, gallery, session_info, status, x_sl, y_sl, page_slider, state])
        reset_btn.click(reset_all, [state],
                        [preview, gallery, session_info, status, x_sl, y_sl, page_slider, state])
        dl_btn.click(download_current, [state], [dl_file, status])

        # tab2
        ocr_btn.click(run_ocr, [ocr_langs, ocr_dpi, ocr_pages, state],
                      [ocr_out, ocr_txt_dl, ocr_pdf_dl, preview, gallery,
                       session_info, status, state],
                      api_name="ocr")
        extract_btn.click(extract_text_txt, [state], [ocr_txt_dl, status])

        # tab3
        find_btn.click(highlight_matches, [find_tb, case_cb, state], [preview2, status])
        repl_btn.click(replace_text, [find_tb, repl_tb, case_cb, cur_cb, state],
                       [preview2, gallery, session_info, status, state], api_name="replace")
        add_btn.click(add_text, [add_tb, font_dd, size_sl, color_pk, x_sl, y_sl, state],
                      [preview2, gallery, session_info, status, state], api_name="add_text")
        preview2.select(on_preview_click, [state], [x_sl, y_sl])

        # tab4
        gallery.select(gallery_select, [state],
                       [preview3, sel_info, x_sl, y_sl, page_slider, state])
        rot_cw.click(rotate_pages, [gr.State(90), rot_all_cb, state],
                     [preview3, gallery, session_info, status, state])
        rot_ccw.click(rotate_pages, [gr.State(-90), rot_all_cb, state],
                      [preview3, gallery, session_info, status, state])
        dup_btn.click(duplicate_page, [state],
                      [preview3, gallery, session_info, status, x_sl, y_sl, page_slider, state])
        del_btn.click(delete_pages, [del_tb, state],
                      [preview3, gallery, session_info, status, x_sl, y_sl, page_slider, state])
        mv_up.click(move_page, [gr.State(-1), state],
                    [preview3, gallery, session_info, status, x_sl, y_sl, page_slider, state])
        mv_dn.click(move_page, [gr.State(1), state],
                    [preview3, gallery, session_info, status, x_sl, y_sl, page_slider, state])
        order_btn.click(reorder_pages, [order_tb, state],
                        [preview3, gallery, session_info, status, x_sl, y_sl, page_slider, state])
        split_btn.click(split_pdf, [split_sl, state], [split_files, status])
        extract_btn.click(extract_pages, [extract_tb, state], [extract_file, status])

        # tab5
        merge_btn.click(merge_files, [merge_files_in, state],
                        [merge_dl, preview4, gallery, session_info, status,
                         x_sl, y_sl, page_slider, state])
        wm_btn.click(add_watermark, [wm_tb, wm_size, wm_color, wm_opa, wm_ang, state],
                     [preview4, gallery, session_info, status, state])
        comp_btn.click(compress_pdf, [state],
                       [comp_dl, preview4, gallery, session_info, status, state])
        enc_btn.click(encrypt_pdf, [pw_tb, state],
                      [prot_dl, preview4, gallery, session_info, status, state])
        dec_btn.click(decrypt_pdf, [pw_tb, state],
                      [prot_dl, preview4, gallery, session_info, status, state])
        scan_form_btn.click(list_forms, [state], [form_dd, status])
        form_set_btn.click(set_form, [form_dd, form_val, state],
                           [preview4, gallery, session_info, status, state])

    return demo


demo = build_demo()

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(
        theme=theme,
        css=CSS,
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
    )