# ── Fix Python 3.13 asyncio bug ───────────────────────────────────────────────
import asyncio.base_events
_orig_del = asyncio.base_events.BaseEventLoop.__del__
def _patched_del(self):
    try: _orig_del(self)
    except Exception: pass
asyncio.base_events.BaseEventLoop.__del__ = _patched_del

# ── Standard library ──────────────────────────────────────────────────────────
import warnings, logging, os, json, re, base64, copy, hashlib
from io import BytesIO
warnings.filterwarnings("ignore")
os.environ["TOKENIZERS_PARALLELISM"] = "false"
logging.getLogger("transformers").setLevel(logging.ERROR)
logging.getLogger("reportlab").setLevel(logging.ERROR)

# ── Third-party ───────────────────────────────────────────────────────────────
import gradio as gr
import fitz
import numpy as np
from pdf2image import convert_from_path
from PIL import Image

# ── ReportLab ─────────────────────────────────────────────────────────────────
from reportlab.pdfgen import canvas
from reportlab.lib.colors import HexColor
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont

# ── Constants ─────────────────────────────────────────────────────────────────
FONT_CACHE_DIR = "/tmp/pdf_fonts"
os.makedirs(FONT_CACHE_DIR, exist_ok=True)
FONT_REGISTRY  = {}

# Stricter placeholder regex — only match [Text with Letters]
PLACEHOLDER_RE = re.compile(r'\[([A-Za-z][A-Za-z0-9\s,\-]+)\]')

STATE = {
    "pages_data": None,
    "page_images": None,
    "placeholders": [],
    "pdf_path": None,
}


# ─────────────────────────────────────────────────────────────────────────────
# HELPERS
# ─────────────────────────────────────────────────────────────────────────────

def color_int_to_rgb(color_int: int) -> tuple:
    r = (color_int >> 16) & 0xFF
    g = (color_int >> 8) & 0xFF
    b = color_int & 0xFF
    return (r/255.0, g/255.0, b/255.0)


def is_repeatable(key: str) -> bool:
    k = key.lower()
    return any(x in k for x in ["child", "beneficiar", "heir", "trustee", "asset", "property"])


def infer_field_type(key: str) -> str:
    k = key.lower()
    if any(x in k for x in ["date", "dob", "day", "year", "month"]): return "date"
    if any(x in k for x in ["email"]): return "email"
    if any(x in k for x in ["phone", "mobile"]): return "tel"
    return "text"


def image_to_base64(img: Image.Image) -> str:
    buf = BytesIO()
    img.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode()


def pdf_to_images(pdf_path: str, dpi: int = 120):
    return convert_from_path(pdf_path, dpi=dpi)


# ─────────────────────────────────────────────────────────────────────────────
# FONT EXTRACTION & REGISTRATION
# ─────────────────────────────────────────────────────────────────────────────

def _fallback_font(name: str) -> str:
    n = name.lower()
    if "bold" in n and ("italic" in n or "oblique" in n): return "Helvetica-BoldOblique"
    if "bold" in n:    return "Helvetica-Bold"
    if "italic" in n or "oblique" in n: return "Helvetica-Oblique"
    if "times" in n:   return "Times-Roman"
    if "courier" in n or "mono" in n: return "Courier"
    return "Helvetica"


def extract_and_register_fonts(pdf_path: str):
    global FONT_REGISTRY
    doc = fitz.open(pdf_path)

    for page in doc:
        for font in page.get_fonts(full=True):
            xref      = font[0]
            base_name = font[3]
            if xref == 0 or base_name in FONT_REGISTRY:
                continue
            clean = base_name.split("+")[-1] if "+" in base_name else base_name
            try:
                font_data = doc.extract_font(xref)
                if not font_data or not font_data[3]:
                    FONT_REGISTRY[base_name] = _fallback_font(clean)
                    continue
                raw   = font_data[3]
                ext   = (font_data[1] or "ttf").lower()
                if ext in ("cff", "type1c"): ext = "otf"
                if not ext: ext = "ttf"

                fhash    = hashlib.md5(raw).hexdigest()[:8]
                fname    = f"{re.sub(r'[^A-Za-z0-9_]','_', clean)}_{fhash}.{ext}"
                fpath    = os.path.join(FONT_CACHE_DIR, fname)
                with open(fpath, "wb") as f:
                    f.write(raw)

                if ext in ("ttf", "otf"):
                    reg = re.sub(r'[^A-Za-z0-9]', '_', clean)
                    if reg not in [v for v in FONT_REGISTRY.values()]:
                        pdfmetrics.registerFont(TTFont(reg, fpath))
                    FONT_REGISTRY[base_name] = reg
                    print(f"  Font registered: {reg}")
                else:
                    FONT_REGISTRY[base_name] = _fallback_font(clean)
            except Exception as e:
                print(f"  Font extract failed ({clean}): {e}")
                FONT_REGISTRY[base_name] = _fallback_font(clean)

    doc.close()
    print(f"Total fonts registered: {len(FONT_REGISTRY)}")


