"""Parallel-only streaming PaDoc demo for Hugging Face ZeroGPU."""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # Must precede torch and every module that imports torch.

import json
import re
import time
from typing import Any

import gradio as gr
import torch
from PIL import Image, ImageDraw, ImageFont

from padoc.modeling import load_padoc_model
from padoc.transformers_infer import SequentialPaDocEngine

MODEL_ID = os.environ.get("MODEL_ID", "Longin-Yu/PaDoc")
DEFAULT_QUERY = "Parse this document."
MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "512"))
MAX_BRANCH_TOKENS = int(os.environ.get("MAX_BRANCH_TOKENS", "512"))
MAX_CONCURRENT_BRANCHES = int(os.environ.get("MAX_CONCURRENT_BRANCHES", "8"))
MAX_TOTAL_BRANCHES = int(os.environ.get("MAX_TOTAL_BRANCHES", "64"))
ZERO_GPU_ENABLED = os.environ.get("SPACES_ZERO_GPU") == "1"

# The CPU staging mode keeps the app RUNNING while a new account's ZeroGPU grant
# is pending. On ZeroGPU, weights are loaded and packed at module scope.
if ZERO_GPU_ENABLED:
    model, processor, fork_map = load_padoc_model(
        MODEL_ID,
        dtype=torch.bfloat16,
        device_map=None,
        attn_implementation="sdpa",
    )
    model = model.to("cuda").eval()
    print(f"[PaDoc] Ready: model={MODEL_ID}, device={model.device}, mode=parallel")
else:
    model = None
    processor = None
    fork_map = None
    print("[PaDoc] CPU staging mode: waiting for ZeroGPU hardware.")


_LAYOUT_RE = re.compile(r"<SP_LAYOUT>(\d+)\s+(\d+)\s+(\d+)\s+(\d+)</SP_LAYOUT>")
_META_RE = re.compile(r"<SP_META>(\{.*?\})</SP_META>")
_COLORS = (
    "#d94f4f",
    "#267a63",
    "#3468a5",
    "#9b5c18",
    "#7654a8",
    "#16808c",
    "#b13d79",
    "#65751f",
)


def parse_layout_boxes(main_text: str) -> list[tuple[int, int, int, int]]:
    """Extract complete layout boxes in normalized [0, 1000] coordinates."""

    boxes = []
    for match in _LAYOUT_RE.finditer(main_text):
        box = tuple(int(value) for value in match.groups())
        if all(0 <= value <= 1000 for value in box) and box[0] < box[2] and box[1] < box[3]:
            boxes.append(box)
    return boxes


def parse_branch_text(branch_text: str) -> tuple[str, str]:
    """Extract the category and visible content from one branch."""

    match = _META_RE.search(branch_text)
    category = "region"
    if match:
        try:
            metadata = json.loads(match.group(1))
            if isinstance(metadata.get("category"), str):
                category = metadata["category"]
        except json.JSONDecodeError:
            pass
    content = _META_RE.sub("", branch_text).strip()
    return category, content


def annotate_image(
    image: Image.Image,
    boxes: list[tuple[int, int, int, int]],
) -> Image.Image:
    """Draw numbered normalized boxes on a copy of the source image."""

    annotated = image.copy().convert("RGB")
    width, height = annotated.size
    draw = ImageDraw.Draw(annotated)
    try:
        font = ImageFont.truetype(
            "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
            max(14, int(min(width, height) / 45)),
        )
    except OSError:
        font = ImageFont.load_default()

    stroke = max(2, int(min(width, height) / 350))
    for index, (x1, y1, x2, y2) in enumerate(boxes):
        color = _COLORS[index % len(_COLORS)]
        pixel_box = (
            int(x1 / 1000 * width),
            int(y1 / 1000 * height),
            int(x2 / 1000 * width),
            int(y2 / 1000 * height),
        )
        draw.rectangle(pixel_box, outline=color, width=stroke)
        label = str(index + 1)
        label_box = draw.textbbox((0, 0), label, font=font)
        label_width = label_box[2] - label_box[0]
        label_height = label_box[3] - label_box[1]
        label_x = pixel_box[0]
        label_y = max(0, pixel_box[1] - label_height - 8)
        draw.rectangle(
            (label_x, label_y, label_x + label_width + 10, label_y + label_height + 8),
            fill=color,
        )
        draw.text((label_x + 5, label_y + 3), label, fill="white", font=font)
    return annotated


