"""PaddleOCR-VL 开发场景代码 OCR — Hugging Face Space (Gradio)

模型：https://huggingface.co/snnh/paddleocr_vl_code_ocr （PaddleOCR-VL-1.6 微调）
提示词：<image>OCR:（与训练一致）

部署说明：本 Space 默认运行在 CPU 硬件上（免费档），0.9B 模型单图推理约 10–60s。
GPU 环境（ZeroGPU / A10G）下 @spaces.GPU 装饰器自动接管，单图数秒。
"""

import os
import re
import inspect
import sys
import threading
import time
import traceback
from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError

import gradio as gr
import torch
# 用 AutoModelForCausalLM 而非 AutoModelForImageTextToText：
# 模型 config.json 的 auto_map 只映射了 AutoModel / AutoModelForCausalLM，
# 没有映射 AutoModelForImageTextToText。用 CausalLM 才能走 trust_remote_code
# 加载模型仓库里的自定义 modeling_paddleocr_vl.PaddleOCRVLForConditionalGeneration
# （该类继承 GenerationMixin，支持 .generate()）。
from transformers import AutoModelForCausalLM, AutoProcessor


_CAUSAL_MASK_COMPAT = None


def _patch_transformers_causal_mask():
    """Map PaddleOCR-VL's create_causal_mask kwargs across transformers builds."""
    global _CAUSAL_MASK_COMPAT
    try:
        import transformers.masking_utils as masking_utils
    except Exception as exc:  # noqa: BLE001
        print(f"[init] causal mask patch skipped: {exc}")
        return None

    original = getattr(masking_utils, "create_causal_mask", None)
    if original is None:
        return None
    if getattr(original, "_ppocr_inputs_embeds_patch", False):
        _CAUSAL_MASK_COMPAT = original
        return original

    signature = inspect.signature(original)
    params = signature.parameters
    accepts_var_kwargs = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values())

    def compat_create_causal_mask(*args, **kwargs):
        if "inputs_embeds" in kwargs and "inputs_embeds" not in params:
            inputs_embeds = kwargs.pop("inputs_embeds")
            if "input_embeds" in params:
                kwargs["input_embeds"] = inputs_embeds
            elif "input_tensor" in params:
                kwargs["input_tensor"] = inputs_embeds
            else:
                # Very old/new signatures may not need the tensor explicitly.
                pass
        if not accepts_var_kwargs:
            kwargs = {k: v for k, v in kwargs.items() if k in params}
        return original(*args, **kwargs)

    compat_create_causal_mask._ppocr_inputs_embeds_patch = True
    masking_utils.create_causal_mask = compat_create_causal_mask
    _CAUSAL_MASK_COMPAT = compat_create_causal_mask
    print("[init] patched transformers create_causal_mask compatibility")
    return compat_create_causal_mask


def _patch_loaded_paddleocr_modules():
    """Remote modeling module imports create_causal_mask by value; update it too."""
    if _CAUSAL_MASK_COMPAT is None:
        return
    for module_name, module in list(sys.modules.items()):
        if module_name.endswith("modeling_paddleocr_vl") and hasattr(module, "create_causal_mask"):
            setattr(module, "create_causal_mask", _CAUSAL_MASK_COMPAT)


_patch_transformers_causal_mask()


def _patch_gradio_schema_bool_handling():
    """Work around gradio_client 5.9.x API-info crashes on bool schemas.

    Gradio 5.9.1 can emit JSON schema fragments such as
    {"additionalProperties": true} for Image/File metadata. Its matching
    gradio_client json_schema_to_python_type path assumes every nested schema
    is a dict and crashes in /gradio_api/info with:
    TypeError: argument of type 'bool' is not iterable.
    That failed /info request makes the frontend report "No API found".
    """
    try:
        import gradio_client.utils as client_utils
    except Exception as exc:  # noqa: BLE001
        print(f"[init] gradio schema patch skipped: {exc}")
        return

    original_get_type = getattr(client_utils, "get_type", None)
    if original_get_type is None or getattr(original_get_type, "_ppocr_bool_patch", False):
        return

    def safe_get_type(schema):
        if isinstance(schema, bool):
            return "Any"
        return original_get_type(schema)

    safe_get_type._ppocr_bool_patch = True
    client_utils.get_type = safe_get_type
    print("[init] patched gradio_client bool schema handling")


