import gradio as gr
import huggingface_hub
import onnxruntime as rt
import numpy as np
import cv2
import os
import time
import threading
import hashlib
import tempfile
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed


# ---------- CPU 优化推理配置 ----------
sess_options = rt.SessionOptions()
sess_options.intra_op_num_threads = 2       # 匹配 2 核 CPU
sess_options.inter_op_num_threads = 2
sess_options.graph_optimization_level = rt.GraphOptimizationLevel.ORT_ENABLE_ALL
sess_options.enable_cpu_mem_arena = True

from onnxruntime.quantization import quantize_static, QuantType, CalibrationMethod
from onnxruntime.quantization import CalibrationDataReader as _CalibReaderBase

providers = ["CPUExecutionProvider"]
model_path = huggingface_hub.hf_hub_download("skytnt/anime-seg", "isnetis.onnx")
model_dir = os.path.dirname(model_path)
quantized_path = os.path.join(model_dir, "isnetis_int8.onnx")


class _CalibReader(_CalibReaderBase):
    """用项目自带的示例图片做静态量化校准，生成 QLinearConv 算子"""
    def __init__(self):
        self._data = []
        for i in range(1, 4):
            path = os.path.join("examples", f"{i:02d}.jpg")
            if not os.path.exists(path):
                continue
            img_bgr = cv2.imread(path)
            if img_bgr is None:
                continue
            img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
            img = (img_rgb / 255).astype(np.float32)
            h, w = img.shape[:2]
            s = 1024
            h2, w2 = (s, int(s * w / h)) if h > w else (int(s * h / w), s)
            ph, pw = s - h2, s - w2
            inp = np.zeros((s, s, 3), dtype=np.float32)
            inp[ph // 2:ph // 2 + h2, pw // 2:pw // 2 + w2] = cv2.resize(img, (w2, h2))
            inp = np.transpose(inp, (2, 0, 1))[np.newaxis, :]
            self._data.append({"img": inp})
        self._idx = 0

    def get_next(self):
        if self._idx < len(self._data):
            val = self._data[self._idx]
            self._idx += 1
            return val
        return None


if not os.path.exists(quantized_path):
    orig_mb = os.path.getsize(model_path) / 1024 / 1024
    print(f"[量化] 原始模型 {orig_mb:.0f}MB，正在生成 INT8 量化版（~5 秒）...")
    try:
        calib_reader = _CalibReader()
        quantize_static(
            model_path,
            quantized_path,
            calibration_data_reader=calib_reader,
            weight_type=QuantType.QInt8,
            activation_type=QuantType.QInt8,
            calibrate_method=CalibrationMethod.MinMax,
            per_channel=False,  # per_channel=True 需要新版 ONNX Runtime 的 DequantizeLinear axis 属性
        )
        q_mb = os.path.getsize(quantized_path) / 1024 / 1024
        print(f"[量化] ✓ 成功！量化后 {q_mb:.0f}MB（压缩 {orig_mb / q_mb:.1f}x）")
    except Exception as e:
        print(f"[量化] ✗ 量化过程失败：{e}，回退到原始 FP32 模型")
        quantized_path = model_path
else:
    print(f"[量化] 发现已缓存的量化模型 {os.path.getsize(quantized_path) / 1024 / 1024:.0f}MB")

# 尝试加载量化模型，失败则删掉缓存重新量化
try:
    rmbg_model = rt.InferenceSession(quantized_path, sess_options=sess_options, providers=providers)
except Exception as e:
    if quantized_path != model_path:
        print(f"[量化] ⚠ 缓存模型加载失败 ({e})，删除并重新量化...")
        os.remove(quantized_path)
        try:
            calib_reader = _CalibReader()
            quantize_static(
                model_path,
                quantized_path,
                calibration_data_reader=calib_reader,
                weight_type=QuantType.QInt8,
                activation_type=QuantType.QInt8,
                calibrate_method=CalibrationMethod.MinMax,
                per_channel=False,
            )
            q_mb = os.path.getsize(quantized_path) / 1024 / 1024
            print(f"[量化] ✓ 重新量化成功！{q_mb:.0f}MB")
            rmbg_model = rt.InferenceSession(quantized_path, sess_options=sess_options, providers=providers)
        except Exception as e2:
            print(f"[量化] ✗ 重新量化也失败：{e2}，回退 FP32")
            rmbg_model = rt.InferenceSession(model_path, sess_options=sess_options, providers=providers)
    else:
        print(f"[量化] ✗ FP32 模型加载也失败：{e}")
        raise

# 结果缓存目录
CACHE_DIR = "cache"
os.makedirs(CACHE_DIR, exist_ok=True)

# 每个线程独立的输入缓冲区（thread-safe）
_thread_local = threading.local()
_INPUT_SIZE = 1024


def get_mask(img, s=1024):
    # 获取当前线程的私有缓冲区
    buf = getattr(_thread_local, "input_buf", None)
    if buf is None:
        buf = np.empty((_INPUT_SIZE, _INPUT_SIZE, 3), dtype=np.float32)
        _thread_local.input_buf = buf

    img = (img / 255).astype(np.float32)
    h, w = h0, w0 = img.shape[:-1]
    h, w = (s, int(s * w / h)) if h > w else (int(s * h / w), s)
    ph, pw = s - h, s - w
    buf.fill(0)
    buf[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w] = cv2.resize(img, (w, h))

    # CHW + batch dim
    img_input = np.ascontiguousarray(buf.transpose(2, 0, 1))[np.newaxis, :]
    mask = rmbg_model.run(None, {"img": img_input})[0][0]
    mask = np.transpose(mask, (1, 2, 0))
    mask = mask[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w]
    mask = cv2.resize(mask, (w0, h0))[:, :, np.newaxis]
    return mask


def rmbg_fn(img):
    # 计算输入图片的像素哈希（用于磁盘缓存）
    img_hash = hashlib.md5(img.tobytes()).hexdigest()
    cache_path = os.path.join(CACHE_DIR, f"{img_hash}.png")

    # 命中缓存
    if os.path.exists(cache_path):
        cached = cv2.imread(cache_path, cv2.IMREAD_UNCHANGED)
        if cached is not None:
            result_rgba = cv2.cvtColor(cached, cv2.COLOR_BGRA2RGBA)
            mask_gray = result_rgba[:, :, 3]
            mask_rgb = np.stack([mask_gray] * 3, axis=2)
            print(f"[单张] 命中缓存 ✓")
            return mask_rgb, result_rgba

    t0 = time.perf_counter()
    mask = get_mask(img)
    img = (mask * img + 255 * (1 - mask)).astype(np.uint8)
    mask = (mask * 255).astype(np.uint8)
    img = np.concatenate([img, mask], axis=2, dtype=np.uint8)
    mask = mask.repeat(3, axis=2)
    elapsed = time.perf_counter() - t0

    # 写入缓存（BGRA，与 batch_process 格式一致）
    cv2.imwrite(cache_path, cv2.cvtColor(img, cv2.COLOR_RGBA2BGRA))

    print(f"[单张] 推理耗时: {elapsed:.2f}s")
    return mask, img


def batch_process(files):
    """批量处理多张图片（并行推理 + 磁盘缓存）"""
    if not files:
        raise gr.Error("请至少选择一张图片")

    # ---- 第 1 步：读入内存 + 检查缓存 ----
    cached_results = {}   # idx → cache_path
    to_process = []       # (img_rgb, stem, idx, file_hash)

    for idx, f in enumerate(files):
        img_path = f.name if hasattr(f, "name") else str(f)
        with open(img_path, "rb") as fh:
            raw = fh.read()

        file_hash = hashlib.md5(raw).hexdigest()
        cache_path = os.path.join(CACHE_DIR, f"{file_hash}.png")

        if os.path.exists(cache_path):
            cached_results[idx] = cache_path
            print(f"  [{idx+1}] {Path(img_path).stem}: 命中缓存 ✓")
            continue

        img_bgr = cv2.imdecode(np.frombuffer(raw, np.uint8), cv2.IMREAD_COLOR)
        if img_bgr is None:
            continue
        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
        stem = Path(img_path).stem
        to_process.append((img_rgb, stem, idx, file_hash))

    if not to_process and not cached_results:
        raise gr.Error("没有有效的图片")

    tmp_dir = tempfile.mkdtemp()
    results = [None] * len(files)

    # 填充缓存命中结果
    for idx, path in cached_results.items():
        dest = os.path.join(tmp_dir, f"cached_{idx}.png")
        cv2.imwrite(dest, cv2.imread(path, cv2.IMREAD_UNCHANGED))
        results[idx] = dest

    if to_process:
        print(f"[批量] 开始并行推理 {len(to_process)} 张图片（2 线程）...")
        t_start = time.perf_counter()

        def _infer_one(img_rgb, stem, idx, file_hash):
            t0 = time.perf_counter()
            mask, result_rgba = rmbg_fn(img_rgb)
            result_bgr = cv2.cvtColor(result_rgba, cv2.COLOR_RGBA2BGRA)

            # 写入磁盘缓存
            cv2.imwrite(os.path.join(CACHE_DIR, f"{file_hash}.png"), result_bgr)

            out_path = os.path.join(tmp_dir, f"{stem}_rmbg.png")
            cv2.imwrite(out_path, result_bgr)
            elapsed = time.perf_counter() - t0
            print(f"  [{idx+1}/{len(files)}] {stem}: {elapsed:.2f}s")
            return idx, out_path

        with ThreadPoolExecutor(max_workers=2) as pool:
            futures = {
                pool.submit(_infer_one, img, stem, idx, fh): idx
                for img, stem, idx, fh in to_process
            }
            for future in as_completed(futures):
                idx, path = future.result()
                results[idx] = path

        t_total = time.perf_counter() - t_start
        print(f"[批量] 完成！推理 {len(to_process)} 张，缓存命中 {len(cached_results)} 张，总耗时 {t_total:.2f}s")

    return results


with gr.Blocks(title="Anime Remove Background") as app:
    gr.Markdown(
        "# Anime Remove Background\n\n"
        "基于 [skytnt/anime-segmentation](https://github.com/SkyTNT/anime-segmentation/) 的动漫图片去背景工具"
    )

    with gr.Tab("单张处理"):
        with gr.Row():
            with gr.Column():
                input_img = gr.Image(label="输入图片")
                examples_data = [[f"examples/{x:02d}.jpg"] for x in range(1, 4)]
                examples = gr.Dataset(components=[input_img], samples=examples_data)
            with gr.Column():
                run_btn = gr.Button("去除背景", variant="primary")
                output_mask = gr.Image(label="蒙版")
                output_img = gr.Image(label="结果（RGBA）", image_mode="RGBA")
        examples.click(lambda x: x[0], [examples], [input_img])
        run_btn.click(rmbg_fn, [input_img], [output_mask, output_img], api_name="rmbg_fn")

    with gr.Tab("批量处理"):
        with gr.Row():
            with gr.Column():
                input_files = gr.File(
                    label="上传多张图片",
                    file_count="multiple",
                    file_types=["image"],
                )
                batch_btn = gr.Button("批量去背景", variant="primary")
            with gr.Column():
                output_gallery = gr.Gallery(label="去背景结果", columns=3)
        batch_btn.click(
            batch_process,
            [input_files],
            [output_gallery],
            api_name="batch_rmbg",
        )

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0")