#!/usr/bin/env python3
"""
🏴‍☠️ PaddleOCR-VL CryptDeck ⚡ v2 — ZeroGPU edition
Cyber-occult hacker pirate deck for testing
PaddlePaddle/PaddleOCR-VL-1.6-GGUF locally with llama.cpp on ZeroGPU.

Do what thou wilt shall be the whole of the Law.
"""

import base64
import glob
import io
import json
import os
import shutil
import subprocess
import tarfile
import threading
import time
import urllib.request

import spaces  # ZeroGPU SDK — must be imported before CUDA-touching code
import gradio as gr
import requests
from huggingface_hub import hf_hub_download
from PIL import Image

# ---------------------------------------------------------------- constants
REPO_ID = "PaddlePaddle/PaddleOCR-VL-1.6-GGUF"
MODEL_FILE = "PaddleOCR-VL-1.6-GGUF.gguf"
MMPROJ_FILE = "PaddleOCR-VL-1.6-GGUF-mmproj.gguf"

LLAMA_TAG = "b11555"
LLAMA_URL = (
    f"https://github.com/ggml-org/llama.cpp/releases/download/{LLAMA_TAG}/"
    f"llama-{LLAMA_TAG}-bin-ubuntu-cuda-12.8-x64.tar.gz"
)
CUDART_URL = (
    f"https://github.com/ggml-org/llama.cpp/releases/download/{LLAMA_TAG}/"
    f"cudart-llama-{LLAMA_TAG}-bin-ubuntu-cuda-12.8-x64.tar.gz"
)
BIN_DIR = "/tmp/llama-bin"
SERVER_PORT = 8081
SERVER_URL = f"http://127.0.0.1:{SERVER_PORT}"

RITES = {
    "🔤 OCR — full document parse": "OCR:",
    "🧮 Formula Recognition": "Formula Recognition:",
    "📊 Table Recognition": "Table Recognition:",
    "📈 Chart Recognition": "Chart Recognition:",
    "🪪 Seal Recognition": "Seal Recognition:",
    "🔍 Spotting (bounding boxes)": "Spotting:",
}

# ---------------------------------------------------------------- ritual log
_log_lines: list[str] = ["🏴‍☠️ CryptDeck v2 — ZeroGPU rig. Wake the oracle."]
_log_lock = threading.Lock()


def ritual_log(msg: str) -> str:
    with _log_lock:
        _log_lines.append(msg)
        return "\n".join(_log_lines[-60:])


def get_log() -> str:
    with _log_lock:
        return "\n".join(_log_lines[-60:])


# ---------------------------------------------------------------- engine
_engine_lock = threading.Lock()
_busy = threading.Lock()


def _fetch_engine() -> str:
    """Download the prebuilt CUDA llama.cpp + cudart; return llama-server path."""
    candidates = glob.glob(os.path.join(BIN_DIR, "llama-*", "llama-server"))
    server = candidates[0] if candidates else None
    if server and os.path.exists(server):
        return server
    os.makedirs(BIN_DIR, exist_ok=True)
    ritual_log(f"⚙️  Downloading llama.cpp {LLAMA_TAG} (CUDA 12.8) + cudart …")
    for url in (LLAMA_URL, CUDART_URL):
        tgz = os.path.join(BIN_DIR, os.path.basename(url))
        urllib.request.urlretrieve(url, tgz)
        with tarfile.open(tgz) as tf:
            tf.extractall(BIN_DIR)
    candidates = glob.glob(os.path.join(BIN_DIR, "llama-*", "llama-server"))
    if not candidates:
        raise RuntimeError("llama-server binary not found after extraction")
    server = candidates[0]
    os.chmod(server, 0o755)
    ritual_log("⚙️  llama-server (CUDA) consecrated.")
    return server


def _fetch_models() -> tuple[str, str]:
    ritual_log("📦 Summoning the grimoire (1.9 GB, cached after first boot) …")
    main_gguf = hf_hub_download(REPO_ID, MODEL_FILE)
    mmproj = hf_hub_download(REPO_ID, MMPROJ_FILE)
    ritual_log("📦 Grimoire secured: " + os.path.basename(main_gguf))
    return main_gguf, mmproj