_patch_gradio_schema_bool_handling()


def _env_truthy(name: str) -> bool:
    return os.environ.get(name, "").strip().lower() in {"1", "true", "yes", "on"}


# ZeroGPU Space 才需要 @spaces.GPU。CPU Basic / paid GPU Space 上强行套这个
# wrapper 可能在 Gradio 事件进入 run_ocr 前抛错，只给前端一个红色“错误”。
IS_ZERO_GPU = _env_truthy("SPACES_ZERO_GPU") or _env_truthy("SPACE_ZERO_GPU")

if IS_ZERO_GPU:
    try:
        import spaces  # type: ignore
        _GPU_DECORATOR = spaces.GPU
        print("[init] ZeroGPU decorator enabled")
    except Exception as exc:  # noqa: BLE001
        print(f"[init] ZeroGPU decorator unavailable: {exc}")

        def _GPU_DECORATOR(fn):
            return fn
else:
    def _GPU_DECORATOR(fn):
        return fn

MODEL_ID = os.environ.get("MODEL_ID", "snnh/paddleocr_vl_code_ocr")
# 官方 ocr 任务的文本部分；<image> 占位符由 processor 的 chat_template 自动插入，
# 训练数据 user.content 为 "<image>OCR:"，经模板 tokenize 后与下方结构化构造等价。
OCR_PROMPT = "OCR:"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
IS_CPU = DEVICE == "cpu"

# 评估主口径是 4096；CPU demo 降低 token 上限，并设置墙钟超时，避免评审端长时间无结果。
MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "1536" if IS_CPU else "4096"))
REPETITION_PENALTY = float(os.environ.get("REPETITION_PENALTY", "1.08"))
RUN_TIMEOUT_SECONDS = float(os.environ.get("RUN_TIMEOUT_SECONDS", "300" if IS_CPU else "120"))
_DEFAULT_MAX_TIME = max(30.0, RUN_TIMEOUT_SECONDS - 45.0) if IS_CPU else 90.0
MAX_TIME_SECONDS = float(os.environ.get("MAX_TIME_SECONDS", str(_DEFAULT_MAX_TIME)))
PROGRESS_POLL_SECONDS = float(os.environ.get("PROGRESS_POLL_SECONDS", "2"))
# 官方推荐：ocr/table/chart/formula/seal 用 ~1M pixels，spotting 用更大尺寸
MAX_PIXELS = 1280 * 28 * 28

# ---------------------------------------------------------------------------
# 后处理：与训练/评估 notebook (Cell 11) 完全一致的 clean_prediction。
# 只处理"失控重复"与代码围栏/前缀噪声，不改写正常 OCR 内容。
# 来源：notebooks/aistudio_paddleocr_vl16_v6_stepmatched_finetune.ipynb
# ---------------------------------------------------------------------------


def collapse_runaway_duplicate_lines(text: str, max_keep: int = 2) -> str:
    lines = text.split('\n')
    kept = []
    last = None
    run = 0
    for line in lines:
        key = line.strip()
        if key and key == last:
            run += 1
        else:
            last = key
            run = 1
        if run <= max_keep:
            kept.append(line)
    return '\n'.join(kept)