def format_regions(
    branches: dict[int, dict[str, Any]],
    boxes: list[tuple[int, int, int, int]],
) -> str:
    """Render current branch streams as stable region sections."""

    if not branches:
        return "_Waiting for forked content branches..._"

    sections = []
    for branch_index in sorted(branches):
        branch = branches[branch_index]
        category, content = parse_branch_text(branch.get("text", ""))
        state = branch.get("state", "queued")
        box_text = ""
        if branch_index < len(boxes):
            box_text = " `[{0}, {1}, {2}, {3}]`".format(*boxes[branch_index])
        sections.append(f"### {branch_index + 1:02d} | {category}{box_text}")
        sections.append(content or f"_{state}..._")
    return "\n\n".join(sections)


def format_status(
    scheduler: dict[str, Any],
    *,
    elapsed: float,
    main_tokens: int,
    branch_count: int,
    done: bool,
) -> str:
    """Format the live parallel scheduler state."""

    phase = "complete" if done else scheduler.get("phase", "starting")
    return (
        f"**Parallel** | {phase} | main {main_tokens} tok | "
        f"{scheduler.get('active_branches', 0)} active | "
        f"{scheduler.get('queued_branches', 0)} queued | "
        f"batch {scheduler.get('batch_size', 0)} | "
        f"{branch_count} branches | {elapsed:.1f}s"
    )


def snapshot(
    *,
    main_text: str,
    branches: dict[int, dict[str, Any]],
    scheduler: dict[str, Any],
    last_event: dict[str, Any],
) -> dict[str, Any]:
    """Build a JSON-safe live result snapshot."""

    return {
        "execution_mode": "parallel",
        "main": main_text,
        "branches": [branches[index] for index in sorted(branches)],
        "scheduler": scheduler,
        "last_event": last_event,
    }


@spaces.GPU(duration=120, size="large")
def parse_document(
    image: Image.Image | None,
    query: str = DEFAULT_QUERY,
):
    """Stream parallel PaDoc parsing for one document image.

    Args:
        image: Document page to parse.
        query: Instruction sent to the document parser.

    Yields:
        Annotated page, scheduler status, main stream, branch streams, and live JSON.
    """

    if image is None:
        raise gr.Error("Select a document image first.")
    if not query or not query.strip():
        raise gr.Error("Query cannot be empty.")
    if model is None or processor is None or fork_map is None:
        raise gr.Error("This Space is waiting for ZeroGPU access.")
    image = image.convert("RGB")

    request_engine = SequentialPaDocEngine(
        model,
        processor,
        fork_map,
        max_new_tokens=MAX_NEW_TOKENS,
        max_branch_tokens=MAX_BRANCH_TOKENS,
        max_concurrent_branches=MAX_CONCURRENT_BRANCHES,
        max_total_branches=MAX_TOTAL_BRANCHES,
        execution_mode="parallel",
        strict=True,
    )
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": query.strip()},
            ],
        }
    ]

    started_at = time.perf_counter()
    main_text = ""
    main_tokens = 0
    branches: dict[int, dict[str, Any]] = {}
    scheduler: dict[str, Any] = {
        "phase": "starting",
        "active_branches": 0,
        "queued_branches": 0,
        "completed_branches": 0,
        "batch_size": 0,
    }
    last_box_count = -1
    first_update = True

    for source_event in request_engine.stream(messages, execution_mode="parallel"):
        event = dict(source_event)
        event_type = event.get("type")

        if event_type == "main":
            main_text += event.get("delta_text", "")
            main_tokens = int(event.get("total", main_tokens))
        elif event_type == "fork":
            index = int(event["branch_index"])
            branches[index] = {
                "branch_index": index,
                "fork_position": event.get("fork_position"),
                "text": event.get("injected_text", ""),
                "state": event.get("branch_state", "queued"),
            }
        elif event_type == "branch":
            index = int(event["branch_index"])
            branch = branches.setdefault(
                index,
                {
                    "branch_index": index,
                    "fork_position": event.get("fork_position"),
                    "text": "",
                    "state": "active",
                },
            )
            branch["text"] += event.get("delta_text", "")
            branch["state"] = "active"
            branch["tokens"] = event.get("total")
        elif event_type == "branch_done":
            index = int(event["branch_index"])
            if index in branches:
                branches[index]["state"] = "done"
                branches[index]["tokens"] = event.get("total")
        elif event_type == "scheduler":
            scheduler = event
        elif event_type == "done":
            main_text = event.get("main", main_text)
            main_tokens = len(event.get("main_token_ids", ()))
            for result_branch in event.get("branches", ()):
                index = int(result_branch["branch_index"])
                branches[index] = {
                    "branch_index": index,
                    "fork_position": result_branch.get("fork_position"),
                    "text": result_branch.get("text", ""),
                    "tokens": len(result_branch.get("token_ids", ())),
                    "state": "done",
                }

        boxes = parse_layout_boxes(main_text)
        if first_update or len(boxes) != last_box_count:
            image_update: Any = annotate_image(image, boxes)
            last_box_count = len(boxes)
            first_update = False
        else:
            image_update = gr.skip()

        elapsed = time.perf_counter() - started_at
        done = event_type == "done"
        status = format_status(
            scheduler,
            elapsed=elapsed,
            main_tokens=main_tokens,
            branch_count=len(branches),
            done=done,
        )
        live_json = (
            event
            if done
            else snapshot(
                main_text=main_text,
                branches=branches,
                scheduler=scheduler,
                last_event=event,
            )
        )
        yield (
            image_update,
            status,
            main_text,
            format_regions(branches, boxes),
            live_json,
        )