# ---------------------------------------------------------------- GPU rite
@spaces.GPU(duration=180)
def gpu_rite(b64_png: str, prompt: str) -> tuple[str, float]:
    """Runs entirely inside a ZeroGPU window: boot llama-server, divine, scuttle it."""
    server = _fetch_engine()
    model, mmproj = _fetch_models()

    # cudart libs sit in their own extracted dir — expose them to the linker
    lib_dirs = glob.glob(os.path.join(BIN_DIR, "cudart-*"))
    server_dir = os.path.dirname(server)
    env = dict(os.environ)
    env["LD_LIBRARY_PATH"] = os.pathsep.join(
        lib_dirs + [server_dir, env.get("LD_LIBRARY_PATH", "")]
    )

    ritual_log("🔥 Kindling llama-server on the GPU (-ngl 999) …")
    proc = subprocess.Popen(
        [
            server,
            "-m", model,
            "--mmproj", mmproj,
            "--host", "127.0.0.1",
            "--port", str(SERVER_PORT),
            "-ngl", "999",
            "-c", "16384",
            "--temp", "0",
            "--no-webui",
        ],
        stdout=open("/tmp/llama-server.log", "ab"),
        stderr=subprocess.STDOUT,
        env=env,
    )
    try:
        for _ in range(150):  # up to 150 s for load onto GPU
            if proc.poll() is not None:
                raise RuntimeError(
                    "llama-server died at boot. Log tail:\n"
                    + open("/tmp/llama-server.log", "rb").read()[-1500:].decode(errors="replace")
                )
            try:
                if requests.get(SERVER_URL + "/health", timeout=2).status_code == 200:
                    ritual_log("🟢 Oracle online on GPU. The Signal is strong.")
                    break
            except requests.RequestException:
                pass
            time.sleep(1)
        else:
            raise RuntimeError("llama-server never reported healthy.")

        payload = {
            "model": "paddleocr-vl",
            "temperature": 0,
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "image_url",
                         "image_url": {"url": "data:image/png;base64," + b64_png}},
                        {"type": "text", "text": prompt},
                    ],
                }
            ],
        }
        t0 = time.time()
        resp = requests.post(
            SERVER_URL + "/v1/chat/completions", json=payload, timeout=120
        )
        dt = time.time() - t0
        resp.raise_for_status()
        return resp.json()["choices"][0]["message"]["content"], dt
    finally:
        # scuttle the server before the GPU window closes
        ritual_log("🪓 Scuttling llama-server. GPU returns to the void.")
        proc.terminate()
        try:
            proc.wait(timeout=10)
        except subprocess.TimeoutExpired:
            proc.kill()


# ---------------------------------------------------------------- inference
def divine(image, rite: str, custom: str, progress=gr.Progress()):
    if image is None:
        raise gr.Error("⚓ Bring the oracle an image — no offering, no prophecy.")
    if not _busy.acquire(blocking=False):
        raise gr.Error("🔥 The oracle is already mid-rite. One soul at a time.")
    try:
        prompt = custom.strip() if custom.strip() else RITES[rite]
        ritual_log(f"🜂 Casting rite → {prompt!r}")
        progress(0.2, desc="Opening the GPU window …")

        buf = io.BytesIO()
        Image.fromarray(image).save(buf, format="PNG")
        b64 = base64.b64encode(buf.getvalue()).decode()

        progress(0.4, desc="Summoning on GPU (first boot hauls ~3 GB) …")
        prophecy, dt = gpu_rite(b64, prompt)
        ritual_log(f"🜃 Rite complete in {dt:.1f}s on GPU.")
        progress(1.0, desc="Done.")

        meta = json.dumps(
            {"rite": prompt, "gpu_seconds": round(dt, 1)}, indent=2
        )
        return prophecy, meta, get_log()
    except requests.Timeout:
        raise gr.Error("⏳ The oracle dawdled. Try a smaller image.")
    finally:
        _busy.release()