def trim_repeated_tail(text: str, min_unit: int = 8, max_unit: int = 80, min_repeats: int = 4) -> str:
    text = text.rstrip()
    if len(text) < min_unit * min_repeats:
        return text
    upper = min(max_unit, len(text) // min_repeats)
    for unit in range(upper, min_unit - 1, -1):
        chunk = text[-unit:]
        if not chunk.strip():
            continue
        repeated = chunk * min_repeats
        if text.endswith(repeated):
            while text.endswith(chunk * 2):
                text = text[:-unit].rstrip()
            return text
    return text


def char_ngram_repeat_ratio(text: str, n: int = 8) -> float:
    compact = re.sub(r'\s+', '', text)
    if len(compact) < n * 3:
        return 0.0
    grams = [compact[i:i + n] for i in range(len(compact) - n + 1)]
    return 1.0 - len(set(grams)) / max(1, len(grams))


def clean_prediction(text):
    text = text.strip()
    text = re.sub(r'^```(?:text|txt|[a-zA-Z0-9_+-]+)?\s*', '', text)
    text = re.sub(r'\s*```$', '', text)
    text = re.sub(r'^(?:Assistant|助手)\s*[:：]\s*', '', text, flags=re.IGNORECASE)
    text = collapse_runaway_duplicate_lines(text, max_keep=2)
    text = trim_repeated_tail(text)
    return text.strip()


DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32

_model_lock = threading.Lock()
_processor = None
_model = None
_inference_executor = ThreadPoolExecutor(max_workers=1)
_active_lock = threading.Lock()
_active_future = None


def get_model():
    """Lazy-load the model so the Space can start before downloading weights."""
    global _processor, _model
    if _processor is not None and _model is not None:
        return _processor, _model

    with _model_lock:
        if _processor is not None and _model is not None:
            return _processor, _model

        print(f"[model] loading {MODEL_ID} on {DEVICE}")
        processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            trust_remote_code=True,
            torch_dtype=DTYPE,
            _attn_implementation="eager",
        ).to(DEVICE).eval()
        _patch_loaded_paddleocr_modules()
        _processor = processor
        _model = model
        print(f"[model] loaded {MODEL_ID} on {DEVICE}")
        return _processor, _model


def _run_ocr_impl(image):
    t0 = time.time()
    try:
        processor, model = get_model()

        # image 由 Gradio 给出，可能是 PIL.Image 或 numpy.ndarray
        if not hasattr(image, "convert"):
            import numpy as np
            from PIL import Image as PILImage
            image = PILImage.fromarray(np.array(image)).convert("RGB")
        else:
            image = image.convert("RGB")

        # 官方标准构造：<image> 占位符 + OCR: 文本，经 apply_chat_template 生成训练分布一致的输入
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image},
                    {"type": "text", "text": OCR_PROMPT},
                ],
            }
        ]
        inputs = processor.apply_chat_template(
            messages,
            add_generation_prompt=True,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
            images_kwargs={
                "size": {
                    "shortest_edge": processor.image_processor.min_pixels,
                    "longest_edge": MAX_PIXELS,
                }
            },
        ).to(DEVICE)

        with torch.inference_mode():
            generation_kwargs = {
                **inputs,
                "max_new_tokens": MAX_NEW_TOKENS,
                "do_sample": False,
                "repetition_penalty": REPETITION_PENALTY,
            }
            if MAX_TIME_SECONDS > 0:
                generation_kwargs["max_time"] = MAX_TIME_SECONDS
            output_ids = model.generate(
                **generation_kwargs,
            )

        # 只取新生成部分，去掉 prompt
        prompt_len = inputs["input_ids"].shape[-1]
        generated = output_ids[0][prompt_len:-1]
        raw_text = processor.decode(generated, skip_special_tokens=True).strip()
        # 与评估口径一致的后处理：剥离代码围栏/前缀，压塌失控重复
        text = clean_prediction(raw_text)
        elapsed = time.time() - t0
        device_note = "CPU（较慢）" if IS_CPU else DEVICE.upper()
        return f"{text}\n\n---\n耗时 {elapsed:.1f}s · {device_note}"
    except Exception as e:  # noqa: BLE001
        traceback.print_exc()
        return f"OCR 失败：{e}"


def _clear_active_future(future):
    global _active_future
    with _active_lock:
        if _active_future is future:
            _active_future = None


def _submit_inference(image):
    global _active_future
    with _active_lock:
        if _active_future is not None and not _active_future.done():
            return None
        _active_future = _inference_executor.submit(_run_ocr_impl, image)
        _active_future.add_done_callback(_clear_active_future)
        return _active_future


