"""
UltraSharp V2 — 图像超分辨率 Gradio 应用
==========================================

## 模型来源

   默认从 Kim2091/UltraSharpV2 公开仓库下载 4x-UltraSharpV2.pth，
   自动缓存到 ~/.cache/huggingface/hub/，无需手动上传。

## 可选环境变量

   MODEL_REPO_ID    覆盖默认仓库（默认 Kim2091/UltraSharpV2）
   MODEL_FILENAME   覆盖默认文件名（默认 4x-UltraSharpV2.pth）
   HF_ENDPOINT      镜像站，如 https://hf-mirror.com（国内加速）
   HF_TOKEN         私有仓库的 token（公开仓库无需设置）

## 本地运行

   python app.py
   # 国内镜像: HF_ENDPOINT=https://hf-mirror.com python app.py

## 部署到 HuggingFace Space

   1. 在 Space 设置中将 Hardware 选为 ZeroGPU
   2. 无需设置 Secrets（模型来自公开仓库）
   3. 如需国内镜像，添加 Secret: HF_ENDPOINT = https://hf-mirror.com
"""

import os
import asyncio
import asyncio.base_events
import gradio as gr
from model_loader import UltraSharpV2

# ---------------------------------------------------------------------------
# 修复 Python 3.12 asyncio 事件循环 GC 时的 "Invalid file descriptor: -1" 报错
#
# 根因: Gradio / spaces 在 import 阶段会创建临时事件循环，这些循环被 GC
# 回收时 __del__ → close() → _close_self_pipe() 尝试对已关闭的 socket
# (fd=-1) 执行 _remove_reader，触发 ValueError。属于 CPython 3.12 的
# 已知问题，对功能无害但日志很吵。此处 patch __del__ 静默吞掉该异常。
# ---------------------------------------------------------------------------
_orig_loop_del = asyncio.base_events.BaseEventLoop.__del__


def _safe_loop_del(self):
    try:
        _orig_loop_del(self)
    except Exception:
        pass


asyncio.base_events.BaseEventLoop.__del__ = _safe_loop_del

# ---------------------------------------------------------------------------
# ZeroGPU 兼容层
# ---------------------------------------------------------------------------
try:
    import spaces

    _zerogpu = spaces.GPU(duration=120)  # 最长 GPU 占用 120s
    IN_ZEROGPU = bool(os.environ.get("SPACES_ZERO_GPU"))
except ImportError:
    spaces = None
    _zerogpu = None
    IN_ZEROGPU = False


def _gpu(fn):
    """安全地应用 @spaces.GPU 装饰器（本地开发时退化为无操作）。"""
    return _zerogpu(fn) if _zerogpu is not None else fn


# ---------------------------------------------------------------------------
# 模型：始终在 CPU 上加载（ZeroGPU 启动时 GPU 不可用）
# ---------------------------------------------------------------------------
model = UltraSharpV2(device="cpu")


# ---------------------------------------------------------------------------
# 推理参数（RTX PRO 6000 Blackwell / 48GB — 无需省显存）
# ---------------------------------------------------------------------------
_TILE_SIZE = 1024
_TILE_OVERLAP = 48


# ---------------------------------------------------------------------------
# 推理函数（生成器模式 — ZeroGPU 硬性要求）
# ---------------------------------------------------------------------------
@_gpu
def on_upscale(image, target_scale):
    if image is None:
        yield None, "请先上传图片"
        return

    model.to_cuda()
    try:
        result, elapsed = model.upscale(
            image, _TILE_SIZE, _TILE_OVERLAP, float(target_scale)
        )
    finally:
        model.to_cpu()

    yield result, f"耗时: {elapsed:.2f}s"


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
with gr.Blocks(title="UltraSharp V2") as demo:
    device_display = "ZeroGPU" if IN_ZEROGPU else model.device.upper()
    gr.Markdown("# UltraSharp V2 - 图像超分辨率")
    gr.Markdown(f"**运行设备**: {device_display} | **模型原生倍率**: {model.scale}x")

    with gr.Row():
        with gr.Column(scale=1):
            input_img = gr.Image(label="输入图片", type="pil", height=400)
            target_scale = gr.Slider(
                label="放大倍率",
                minimum=1.0,
                maximum=4.0,
                value=4.0,
                step=0.05,
                info="> 模型原生倍率时, 输出先 4x 推理再 Lanczos 缩放",
            )

        with gr.Column(scale=1):
            run_btn = gr.Button("开始推理", variant="primary")
            output_img = gr.Image(label="推理结果", height=400)
            status = gr.Textbox(label="状态", interactive=False)

    run_btn.click(
        fn=on_upscale,
        inputs=[input_img, target_scale],
        outputs=[output_img, status],
    )


# ---------------------------------------------------------------------------
# ZeroGPU 必须启用 queue（默认并发 1，队列上限 10）
# ---------------------------------------------------------------------------
demo.queue(max_size=10, default_concurrency_limit=1)

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