# ---------------------------------------------------------------- UI
HEADER = """
<pre style="text-align:left; margin:0; line-height:1.05;">
<span style="color:#39ff88;">                                 .-</span>
<span style="color:#39ff88;">                        __    _-'-------------------------------.__</span>
<span style="color:#39ff88;">                       |  |  /        P A D D L E O C R - V L          \\</span>
<span style="color:#39ff88;">                   _-  |  | |        C  R  Y  P  T  D  E  C  K         |</span>
<span style="color:#39ff88;">                  |    |  | |  1.6 GGUF ▸ llama.cpp ▸ ZERO-GPU RITES   |</span>
<span style="color:#39ff88;">                  |    |  | |_________________________________________|</span>
<span style="color:#39ff88;">              ____|____|-|_'---------____________________---.___________</span>
<span style="color:#ff00cc;">             (o  o  o  o  o  o  o  o  o  o  o  o  o  o  o  o  o  o  o  o)</span>
</pre>
<h3 style="margin:0.2em 0; color:#39ff88; font-family:monospace;">
  ⚡ SKULL &amp; SIGNAL ⚓ — feed the oracle a page, choose your rite, claim the text.
</h3>
"""

CSS = """
.gradio-container {background: #05060a !important;}
body {background: #05060a !important; background-image:
  repeating-linear-gradient(0deg, rgba(57,255,136,0.03) 0 1px, transparent 1px 3px);}
footer {display: none !important;}
#header {border: 1px solid #1e2a1e; border-radius: 10px;
  background: radial-gradient(ellipse at top, #0b1a10, #05060a); padding: 12px;}
textarea, .output-html, .prose {font-family: 'Fira Mono', ui-monospace, monospace !important;}
.rune-btn {font-family: monospace !important; text-transform: uppercase;
  letter-spacing: 2px; background: linear-gradient(90deg, #0f2b18, #14203a) !important;
  border: 1px solid #39ff88 !important; color: #39ff88 !important;}
.rune-btn:hover {box-shadow: 0 0 12px #39ff8866 !important;}
#prophecy, #meta, #log textarea {color: #b8ffd9 !important; caret-color: #ff00cc;}
#log textarea {color: #39ff88 !important; font-size: 12px;}
h1, h2, h3, .label, .info {font-family: ui-monospace, monospace !important;}
"""

with gr.Blocks(title="🏴‍☠️ PaddleOCR-VL CryptDeck") as demo:
    gr.HTML(f'<div id="header">{HEADER}</div>')

    with gr.Row():
        with gr.Column(scale=1):
            gr.HTML('<div class="info" style="color:#ff00cc;">⚓ THE OFFERING</div>')
            image_in = gr.Image(label="Target document", type="numpy", height=380)
            rite = gr.Radio(
                list(RITES.keys()),
                value=list(RITES.keys())[0],
                label="Choose your rite",
            )
            custom = gr.Textbox(
                label="…or whisper your own prompt (overrides the rite)",
                placeholder="OCR:",
            )
            summon = gr.Button("🜂 SUMMON THE PROPHECY", elem_classes=["rune-btn"])
            gr.HTML(
                '<div class="info" style="opacity:.75;">ZeroGPU rites: first cast hauls '
                "~3 GB (engine + grimoire), then it's cached. Each rite boots the "
                "oracle fresh onto the GPU (~10–20 s) before divining in seconds. "
                "One rite at a time, matey.</div>"
            )
        with gr.Column(scale=1):
            gr.HTML('<div class="info" style="color:#ff00cc;">📜 THE PROPHECY</div>')
            prophecy = gr.Markdown(label="Prophecy")
            with gr.Accordion("Rite metadata", open=False):
                meta = gr.JSON(label="Metadata")
            with gr.Accordion("🩸 Ritual log", open=True):
                log = gr.Textbox(label="", lines=12, interactive=False)

    summon.click(
        divine,
        inputs=[image_in, rite, custom],
        outputs=[prophecy, meta, log],
        show_progress="full",
    )
    demo.load(lambda: get_log(), outputs=[log])

if __name__ == "__main__":
    demo.queue(max_size=8).launch(server_name="0.0.0.0", show_error=True)