CSS = """
#app-shell { max-width: 1240px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
#stream-status { min-height: 30px; }
#main-stream textarea { font-family: ui-monospace, SFMono-Regular, Menlo, monospace; }
.gradio-container { letter-spacing: 0; }
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="PaDoc") as demo:
    with gr.Column(elem_id="app-shell"):
        gr.Markdown(
            "# PaDoc\n"
            "[Model](https://huggingface.co/Longin-Yu/PaDoc) | "
            "[Code](https://github.com/Longin-Yu/Padoc) | "
            "[Paper](https://arxiv.org/abs/2608.06146)"
        )
        with gr.Row(equal_height=False):
            with gr.Column(scale=5, min_width=320):
                image_input = gr.Image(
                    label="Document",
                    type="pil",
                    sources=["upload", "clipboard"],
                    height=470,
                )
                query_input = gr.Textbox(
                    label="Query",
                    value=DEFAULT_QUERY,
                    lines=2,
                )
                run_button = gr.Button("Parse document", variant="primary")
            with gr.Column(scale=7, min_width=360):
                annotated_output = gr.Image(
                    label="Detected regions",
                    interactive=False,
                    height=470,
                )
                status_output = gr.Markdown(
                    (
                        "**Parallel** | ready"
                        if ZERO_GPU_ENABLED
                        else "**Parallel** | waiting for ZeroGPU access"
                    ),
                    elem_id="stream-status",
                )

        with gr.Tabs():
            with gr.Tab("Regions"):
                regions_output = gr.Markdown("_Waiting for a document..._")
            with gr.Tab("Main stream"):
                main_output = gr.Textbox(
                    label="Main sequence",
                    lines=12,
                    interactive=False,
                    show_copy_button=True,
                    elem_id="main-stream",
                )
            with gr.Tab("JSON"):
                json_output = gr.JSON(label="Live result")

        gr.Examples(
            examples=[
                ["examples/sample_memo.png", DEFAULT_QUERY],
                ["examples/sample_invoice.png", DEFAULT_QUERY],
            ],
            inputs=[image_input, query_input],
            outputs=[
                annotated_output,
                status_output,
                main_output,
                regions_output,
                json_output,
            ],
            fn=parse_document,
            cache_examples=True,
            cache_mode="lazy",
        )

    run_button.click(
        fn=parse_document,
        inputs=[image_input, query_input],
        outputs=[
            annotated_output,
            status_output,
            main_output,
            regions_output,
            json_output,
        ],
        api_name="parse",
        concurrency_limit=1,
        concurrency_id="padoc-gpu",
        show_progress="minimal",
    )

demo.queue(default_concurrency_limit=1, max_size=20)
demo.launch(mcp_server=True)