import io
import os
import shutil
import subprocess
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

# ---- 自动拉取 Xet 指针文件 ----
def _is_valid_image(fpath):
    try:
        with open(fpath, "rb") as f:
            head = f.read(16)
        if len(head) < 4:
            return False
        if head[:8] == b"\x89PNG\r\n\x1a\n":
            return True
        if head[:3] == b"\xff\xd8\xff":
            return True
        if head[:6] in (b"GIF87a", b"GIF89a"):
            return True
        if head[:4] == b"RIFF" and head[8:12] == b"WEBP":
            return True
        if head[:2] == b"BM":
            return True
        return False
    except Exception:
        return False


def _is_valid_image_bytes(data):
    if len(data) < 4:
        return False
    if data[:8] == b"\x89PNG\r\n\x1a\n":
        return True
    if data[:3] == b"\xff\xd8\xff":
        return True
    if data[:6] in (b"GIF87a", b"GIF89a"):
        return True
    if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
        return True
    if data[:2] == b"BM":
        return True
    return False


def _download_one(urls, fpath):
    """尝试从多个源下载，成功一个即返回。"""
    if isinstance(urls, str):
        urls = [urls]
    for url in urls:
        for attempt in range(3):
            try:
                req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
                with urllib.request.urlopen(req, timeout=120) as resp:
                    data = resp.read()
                if data and _is_valid_image_bytes(data):
                    with open(fpath, "wb") as f:
                        f.write(data)
                    return True
            except Exception:
                if attempt < 2:
                    import time
                    time.sleep(1)
    return False


def _materialize_xet():
    repo_id = "donotthink/meme"
    hf_base = f"https://huggingface.co/spaces/{repo_id}/resolve/main"
    gh_base = "https://raw.githubusercontent.com/Memes02/meme-generator/main"
    exts = (".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp")
    need_download = []
    total = 0
    for root, _dirs, files in os.walk("meme_generator"):
        for fname in files:
            if not fname.lower().endswith(exts):
                continue
            total += 1
            fpath = os.path.join(root, fname)
            if not _is_valid_image(fpath):
                need_download.append(fpath)
    print(f"[xet] scanned {total} images, {len(need_download)} need download", flush=True)
    if not need_download:
        return
    count = 0
    failed_files = []
    with ThreadPoolExecutor(max_workers=6) as executor:
        futures = {}
        for fpath in need_download:
            rel = os.path.relpath(fpath, ".").replace(os.sep, "/")
            # 先试 HF，失败再试 GitHub
            urls = [f"{hf_base}/{rel}", f"{gh_base}/{rel}"]
            futures[executor.submit(_download_one, urls, fpath)] = fpath
        for fut in as_completed(futures):
            fpath = futures[fut]
            if fut.result():
                count += 1
            else:
                failed_files.append(fpath)
            if (count + len(failed_files)) % 100 == 0:
                print(f"[xet] progress {count + len(failed_files)}/{len(need_download)}", flush=True)
    print(f"[xet] materialized {count} files, {len(failed_files)} failures", flush=True)
    if failed_files:
        print(f"[xet] failed files ({len(failed_files)}):", flush=True)
        for f in failed_files[:30]:
            print(f"  - {f}", flush=True)
        if len(failed_files) > 30:
            print(f"  ... and {len(failed_files) - 30} more", flush=True)


_materialize_xet()

# ---- 注册字体 ----
_font_src = Path("resources/fonts")
_font_dst = Path.home() / ".fonts"
if _font_src.exists():
    _font_dst.mkdir(parents=True, exist_ok=True)
    for _f in _font_src.iterdir():
        if _f.is_file():
            shutil.copy2(_f, _font_dst / _f.name)
    subprocess.run(["fc-cache", "-fv"], capture_output=True)

# ---- ZeroGPU 兼容层：必须有 @spaces.GPU 装饰的函数并实际调用，否则 ZeroGPU 会停止应用 ----
import spaces
import asyncio

@spaces.GPU
def _zero_gpu_warmup():
    """占位函数，让 ZeroGPU 检测到 GPU 函数调用。实际不使用 GPU。"""
    return "warmup-ok"

# 启动前实际调用一次，触发 spaces 库的 GPU 初始化
try:
    _r = _zero_gpu_warmup()
    if asyncio.iscoroutine(_r):
        asyncio.run(_r)
    print("[zero-gpu] warmup done", flush=True)
except Exception as _e:
    print(f"[zero-gpu] warmup skipped: {_e}", flush=True)

# ---- 启动 meme_generator 原生 FastAPI（无 Gradio，无 CSRF）----
import uvicorn
from meme_generator.app import app, register_routers

register_routers()

print("[meme] API server starting on 0.0.0.0:7860", flush=True)
print("[meme] endpoints: /memes/keys, /memes/{key}/, /memes/{key}/info, /docs", flush=True)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)