def resolve_font(pdf_font_name: str) -> str:
    if pdf_font_name in FONT_REGISTRY:
        return FONT_REGISTRY[pdf_font_name]
    clean = pdf_font_name.split("+")[-1] if "+" in pdf_font_name else pdf_font_name
    for k, v in FONT_REGISTRY.items():
        if clean in k:
            return v
    return _fallback_font(pdf_font_name)


# ─────────────────────────────────────────────────────────────────────────────
# PDF EXTRACTION — Store EXACT positions
# ─────────────────────────────────────────────────────────────────────────────

def extract_pdf(pdf_path: str) -> list:
    doc = fitz.open(pdf_path)
    pages = []
    for page in doc:
        w, h   = page.rect.width, page.rect.height
        data   = page.get_text("dict", flags=fitz.TEXT_PRESERVE_WHITESPACE)
        elems  = []
        for block in data.get("blocks", []):
            if block.get("type") != 0:
                continue
            for line in block.get("lines", []):
                for span in line.get("spans", []):
                    txt = span.get("text", "")
                    if not txt.strip():
                        continue
                    x0, y0, x1, y1 = span["bbox"]
                    fl = span.get("flags", 0)
                    elems.append({
                        "content":   txt,
                        "font_size": round(span.get("size", 12), 2),
                        "font_name": span.get("font", "Helvetica"),
                        "bold":      bool(fl & 16),
                        "italic":    bool(fl & 2),
                        "color":     span.get("color", 0),
                        "bbox":      {"x1": x0, "y1": y0, "x2": x1, "y2": y1},
                    })
        pages.append({
            "page_number": page.number + 1,
            "dimensions":  {"width": w, "height": h},
            "elements":    elems,
        })
    doc.close()
    return pages


# ─────────────────────────────────────────────────────────────────────────────
# PLACEHOLDER DETECTION — Stricter rules
# ─────────────────────────────────────────────────────────────────────────────

def auto_detect_placeholders(pages_data: list) -> list:
    found = {}
    for page in pages_data:
        for elem in page["elements"]:
            for m in PLACEHOLDER_RE.findall(elem["content"]):
                key = m.strip()
                # Ignore pure underscores/symbols
                if not any(c.isalpha() for c in key):
                    continue
                if key not in found:
                    found[key] = {
                        "key":        key,
                        "question":   f"Enter {key}:",
                        "field_type": infer_field_type(key),
                        "repeatable": is_repeatable(key),
                        "default":    "",
                    }
    return list(found.values())


# ─────────────────────────────────────────────────────────────────────────────
# HTML PREVIEW
# ─────────────────────────────────────────────────────────────────────────────