@_GPU_DECORATOR
def run_ocr(image, progress=gr.Progress(track_tqdm=False)):
    """对上传图片做开发场景 OCR。只转写可见文字。"""
    if image is None:
        return "请先上传一张图片。"

    future = _submit_inference(image)
    if future is None:
        return (
            "上一轮 OCR 仍在运行，请稍后再试。\n\n"
            "---\n"
            "当前 Demo 已限制为单任务队列，避免 CPU 免费硬件被重复点击拖到长时间无响应。"
        )

    started_at = time.time()
    progress(0.03, desc="已进入队列，准备加载模型")
    while True:
        elapsed = time.time() - started_at
        remaining = RUN_TIMEOUT_SECONDS - elapsed
        if remaining <= 0:
            return (
                "OCR 超时，已停止等待本次结果。\n\n"
                "---\n"
                f"等待 {RUN_TIMEOUT_SECONDS:.0f}s 未完成。CPU 免费硬件可能无法稳定满足评审体验，"
                "建议切换 ZeroGPU / A10G 或本地 GPU 部署后重试。"
            )

        try:
            return future.result(timeout=min(PROGRESS_POLL_SECONDS, remaining))
        except FuturesTimeoutError:
            ratio = min(0.95, 0.05 + 0.90 * elapsed / max(1.0, RUN_TIMEOUT_SECONDS))
            desc = "首次加载模型权重" if _model is None else "正在识别，请勿重复点击"
            progress(ratio, desc=desc)


EXAMPLES_NOTICE = (
    "上传 IDE 截图、终端、Traceback、配置文件、Git diff、API 表格或困难样本。"
    "输出只含可见文字，不解释、不补全。"
)

CPU_NOTICE = (
    "⚠️ 当前为 **CPU 免费硬件**，已启用单任务队列和超时保护。首次点击需要加载模型权重，"
    f"单次最多等待约 {RUN_TIMEOUT_SECONDS:.0f} 秒；识别期间请勿重复点击。如需更快体验，可在 ZeroGPU/A10G 或本地 GPU 部署"
    "（见 GitHub 仓库 `demo/openai_compatible_ocr_demo.py`）。"
    if IS_CPU else ""
)

with gr.Blocks(title="PaddleOCR-VL 开发场景代码 OCR") as demo:
    gr.Markdown(
        "# PaddleOCR-VL 开发场景代码 OCR\n"
        "PaddleOCR 全球衍生模型挑战赛提交 Demo · 基于 PaddleOCR-VL-1.6 微调\n\n"
        f"{EXAMPLES_NOTICE}"
    )
    if CPU_NOTICE:
        gr.Markdown(CPU_NOTICE)
    with gr.Row():
        with gr.Column():
            img_in = gr.Image(label="开发场景截图", type="pil", height=420)
            btn = gr.Button("开始 OCR", variant="primary")
        with gr.Column():
            txt_out = gr.Textbox(
                label="识别结果",
                lines=22,
                show_copy_button=True,
                placeholder="识别结果会显示在这里...（CPU 模式已启用队列和超时保护）",
            )
    gr.Markdown(
        f"模型：`{MODEL_ID}` · 提示词：`<image>OCR:` · "
        f"`max_new_tokens={MAX_NEW_TOKENS}` · `repetition_penalty={REPETITION_PENALTY}` · "
        f"`max_time={MAX_TIME_SECONDS:.0f}s` · `run_timeout={RUN_TIMEOUT_SECONDS:.0f}s` · "
        f"后处理：`clean_prediction`（与评估口径一致，压塌失控重复）"
    )
    btn.click(run_ocr, inputs=img_in, outputs=txt_out, api_name="ocr", queue=True, show_api=False)
    legacy_api_btn = gr.Button("legacy API", visible=False)
    legacy_api_btn.click(
        run_ocr,
        inputs=img_in,
        outputs=txt_out,
        api_name="run_ocr",
        queue=True,
        show_api=False,
    )
    demo.queue(max_size=4, default_concurrency_limit=1)

if __name__ == "__main__":
    demo.launch()