def build_preview_html(pages_data: list, page_images: list) -> str:
    html = []
    for page_data, img in zip(pages_data, page_images):
        b64   = image_to_base64(img)
        w = page_data["dimensions"]["width"]
        h = page_data["dimensions"]["height"]
        elems = ""
        for e in page_data["elements"]:
            bbox = e["bbox"]
            x1, y1, x2, y2 = bbox["x1"], bbox["y1"], bbox["x2"], bbox["y2"]
            left = (x1 / w) * 100
            top  = (y1 / h) * 100
            width = ((x2 - x1) / w) * 100
            height = ((y2 - y1) / h) * 100
            
            fs  = e.get("font_size", 12)
            txt = (e["content"].replace("&","&")
                               .replace("<","&lt;")
                               .replace(">","&gt;"))
            has_ph = bool(PLACEHOLDER_RE.search(e["content"]))
            bg = "rgba(255,220,0,0.5)" if has_ph else "transparent"
            elems += (
                f'<div style="position:absolute;'
                f'left:{left:.3f}%;top:{top:.3f}%;'
                f'width:{width:.3f}%;height:{height:.3f}%;'
                f'font-size:{fs}px;font-family:Arial,sans-serif;'
                f'white-space:pre-wrap;overflow:visible;'
                f'color:transparent;background:{bg};'
                f'cursor:text;user-select:text;" title="{txt}">{txt}</div>'
            )
        html.append(
            f'<div style="position:relative;width:100%;max-width:850px;'
            f'margin:0 auto 20px;box-shadow:0 2px 8px rgba(0,0,0,0.2);">'
            f'<img src="data:image/png;base64,{b64}" style="width:100%;display:block;"/>'
            f'<div style="position:absolute;top:0;left:0;width:100%;height:100%;">'
            f'{elems}</div></div>'
        )
    return (
        '<div style="background:#ddd;padding:12px;">'
        '<p style="text-align:center;font-family:Arial;font-size:13px;">'
        '🟡 Yellow = detected placeholders (underscores filtered out)</p>'
        + "".join(html) + "</div>"
    )


# ─────────────────────────────────────────────────────────────────────────────
# PDF REGENERATION — ABSOLUTE POSITIONING (exact layout preservation)
# ─────────────────────────────────────────────────────────────────────────────

def handle_repeatable_block(pages_data: list, placeholders: list, answers: dict) -> list:
    """
    Clone pages for each repeatable child/beneficiary.
    Returns: expanded pages_data with substituted text.
    """
    # Find which placeholders are repeatable
    repeat_map = {}
    for p in placeholders:
        if p.get("repeatable"):
            key = p["key"]
            bracket = f"[{key}]"
            ans = answers.get(key, "")
            items = [x.strip() for x in ans.split(";") if x.strip()] if ans else [bracket]
            repeat_map[bracket] = items

    if not repeat_map:
        # No repeatable fields — just return with substitutions
        return apply_substitutions(pages_data, placeholders, answers)

    # Find pages containing repeatable placeholders
    pages_with_repeats = set()
    for pid, page in enumerate(pages_data):
        for elem in page["elements"]:
            for rk in repeat_map.keys():
                if rk in elem["content"]:
                    pages_with_repeats.add(pid)
                    break

    if not pages_with_repeats:
        return apply_substitutions(pages_data, placeholders, answers)

    # For now, simple approach: duplicate ENTIRE document for each repeat item
    # This ensures continuation text flows properly
    result_pages = []
    first_repeat_key = list(repeat_map.keys())[0]
    repeat_items = repeat_map[first_repeat_key]

    for item_idx, item_value in enumerate(repeat_items):
        # Clone all pages
        cloned = copy.deepcopy(pages_data)
        # Substitute this specific item
        single_item_answers = {**answers}
        for rk in repeat_map.keys():
            # Extract the key name
            key_name = rk[1:-1]  # remove [ ]
            single_item_answers[key_name] = item_value
        cloned = apply_substitutions(cloned, placeholders, single_item_answers)
        result_pages.extend(cloned)

    return result_pages


def apply_substitutions(pages_data: list, placeholders: list, answers: dict) -> list:
    """Apply text replacements to all elements."""
    result = copy.deepcopy(pages_data)
    replacement_map = {}
    for p in placeholders:
        key = p["key"]
        bracket = f"[{key}]"
        ans = answers.get(key, "")
        replacement_map[bracket] = str(ans) if ans else bracket

    for page in result:
        for elem in page["elements"]:
            for placeholder, replacement in replacement_map.items():
                elem["content"] = elem["content"].replace(placeholder, replacement)

    return result


def regenerate_pdf_absolute(pages_data: list, placeholders: list, answers: dict, out_path: str) -> str:
    """
    Regenerate PDF using Canvas with ABSOLUTE positioning.
    This preserves exact layout, fonts, sizes, colors, positions.
    """
    # Handle repeatable sections by duplicating pages
    expanded_pages = handle_repeatable_block(pages_data, placeholders, answers)

    c = canvas.Canvas(out_path)

    for page in expanded_pages:
        pw = page["dimensions"]["width"]
        ph = page["dimensions"]["height"]
        c.setPageSize((pw, ph))

        for elem in page["elements"]:
            bbox = elem["bbox"]
            x = bbox["x1"]
            y = ph - bbox["y2"]  # PDF coordinate system: origin at bottom-left

            # Font
            font_name = elem.get("font_name", "Helvetica")
            rl_font   = resolve_font(font_name)
            font_size = elem.get("font_size", 12)

            # Color
            r, g, b = color_int_to_rgb(elem.get("color", 0))
            c.setFillColorRGB(r, g, b)

            # Set font
            try:
                c.setFont(rl_font, font_size)
            except:
                c.setFont("Helvetica", font_size)

            # Draw text
            text = elem["content"]
            try:
                c.drawString(x, y, text)
            except:
                # Fallback for special characters
                c.drawString(x, y, text.encode('latin-1', 'replace').decode('latin-1'))

        c.showPage()

    c.save()
    return out_path


# ─────────────────────────────────────────────────────────────────────────────
# GRADIO HANDLERS
# ─────────────────────────────────────────────────────────────────────────────

def admin_upload(pdf_file, progress=gr.Progress()):
    if pdf_file is None:
        return "<p>No file uploaded.</p>", "[]", gr.update(visible=False)
    try:
        pdf_path = pdf_file if isinstance(pdf_file, str) else pdf_file.name
        STATE["pdf_path"] = pdf_path

        progress(0.1, desc="Extracting fonts...")
        extract_and_register_fonts(pdf_path)

        progress(0.3, desc="Extracting text...")
        pages = extract_pdf(pdf_path)
        STATE["pages_data"] = pages

        progress(0.5, desc="Detecting placeholders...")
        phs = auto_detect_placeholders(pages)
        STATE["placeholders"] = phs

        progress(0.75, desc="Building preview...")
        imgs = pdf_to_images(pdf_path)
        STATE["page_images"] = imgs

        preview  = build_preview_html(pages, imgs)
        ph_json  = json.dumps(phs, indent=2)

        progress(1.0, desc="Done!")
        return preview, ph_json, gr.update(visible=True)
    except Exception as e:
        import traceback
        msg = traceback.format_exc()
        print(msg)
        return f"<p style='color:red'>Error: {e}</p>", "[]", gr.update(visible=False)


def admin_save_config(ph_json_str: str):
    try:
        phs = json.loads(ph_json_str)
        STATE["placeholders"] = phs
        with open("form_config.json", "w") as f:
            json.dump({"placeholders": phs, "pages_data": STATE["pages_data"]}, f, indent=2)
        return f"✅ Saved {len(phs)} placeholder(s). Users can now fill the form."
    except Exception as e:
        return f"❌ Error saving config: {e}"


def load_form():
    try:
        with open("form_config.json") as f:
            cfg = json.load(f)
        STATE["pages_data"]  = cfg["pages_data"]
        STATE["placeholders"] = cfg["placeholders"]
        phs = cfg["placeholders"]
        defaults = {p["key"]: p.get("default", "") for p in phs}
        return "✅ Form loaded!", phs, json.dumps(defaults, indent=2)
    except FileNotFoundError:
        return "⚠️ No form configured yet. Admin must upload a PDF first.", [], "{}"
    except Exception as e:
        return f"❌ Error: {e}", [], "{}"


def generate_pdf(answers_json: str, progress=gr.Progress()):
    try:
        answers = json.loads(answers_json)
    except Exception as e:
        return None, f"<p style='color:red'>Invalid JSON: {e}</p>"

    phs        = STATE.get("placeholders", [])
    pages_data = STATE.get("pages_data")

    if not pages_data:
        return None, "<p style='color:red'>No PDF loaded. Ask admin to upload first.</p>"

    try:
        progress(0.3, desc="Rebuilding PDF with exact layout...")
        # Re-register fonts if needed
        pdf_path = STATE.get("pdf_path")
        if pdf_path and os.path.exists(pdf_path):
            extract_and_register_fonts(pdf_path)

        out_path = "/tmp/filled_output.pdf"
        regenerate_pdf_absolute(pages_data, phs, answers, out_path)

        progress(0.8, desc="Rendering preview...")
        imgs   = pdf_to_images(out_path)
        pages_preview = "".join(
            f'<div style="margin:10px auto;max-width:800px;box-shadow:0 2px 8px rgba(0,0,0,.2);">'
            f'<img src="data:image/png;base64,{image_to_base64(img)}" style="width:100%;display:block;"/>'
            f'</div>'
            for img in imgs
        )
        preview = (
            f'<div style="background:#ddd;padding:12px;">'
            f'<h3 style="text-align:center;font-family:Arial;">✅ Generated — {len(imgs)} page(s)</h3>'
            f'{pages_preview}</div>'
        )
        progress(1.0, desc="Done!")
        return out_path, preview
    except Exception as e:
        import traceback
        return None, f"<pre style='color:red'>{traceback.format_exc()}</pre>"


# ─────────────────────────────────────────────────────────────────────────────
# GRADIO UI
# ─────────────────────────────────────────────────────────────────────────────

with gr.Blocks(title="Smart PDF Form Builder", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 📄 Smart PDF Form Builder")
    gr.Markdown("**Admin** uploads & configures → **Users** fill & download exact-layout PDF")

    with gr.Tabs():

        # ── ADMIN ──────────────────────────────────────────────────────────
        with gr.Tab("🔧 Admin"):
            gr.Markdown("### Upload PDF Template")
            a_file = gr.File(label="Upload PDF", file_types=[".pdf"], type="filepath")
            a_btn  = gr.Button("Extract & Detect Placeholders", variant="primary")
            a_prev = gr.HTML(label="Preview")

            with gr.Accordion("Edit Placeholder Config", open=True, visible=False) as a_acc:
                gr.Markdown("""
                Edit JSON — fields:
                - `key`: text inside `[ ]` in your PDF
                - `question`: shown to user
                - `field_type`: `text` / `date` / `number`
                - `repeatable`: `true` for children/beneficiaries etc.
                - `default`: pre-filled value
                """)
                a_cfg    = gr.Code(language="json", label="Config", lines=15)
                a_save   = gr.Button("💾 Save for Users", variant="primary")
                a_status = gr.Textbox(label="Status", interactive=False)

            a_btn.click(admin_upload,  [a_file], [a_prev, a_cfg, a_acc])
            a_save.click(admin_save_config, [a_cfg], [a_status])

        # ── USER ───────────────────────────────────────────────────────────
        with gr.Tab("📝 Fill Form"):
            u_load   = gr.Button("🔄 Load Form", variant="secondary")
            u_status = gr.Textbox(label="Status", interactive=False)
            u_phs    = gr.State([])

            gr.Markdown("### Your Answers")
            gr.Markdown(
                "Edit the JSON below. For **repeatable fields** (children, beneficiaries), "
                "separate multiple entries with `;`\n\n"
                "Example: `\"Child Full Legal Name, DOB\": \"Emma Johnson, 01/01/2010; Liam Johnson, 03/15/2012\"`"
            )
            u_answers = gr.Code(language="json", label="Answers (JSON)", lines=12, value="{}")

            @gr.render(inputs=[u_phs])
            def render_questions(phs):
                if not phs:
                    gr.Markdown("*Click 'Load Form' to see questions.*")
                    return
                for p in phs:
                    hint = " *(separate multiple with `;`)*" if p.get("repeatable") else ""
                    gr.Markdown(f"**{p['question']}**{hint}")

            u_gen  = gr.Button("📄 Generate Filled PDF", variant="primary")

            with gr.Row():
                u_out  = gr.File(label="📥 Download PDF")
                u_prev = gr.HTML(label="Preview")

            def _load():
                status, phs, defaults = load_form()
                return status, phs, defaults

            u_load.click(_load, [], [u_status, u_phs, u_answers])
            u_gen.click(generate_pdf, [u_answers], [u_out, u_prev])

# ── Launch ────────────────────────────────────────────────────────────────────
gr.close_all()
demo.queue(max_size=10).launch(
    server_name="0.0.0.0",
    server_port=int(os.environ.get("GRADIO_SERVER_PORT", 7860)),
    ssr_mode=False,
    show_error=True,
)