#!/usr/bin/env python3
"""
Web interface for video-subtitle-remover.

Usage:
    cp .env.example .env  # 编辑 .env 设置 ADMIN_PASSWORD
    python web_app.py

Platform worker: 在「系统设置」页面配置，无需环境变量。
"""

from __future__ import annotations

import asyncio
import collections
import cv2
import hashlib
import hmac
import itertools
import json
import logging
import os
import queue as q_module
import sqlite3
import sys
import threading
import time
import traceback
import uuid
from pathlib import Path
from typing import Any, Optional

# 加载 .env（同目录）
try:
    from dotenv import load_dotenv
    load_dotenv(Path(__file__).parent / ".env")
except ImportError:
    pass

import httpx
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException, Request
from fastapi.responses import FileResponse, HTMLResponse, RedirectResponse, Response
from pydantic import BaseModel

# ── project path ─────────────────────────────────────────────────────────────
VSR_PATH = os.environ.get("VSR_PATH", str(Path(__file__).parent))
if VSR_PATH not in sys.path:
    sys.path.insert(0, VSR_PATH)

try:
    from backend.config import config, tr
    from backend.tools.constant import InpaintMode
    from backend.main import SubtitleRemover
    import configparser as _cp
    _ch_ini = os.path.join(VSR_PATH, 'backend', 'interface', 'ch.ini')
    tr.read(_ch_ini, encoding='utf-8')
    VSR_AVAILABLE = True
    _import_err_msg = ""
except ImportError as _import_err:
    VSR_AVAILABLE = False
    _import_err_msg = str(_import_err)

# ── paths ─────────────────────────────────────────────────────────────────────
BASE_DIR = Path(__file__).parent
CACHE_DIR = BASE_DIR / "web_cache" / "downloads"
OUTPUT_DIR = BASE_DIR / "web_cache" / "outputs"
DB_PATH = BASE_DIR / "web_cache" / "tasks.db"
CACHE_DIR.mkdir(parents=True, exist_ok=True)
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

# ── auth ──────────────────────────────────────────────────────────────────────

_COOKIE = "vsr_session"
_COOKIE_MAX_AGE = 86400 * 30  # 30 天

_LOGIN_HTML = """<!DOCTYPE html>
<html lang="zh">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1">
<title>登录 · 视频字幕去除器</title>
<style>
*{{box-sizing:border-box;margin:0;padding:0}}
body{{display:flex;align-items:center;justify-content:center;min-height:100vh;
     background:#0d1117;font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",sans-serif}}
.card{{background:#161b22;border:1px solid #30363d;border-radius:12px;
      padding:40px 36px;width:100%;max-width:380px}}
h1{{font-size:1.2rem;font-weight:600;color:#e6edf3;margin-bottom:28px;text-align:center}}
label{{display:block;font-size:.85rem;color:#8b949e;margin-bottom:6px}}
input[type=password]{{width:100%;padding:10px 14px;background:#0d1117;border:1px solid #30363d;
                     border-radius:8px;font-size:1rem;color:#e6edf3;outline:none;
                     transition:border-color .2s}}
input[type=password]:focus{{border-color:#58a6ff}}
button{{width:100%;padding:11px;background:#1f6feb;color:#fff;border:none;border-radius:8px;
       font-size:1rem;cursor:pointer;margin-top:18px;transition:background .2s}}
button:hover{{background:#388bfd}}
.err{{color:#f85149;font-size:.85rem;margin-top:12px;text-align:center}}
</style>
</head>
<body>
<div class="card">
  <h1>视频字幕去除器</h1>
  <form method="post" action="/auth/login">
    <label for="password">密码</label>
    <input id="password" name="password" type="password" autofocus placeholder="请输入管理员密码">
    {error}
    <button type="submit">登录</button>
  </form>
</div>
</body>
</html>"""


def _admin_password() -> str:
    return (os.environ.get("ADMIN_PASSWORD") or "").strip()


def _make_token() -> str:
    pw = _admin_password()
    ts = str(int(time.time()))
    sig = hmac.new(pw.encode(), ts.encode(), hashlib.sha256).hexdigest()
    return f"{ts}.{sig}"


def _verify_token(token: str) -> bool:
    pw = _admin_password()
    if not pw:
        return True
    try:
        ts, sig = token.split(".", 1)
        expected = hmac.new(pw.encode(), ts.encode(), hashlib.sha256).hexdigest()
        return hmac.compare_digest(sig, expected) and int(ts) > time.time() - _COOKIE_MAX_AGE
    except Exception:
        return False


def _auth_ok(request: Request) -> bool:
    if not _admin_password():
        return True
    return _verify_token(request.cookies.get(_COOKIE, ""))


# ── database ──────────────────────────────────────────────────────────────────

def _db() -> sqlite3.Connection:
    conn = sqlite3.connect(str(DB_PATH))
    conn.row_factory = sqlite3.Row
    return conn


def init_db():
    conn = _db()
    conn.execute("""
        CREATE TABLE IF NOT EXISTS tasks (
            id TEXT PRIMARY KEY,
            url TEXT,
            is_cache_file BOOLEAN,
            cache_path TEXT,
            inpaint_mode TEXT,
            subtitle_areas TEXT,
            preview_mode BOOLEAN DEFAULT 0,
            detect_mode TEXT DEFAULT 'text',
            status TEXT,
            progress INTEGER,
            download_progress INTEGER,
            cached BOOLEAN,
            output_path TEXT,
            error TEXT,
            platform_task_id INTEGER,
            created_at REAL,
            started_at REAL,
            completed_at REAL,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    conn.execute("""
        CREATE TABLE IF NOT EXISTS settings (
            key TEXT PRIMARY KEY,
            value TEXT NOT NULL
        )
    """)
    conn.execute("""
        CREATE TABLE IF NOT EXISTS task_logs (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            task_id TEXT NOT NULL,
            message TEXT NOT NULL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    conn.execute("CREATE INDEX IF NOT EXISTS idx_task_logs_task_id ON task_logs(task_id)")
    for col, default in [
        ("preview_mode", "BOOLEAN DEFAULT 0"),
        ("detect_mode", "TEXT DEFAULT 'text'"),
        ("platform_task_id", "INTEGER"),
        ("started_at", "REAL"),
        ("completed_at", "REAL"),
        ("s3_failed", "BOOLEAN DEFAULT 0"),
    ]:
        try:
            conn.execute(f"ALTER TABLE tasks ADD COLUMN {col} {default}")
        except Exception:
            pass
    conn.commit()
    conn.close()


init_db()

WORKER_KEYS = ("platform_base_url", "platform_app_id", "platform_app_secret",
               "worker_concurrency", "worker_poll_interval")


def _settings_get_all() -> dict[str, str]:
    conn = _db()
    rows = conn.execute("SELECT key, value FROM settings").fetchall()
    conn.close()
    return {r["key"]: r["value"] for r in rows}


def _settings_set_many(data: dict[str, str]) -> None:
    conn = _db()
    for k, v in data.items():
        conn.execute("INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)", (k, v))
    conn.commit()
    conn.close()


# ── inpaint mode map ──────────────────────────────────────────────────────────

# STTN 时序 inpaint 无法去除持久水印（参考帧里也有水印），已移除。
# 未知/不支持的模式统一回退到 LaMa（空间 inpaint，对持久水印有效）。
INPAINT_MODE_MAP = {
    "lama":         lambda: InpaintMode.LAMA,
    "propainter":   lambda: InpaintMode.PROPAINTER,
    "opencv":       lambda: InpaintMode.OPENCV,
}
DEFAULT_INPAINT = "lama"

# ── task state ────────────────────────────────────────────────────────────────

tasks: dict[str, dict] = {}
task_lock = threading.Lock()
task_queue: q_module.Queue = q_module.Queue()
ws_connections: list[WebSocket] = []
_loop: Optional[asyncio.AbstractEventLoop] = None


def url_to_cache_path(url: str) -> Path:
    h = hashlib.md5(url.encode()).hexdigest()[:12]
    raw_name = url.rstrip("/").split("/")[-1].split("?")[0][:80]
    name = raw_name if raw_name and "." in raw_name else f"{h}.mp4"
    return CACHE_DIR / f"{h}_{name}"


def _task_snapshot(task: dict) -> dict:
    return {k: v for k, v in task.items() if k != "logs"}


def _push(coro):
    if _loop and not _loop.is_closed():
        asyncio.run_coroutine_threadsafe(coro, _loop)


async def _broadcast(msg: dict):
    dead = []
    text = json.dumps(msg)
    for ws in ws_connections:
        try:
            await ws.send_text(text)
        except Exception:
            dead.append(ws)
    for ws in dead:
        if ws in ws_connections:
            ws_connections.remove(ws)


def update_task(task_id: str, **kwargs):
    with task_lock:
        if task_id not in tasks:
            return
        tasks[task_id].update(kwargs)
        snap = _task_snapshot(tasks[task_id])
    _save_task_to_db(tasks[task_id])
    _push(_broadcast({"type": "task_update", "task": snap}))


def _save_task_to_db(task: dict):
    try:
        conn = _db()
        conn.execute("""
            INSERT OR REPLACE INTO tasks
            (id, url, is_cache_file, cache_path, inpaint_mode, subtitle_areas,
             preview_mode, detect_mode, status, progress, download_progress,
             cached, output_path, error, platform_task_id,
             created_at, started_at, completed_at, s3_failed)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """, (
            task["id"], task["url"], task.get("is_cache_file", False),
            task.get("cache_path", ""), task["inpaint_mode"],
            json.dumps(task["subtitle_areas"]), task.get("preview_mode", False),
            task.get("detect_mode", "text"), task["status"], task["progress"],
            task["download_progress"], task["cached"], task["output_path"],
            task["error"], task.get("platform_task_id"),
            task.get("created_at"), task.get("started_at"), task.get("completed_at"),
            task.get("s3_failed", False),
        ))
        conn.commit()
        conn.close()
    except Exception as e:
        print(f"Failed to save task to DB: {e}")


def _load_tasks_from_db():
    try:
        conn = _db()
        cursor = conn.execute("SELECT * FROM tasks ORDER BY created_at DESC LIMIT 100")
        keys = [d[0] for d in cursor.description]
        for row in cursor.fetchall():
            r = dict(zip(keys, row))
            task = {
                "id": r["id"], "url": r["url"],
                "is_cache_file": r.get("is_cache_file", False),
                "cache_path": r.get("cache_path", ""),
                "inpaint_mode": r["inpaint_mode"],
                "subtitle_areas": json.loads(r["subtitle_areas"] or "[]"),
                "preview_mode": r.get("preview_mode", False),
                "detect_mode": r.get("detect_mode", "text"),
                "status": r["status"], "progress": r["progress"],
                "phase": "pending", "detect_pct": 0, "process_pct": 0, "merge_pct": 0,
                "download_progress": r["download_progress"],
                "cached": r["cached"], "output_path": r["output_path"],
                "error": r["error"], "platform_task_id": r.get("platform_task_id"),
                "created_at": r.get("created_at"), "started_at": r.get("started_at"),
                "completed_at": r.get("completed_at"),
                "s3_failed": bool(r.get("s3_failed", False)),
                "logs": [],
            }
            # 重启时 downloading/processing 说明上次处理被中断，标为 error
            if task["status"] in ("downloading", "processing"):
                task["status"] = "error"
                task["error"] = "服务重启，任务中断"
                _save_task_to_db(task)
            with task_lock:
                tasks[task["id"]] = task
            if task["status"] == "pending":
                task_queue.put(task["id"])
        conn.close()
    except Exception as e:
        print(f"Failed to load tasks from DB: {e}")


def push_log(task_id: str, message: str):
    with task_lock:
        if task_id in tasks:
            tasks[task_id]["logs"].append(message)
    # 持久化到 DB
    try:
        conn = _db()
        conn.execute("INSERT INTO task_logs (task_id, message) VALUES (?, ?)", (task_id, message))
        conn.commit()
        conn.close()
    except Exception:
        pass
    _push(_broadcast({"type": "log", "task_id": task_id, "message": message}))


# ── local worker ──────────────────────────────────────────────────────────────

def _download(task_id: str, url: str, dest: Path) -> bool:
    update_task(task_id, status="downloading", download_progress=0)
    push_log(task_id, f"下载中: {url}")
    last_pct = -1
    try:
        with httpx.stream("GET", url, follow_redirects=True, timeout=7200) as r:
            r.raise_for_status()
            total = int(r.headers.get("content-length", 0))
            done = 0
            with open(dest, "wb") as f:
                for chunk in r.iter_bytes(chunk_size=1024 * 512):
                    f.write(chunk)
                    done += len(chunk)
                    if total:
                        pct = int(done / total * 100)
                        if pct >= last_pct + 2:
                            last_pct = pct
                            update_task(task_id, download_progress=pct)
                            push_log(task_id, f"下载 {pct}%（{done//1048576} MB / {total//1048576} MB）")
        push_log(task_id, f"下载完成 → {dest.name}")
        return True
    except Exception as e:
        push_log(task_id, f"下载失败: {e}")
        dest.unlink(missing_ok=True)
        return False


def _ensure_h264(task_id: str, video_path: str) -> str:
    """如果视频无法被 OpenCV 解码（如 AV1/HEVC），先用 FFmpeg 转成 H.264。
    返回可用的视频路径（原路径或转码后路径）。"""
    cap = cv2.VideoCapture(video_path)
    ret, _ = cap.read()
    cap.release()
    if ret:
        return video_path  # 可正常解码，无需转换

    # 无法解码，转码
    converted = str(Path(video_path).with_suffix('')) + '_h264.mp4'
    if Path(converted).exists():
        push_log(task_id, f"使用已有 H.264 转码版本: {Path(converted).name}")
        return converted

    push_log(task_id, "⚠ 视频格式（可能是 AV1/HEVC）OpenCV 无法解码，自动转换为 H.264…")
    try:
        import subprocess, shutil
        # 优先系统 ffmpeg（有 NVENC），项目自带的没有编译 NVENC 支持
        sys_ffmpeg = shutil.which('ffmpeg') or '/usr/bin/ffmpeg'
        push_log(task_id, f"FFmpeg 转码中（H.264 NVENC）…")
        cmd = [
            sys_ffmpeg, '-y', '-i', video_path,
            '-c:v', 'h264_nvenc', '-cq', '18', '-preset', 'p4',
            '-c:a', 'copy', '-loglevel', 'error', converted,
        ]
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=3600)
        if result.returncode != 0 or not Path(converted).exists():
            err = (result.stderr or '').strip()[:200]
            push_log(task_id, f"NVENC 失败（{err or 'unknown'}），回退 libx264…")
            cmd_x264 = [
                sys_ffmpeg, '-y', '-i', video_path,
                '-c:v', 'libx264', '-crf', '18', '-preset', 'fast',
                '-c:a', 'copy', '-loglevel', 'error', converted,
            ]
            result = subprocess.run(cmd_x264, capture_output=True, text=True, timeout=3600)
        if result.returncode == 0 and Path(converted).exists():
            size_mb = Path(converted).stat().st_size / 1048576
            push_log(task_id, f"✓ 转码完成 → {Path(converted).name} ({size_mb:.1f} MB)")
            return converted
        else:
            push_log(task_id, f"转码失败: {result.stderr[:200]}，尝试用原始文件继续")
            return video_path
    except Exception as e:
        push_log(task_id, f"转码异常: {e}，尝试用原始文件继续")
        return video_path


def _process(task_id: str, video_path: str, inpaint_mode_str: str, subtitle_areas: list,
             preview_mode: bool = False, detect_mode: str = "text"):
    if not VSR_AVAILABLE:
        push_log(task_id, f"错误：无法加载后端模块。{_import_err_msg}")
        update_task(task_id, status="error", error="VSR backend not available")
        return

    # 未知/已废弃模式（含旧的 sttn-*）统一回退 LaMa
    if inpaint_mode_str not in INPAINT_MODE_MAP:
        inpaint_mode_str = DEFAULT_INPAINT
    modeLabels = {'lama': 'LaMa', 'propainter': 'ProPainter', 'opencv': 'OpenCV'}
    push_log(task_id, f"启动 — 算法：{modeLabels.get(inpaint_mode_str, inpaint_mode_str)}")
    update_task(task_id, status="processing", progress=0, started_at=time.time())

    # ── 预检：AV1/HEVC 等 OpenCV 无法解码的格式，先转换为 H.264 ──────────────
    video_path = _ensure_h264(task_id, video_path)

    class _WebRemover(SubtitleRemover):
        _static_boxes = []          # 静态框（手动指定时用），注入所有帧
        _per_frame_boxes = None     # {frame_no: [(xmin,xmax,ymin,ymax)]} 逐帧框

        def append_output(self_, *args):
            msg = " ".join(str(a) for a in args)
            print(msg, flush=True)
            push_log(task_id, msg)

        def update_preview_with_comp(self_, frame_ori, frame_comp):
            pass

        def run(self_):
            if self_._per_frame_boxes or self_._static_boxes:
                from backend.tools.subtitle_detect import SubtitleDetect as _SD
                _orig = _SD.find_subtitle_frame_no
                _static = list(self_._static_boxes or [])
                _perframe = self_._per_frame_boxes or {}

                def _patched(sd_self, sub_remover=None):
                    # 逐帧框模式：直接用检测到的逐帧字典，不跑 OCR（淡水印 OCR 检测不到）
                    cap = cv2.VideoCapture(sd_self.video_path)
                    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
                    cap.release()
                    result = {}
                    for frame_no in range(1, total + 1):
                        boxes = list(_perframe.get(frame_no, []))
                        if _static:
                            boxes = boxes + _static
                        if boxes:
                            result[frame_no] = boxes
                    push_log(task_id, f"水印逐帧框：{len(result)}/{total} 帧有水印框")
                    return result
                _SD.find_subtitle_frame_no = _patched
                try:
                    super().run()
                finally:
                    _SD.find_subtitle_frame_no = _orig
            else:
                super().run()

        def merge_audio_to_video(self_):
            push_log(task_id, "正在合并音频...")
            update_task(task_id, detect_pct=100, process_pct=100, merge_pct=50, phase="merging")
            super().merge_audio_to_video()
            update_task(task_id, detect_pct=100, process_pct=100, merge_pct=100, phase="merging")
            push_log(task_id, "音频合并完成")

    preview_clip_path = None
    if preview_mode:
        try:
            from backend.tools.ffmpeg_cli import FFmpegCLI
            import subprocess
            preview_clip_path = str(Path(video_path).parent / f"_preview_{task_id}.mp4")
            cut_cmd = [FFmpegCLI.instance().ffmpeg_path, "-y", "-i", video_path,
                       "-t", "10", "-c", "copy", "-loglevel", "error", preview_clip_path]
            push_log(task_id, "预览模式：截取前10秒...")
            subprocess.run(cut_cmd, check=True, timeout=60)
            push_log(task_id, "截取完成，开始处理预览片段")
            video_path = preview_clip_path
        except Exception as e:
            push_log(task_id, f"截取失败，改为处理完整视频: {e}")
            preview_clip_path = None

    try:
        sr = _WebRemover(video_path, gui_mode=False)
        if preview_mode:
            push_log(task_id, f"预览片段：{sr.frame_count} 帧 / {sr.fps:.1f}fps")

        mode_factory = INPAINT_MODE_MAP.get(inpaint_mode_str, INPAINT_MODE_MAP[DEFAULT_INPAINT])
        config.set(config.inpaintMode, mode_factory())

        if not subtitle_areas:
            if detect_mode == "smart":
                # 水印模式：逐帧检测半透明/移动文字水印 + 彩色 logo（每秒一帧采样）
                push_log(task_id, "水印检测中（逐帧 top-hat + 彩色 logo）...")
                from backend.tools.smart_detect import detect_watermark_per_frame
                per_frame = detect_watermark_per_frame(video_path)
                if per_frame:
                    n_boxes = sum(len(v) for v in per_frame.values())
                    sr._per_frame_boxes = per_frame
                    push_log(task_id, f"水印检测：{len(per_frame)} 帧共 {n_boxes} 个水印框")
                else:
                    # 未检测到水印：这类淡水印 OCR 也看不到，跑 OCR 纯浪费 → 原样输出
                    push_log(task_id, "未检测到水印，原样输出原视频")
                    out_stem = Path(video_path).stem
                    sr.video_out_path = str(OUTPUT_DIR / f"{out_stem}_{task_id[:8]}_no_sub.mp4")
                    import shutil
                    shutil.copy2(video_path, sr.video_out_path)
                    update_task(task_id, status="done", progress=100,
                                detect_pct=100, process_pct=100, merge_pct=100,
                                output_path=sr.video_out_path, completed_at=time.time())
                    push_log(task_id, f"完成（未检测到水印，输出原视频）→ {sr.video_out_path}")
                    return
            else:
                # 文字模式：标准 OCR 逐帧检测（走原生 find_subtitle_frame_no）
                push_log(task_id, "文字检测（OCR）将逐帧进行")

        if subtitle_areas:
            if detect_mode == "smart":
                sr._static_boxes = (sr._static_boxes or []) + [
                    (int(a[2]/100*sr.frame_width), int(a[3]/100*sr.frame_width),
                     int(a[0]/100*sr.frame_height), int(a[1]/100*sr.frame_height))
                    for a in subtitle_areas
                ]
                push_log(task_id, f"手动指定 {len(subtitle_areas)} 个水印区域 → 强制每帧去除")
            else:
                sr.sub_areas = [
                    (int(a[0]/100*sr.frame_height), int(a[1]/100*sr.frame_height),
                     int(a[2]/100*sr.frame_width), int(a[3]/100*sr.frame_width))
                    for a in subtitle_areas
                ]

        out_stem = Path(video_path).stem
        sr.video_out_path = str(OUTPUT_DIR / f"{out_stem}_{task_id[:8]}_no_sub.mp4")

        _state = {"phase": "detecting", "processing": False, "prev_raw": 0}

        def _on_progress(pct, finished):
            try:
                raw = int(pct)
                prev = _state["prev_raw"]
                _state["prev_raw"] = raw
                if not _state["processing"] and prev >= 40 and raw < 15:
                    _state["processing"] = True
                if _state["processing"] or raw > 50:
                    _state["processing"] = True
                    update_task(task_id, detect_pct=100, process_pct=min(raw, 98), merge_pct=0, phase="processing")
                else:
                    update_task(task_id, detect_pct=min(raw * 2, 99), process_pct=0, merge_pct=0, phase="detecting")
            except Exception as e:
                print(f"[ERROR] Progress callback failed: {e}", flush=True)

        push_log(task_id, f"检测模式：{'智能（SAM2）' if detect_mode == 'smart' else '文字（OCR）'}")
        sr.add_progress_listener(_on_progress)
        sr.run()

        if preview_mode and hasattr(sr, 'video_temp_file') and sr.video_temp_file:
            import shutil
            push_log(task_id, "预览模式：跳过音频合并，直接输出")
            try:
                shutil.copy2(sr.video_temp_file.name, sr.video_out_path)
            except Exception as e:
                push_log(task_id, f"复制文件失败: {e}")

        update_task(task_id, status="done", progress=100,
                    output_path=sr.video_out_path, completed_at=time.time())
        push_log(task_id, f"完成 → {sr.video_out_path}")

    except Exception:
        err = traceback.format_exc()
        push_log(task_id, f"错误：\n{err}")
        update_task(task_id, status="error", error=err.splitlines()[-1])
    finally:
        if preview_clip_path and Path(preview_clip_path).exists():
            try:
                Path(preview_clip_path).unlink()
            except Exception:
                pass


def _run_task(task_id: str):
    with task_lock:
        task = tasks.get(task_id)
    if not task:
        return

    if task.get("is_cache_file"):
        cache_path = Path(task["cache_path"])
        if not cache_path.exists():
            update_task(task_id, status="error", error="缓存文件不存在")
            return
        push_log(task_id, f"使用本地缓存: {cache_path.name}")
        update_task(task_id, cached=True)
    else:
        url = task["url"]
        cache_path = url_to_cache_path(url)
        if cache_path.exists():
            push_log(task_id, f"缓存命中: {cache_path.name}")
            update_task(task_id, cached=True)
        else:
            if not _download(task_id, url, cache_path):
                update_task(task_id, status="error", error="下载失败")
                return

    _process(task_id, str(cache_path), task["inpaint_mode"], task["subtitle_areas"],
             task.get("preview_mode", False), task.get("detect_mode", "text"))


def _worker():
    while True:
        task_id = task_queue.get()
        if task_id is None:
            break
        try:
            _run_task(task_id)
        except Exception:
            traceback.print_exc()
        finally:
            task_queue.task_done()


# ── platform worker ───────────────────────────────────────────────────────────

logger = logging.getLogger(__name__)


class WorkerLogBuffer:
    _seq = itertools.count(1)

    def __init__(self, maxlen: int = 1000) -> None:
        self._buf: collections.deque[dict] = collections.deque(maxlen=maxlen)
        self._lock = threading.Lock()

    def push(self, level: str, msg: str) -> None:
        with self._lock:
            self._buf.append({"seq": next(self._seq), "ts": time.time(), "level": level, "msg": msg})

    def since(self, seq: int) -> tuple[list[dict], int]:
        with self._lock:
            lines = [e for e in self._buf if e["seq"] > seq]
        return lines, (lines[-1]["seq"] if lines else seq)


class _WorkerLogHandler(logging.Handler):
    """把 Python logging 桥接到 WorkerLogBuffer，供前端轮询。"""
    def __init__(self, buf: WorkerLogBuffer) -> None:
        super().__init__()
        self._buf = buf

    def emit(self, record: logging.LogRecord) -> None:
        try:
            self._buf.push(record.levelname, self.format(record))
        except Exception:
            pass


class VSRTaskWorker:
    """轮询外部平台认领 VSR 任务，执行后上报结果。"""

    def __init__(self, *, base_url: str, app_id: str, app_secret: str,
                 slot_id: int, slot_total: int, poll_interval: float = 5.0) -> None:
        self._base_url = base_url.rstrip("/")
        self._app_id = app_id
        self._app_secret = app_secret
        self._slot_id = slot_id
        self._slot_total = slot_total
        self._poll_interval = poll_interval
        self._active: set[asyncio.Task] = set()

    async def run(self) -> None:
        logger.info(f"[vsr_worker] 启动 slot={self._slot_id} → {self._base_url}")
        async with httpx.AsyncClient(base_url=self._base_url, timeout=30.0) as client:
            self._client = client
            while True:
                try:
                    await self._tick()
                except Exception:
                    logger.exception("[vsr_worker] tick 异常")
                await asyncio.sleep(self._poll_interval)

    async def _tick(self) -> None:
        if self._active:
            return
        task = await self._claim()
        if task is None:
            return
        t = asyncio.create_task(self._run_task(task))
        self._active.add(t)
        t.add_done_callback(self._active.discard)

    async def _claim(self) -> dict[str, Any] | None:
        resp = await self._signed_post(
            "/openapi/watermark-task/claim",
            {"slot_id": self._slot_id, "slot_total": self._slot_total},
        )
        if resp is None:
            return None
        task = resp.get("task")
        if task is None:
            logger.debug(f"[vsr_worker] slot={self._slot_id} → 无待处理任务")
        return task

    async def _run_task(self, ptask: dict[str, Any]) -> None:
        platform_task_id: int = ptask["id"]
        video_url: str = ptask.get("media_url", "") or ptask.get("video_url", "")
        logger.info(f"[vsr_worker] 认领任务 #{platform_task_id} url={video_url[:60]}")

        task_id = uuid.uuid4().hex[:8]
        cache_path = url_to_cache_path(video_url)
        cached = cache_path.exists()

        new_task = {
            "id": task_id,
            "url": video_url,
            "is_cache_file": False,
            "cache_path": str(cache_path),
            "inpaint_mode": ptask.get("inpaint_mode", "lama"),
            "subtitle_areas": ptask.get("subtitle_areas", []),
            "preview_mode": ptask.get("preview_mode", False),
            "detect_mode": ptask.get("detect_mode", "text"),
            "status": "pending",
            "progress": 0,
            "phase": "pending",
            "detect_pct": 0, "process_pct": 0, "merge_pct": 0,
            "download_progress": 0,
            "cached": cached,
            "output_path": None,
            "error": None,
            "s3_failed": False,
            "platform_task_id": platform_task_id,
            "logs": [],
            "created_at": time.time(),
            "started_at": None,
            "completed_at": None,
        }
        with task_lock:
            tasks[task_id] = new_task

        # 立即广播给前端，使任务出现在列表中
        _save_task_to_db(new_task)
        _push(_broadcast({"type": "task_update", "task": _task_snapshot(new_task)}))

        # 在线程中同步执行任务
        def _do():
            _run_task(task_id)

        loop = asyncio.get_running_loop()

        # 进度上报后台 loop
        async def _progress_reporter():
            reported_logs = 0
            while True:
                await asyncio.sleep(2)
                with task_lock:
                    t = tasks.get(task_id)
                if not t:
                    break
                logs = t.get("logs", [])
                for msg in logs[reported_logs:]:
                    await self._report_progress(platform_task_id, msg)
                reported_logs = len(logs)
                if t["status"] in ("done", "error"):
                    break

        reporter = asyncio.create_task(_progress_reporter())
        try:
            await loop.run_in_executor(None, _do)
        except Exception:
            logger.exception(f"[vsr_worker] 任务 #{platform_task_id} 执行异常")
            update_task(task_id, status="error", error="worker 执行异常，详见 Worker 日志")
        finally:
            reporter.cancel()

        with task_lock:
            final = tasks.get(task_id, {})

        if final.get("status") == "done" and final.get("output_path"):
            output_url = await self._upload_result_video(final["output_path"], task_id=task_id)
            if output_url:
                # 上传 S3 成功，才算真正完成
                await self._report_result(platform_task_id, status="done", output_url=output_url)
                push_log(task_id, f"✓ 已上传 S3：{output_url}")
            else:
                # 处理成功但上传失败 → 保留 done 状态（本地文件可播放/下载），仅报错给平台
                logger.error(f"[vsr_worker] 任务 #{platform_task_id} 处理成功但上传 S3 失败")
                update_task(task_id, s3_failed=True, error="S3 上传失败，视频已在本地，可直接下载")
                push_log(task_id, "✗ S3 上传失败，但本地文件完好，可点击下载按钮获取")
                await self._report_result(platform_task_id, status="failed",
                                          error_message="处理完成但视频上传 S3 失败")
        else:
            await self._report_result(platform_task_id, status="failed",
                                      error_message=final.get("error", "unknown"))

    async def _report_progress(self, platform_task_id: int, message: str) -> None:
        await self._signed_post(
            f"/openapi/watermark-task/{platform_task_id}/progress",
            {"message": message},
            headers={"X-Worker-Slot": str(self._slot_id)},
        )

    async def _report_result(self, platform_task_id: int, *, status: str,
                             output_url: str | None = None,
                             error_message: str = "") -> None:
        await self._signed_post(
            f"/openapi/watermark-task/{platform_task_id}/result",
            {"status": status, "output_url": output_url, "error_message": error_message},
            headers={"X-Worker-Slot": str(self._slot_id)},
        )
        logger.info(f"[vsr_worker] 任务 #{platform_task_id} 上报结果 status={status}")

    async def _upload_result_video(self, local_path: str, task_id: str | None = None) -> str | None:
        """处理完成后将输出视频上传到平台，返回 play_url（S3 地址）。流式上传，不一次性读入内存。"""
        import hashlib as _hl
        def _log(msg: str):
            logger.info(msg)
            if task_id:
                push_log(task_id, msg)
        try:
            file_path = Path(local_path)
            file_size = file_path.stat().st_size
            filename = file_path.name

            # 流式计算 SHA256，避免 read_bytes() 把几百 MB 全加载进内存
            h = _hl.sha256()
            with open(local_path, "rb") as f:
                for chunk in iter(lambda: f.read(8 * 1024 * 1024), b""):
                    h.update(chunk)
            sha256 = h.hexdigest()
            _log(f"[S3] 准备上传 {filename}，大小={file_size//1048576}MB")

            sign_resp = await self._signed_post(
                "/openapi/file/upload/get_upload_sign",
                {
                    "sha256": sha256,
                    "size": file_size,
                    "filename": filename,
                    "content_type": "video/mp4",
                    "file_type": "video",
                    "biz_ref": "vsr-output",
                    "no_transcode": True,
                },
            )
            if sign_resp is None:
                _log("[S3] get_upload_sign 返回空")
                return None
            action = sign_resp.get("action", "")
            if action == "skip":
                _log("[S3] 文件已存在，直接复用")
                return sign_resp.get("play_url")
            presigned_url = sign_resp.get("presigned_url", "")
            object_key = sign_resp.get("object_key", "")
            if not presigned_url or not object_key:
                _log(f"[S3] upload sign 未返回 presigned_url/object_key: {sign_resp}")
                return None

            # 用 asyncio executor 做同步分块读取，避免 sync file 阻塞 async event loop
            _log(f"[S3] 开始上传 {file_size//1048576}MB → S3…")
            CHUNK = 8 * 1024 * 1024  # 8MB chunks

            async def _file_chunks():
                loop = asyncio.get_running_loop()
                with open(local_path, "rb") as fobj:
                    while True:
                        chunk = await loop.run_in_executor(None, fobj.read, CHUNK)
                        if not chunk:
                            break
                        yield chunk

            async with httpx.AsyncClient(timeout=7200) as upload_client:
                put_resp = await upload_client.put(
                    presigned_url,
                    content=_file_chunks(),
                    headers={"Content-Type": "video/mp4", "Content-Length": str(file_size)},
                )
            if put_resp.status_code not in (200, 204):
                _log(f"[S3] PUT 失败 HTTP {put_resp.status_code}: {put_resp.text[:200]}")
                return None
            _log("[S3] 上传完成，注册资产…")

            reg_resp = await self._signed_post(
                "/openapi/file/assets/register",
                {
                    "sha256": sha256,
                    "object_key": object_key,
                    "size": file_size,
                    "filename": filename,
                    "content_type": "video/mp4",
                    "file_type": "video",
                    "biz_ref": "vsr-output",
                    "no_transcode": True,
                },
            )
            play_url = reg_resp.get("play_url") if reg_resp else None
            _log(f"[S3] 资产注册完成 play_url={play_url}")
            return play_url
        except Exception as e:
            import traceback as _tb
            _log(f"[S3] 上传失败: {e}\n{_tb.format_exc()[-400:]}")
            logger.warning(f"[vsr_worker] 上传结果视频失败: {e}")
            return None

    async def _signed_post(self, path: str, body: dict[str, Any],
                           headers: dict[str, str] | None = None) -> dict[str, Any] | None:
        body_str = json.dumps(body, separators=(",", ":"), ensure_ascii=False)
        timestamp = str(int(time.time()))
        nonce = uuid.uuid4().hex
        payload_str = "\n".join([self._app_id, timestamp, nonce, body_str])
        signature = hmac.new(self._app_secret.encode(), payload_str.encode(), hashlib.sha256).hexdigest()
        req_headers = {
            "Content-Type": "application/json",
            "X-App-Id": self._app_id,
            "X-Timestamp": timestamp,
            "X-Nonce": nonce,
            "X-Signature": signature,
        }
        if headers:
            req_headers.update(headers)
        try:
            resp = await self._client.post(path, content=body_str, headers=req_headers)
            resp.raise_for_status()
            return resp.json()
        except httpx.HTTPStatusError as exc:
            logger.warning(f"[vsr_worker] POST {path} → {exc.response.status_code}: {exc.response.text[:200]}")
        except Exception:
            logger.exception(f"[vsr_worker] POST {path} 失败")
        return None


class WorkerManager:
    def __init__(self) -> None:
        self._tasks: list[asyncio.Task] = []
        self.log_buffer = WorkerLogBuffer()
        _handler = _WorkerLogHandler(self.log_buffer)
        _handler.setFormatter(logging.Formatter("%(asctime)s  %(message)s", datefmt="%H:%M:%S"))
        _wl = logging.getLogger(__name__)
        _wl.setLevel(logging.DEBUG)
        _wl.addHandler(_handler)

    def start(self, *, base_url: str, app_id: str, app_secret: str,
              concurrency: int, poll_interval: float) -> None:
        concurrency = max(1, concurrency)
        for slot_id in range(concurrency):
            worker = VSRTaskWorker(
                base_url=base_url, app_id=app_id, app_secret=app_secret,
                slot_id=slot_id, slot_total=concurrency, poll_interval=poll_interval,
            )
            self._tasks.append(asyncio.create_task(worker.run()))
        logger.info(f"[WorkerManager] 启动 {concurrency} 个 worker slot")
        self.log_buffer.push("info", f"启动 {concurrency} 个 worker slot → {base_url}")

    def stop(self) -> None:
        for t in self._tasks:
            t.cancel()
        self._tasks.clear()

    def restart(self, *, base_url: str, app_id: str, app_secret: str,
                concurrency: int, poll_interval: float) -> None:
        self.stop()
        if base_url and app_id and app_secret and concurrency >= 1:
            self.start(base_url=base_url, app_id=app_id, app_secret=app_secret,
                       concurrency=concurrency, poll_interval=poll_interval)

    @property
    def running(self) -> bool:
        return bool(self._tasks)


# ── FastAPI app ───────────────────────────────────────────────────────────────

app = FastAPI(title="VSR Web")
_worker_manager = WorkerManager()

_worker_thread = threading.Thread(target=_worker, daemon=True, name="vsr-worker")
_worker_thread.start()


@app.middleware("http")
async def _auth_middleware(request: Request, call_next):
    path = request.url.path
    skip = {"/login", "/auth/login", "/auth/logout", "/health", "/ws"}
    if path in skip or path.startswith("/static"):
        return await call_next(request)
    if not _auth_ok(request):
        if path.startswith("/api/"):
            return Response(content='{"detail":"Unauthorized"}',
                            status_code=401, media_type="application/json")
        return RedirectResponse(url="/login", status_code=302)
    return await call_next(request)


@app.on_event("startup")
async def _startup():
    global _loop
    _loop = asyncio.get_event_loop()
    _load_tasks_from_db()

    saved = _settings_get_all()
    base_url = saved.get("platform_base_url", "").strip()
    app_id = saved.get("platform_app_id", "").strip()
    app_secret = saved.get("platform_app_secret", "").strip()
    concurrency = int(saved.get("worker_concurrency", "1"))
    poll_interval = float(saved.get("worker_poll_interval", "5"))
    _worker_manager.restart(base_url=base_url, app_id=app_id, app_secret=app_secret,
                            concurrency=concurrency, poll_interval=poll_interval)


@app.on_event("shutdown")
async def _shutdown():
    _worker_manager.stop()


# ── auth endpoints ────────────────────────────────────────────────────────────

@app.get("/login", include_in_schema=False)
async def login_page():
    return HTMLResponse(_LOGIN_HTML.format(error=""))


@app.post("/auth/login", include_in_schema=False)
async def do_login(request: Request):
    form = await request.form()
    pw = (form.get("password") or "").strip()
    if hmac.compare_digest(pw.encode(), _admin_password().encode()):
        resp = RedirectResponse(url="/", status_code=302)
        resp.set_cookie(_COOKIE, _make_token(), max_age=_COOKIE_MAX_AGE, httponly=True, samesite="lax")
        return resp
    html = _LOGIN_HTML.format(error='<p class="err">密码错误，请重试</p>')
    return HTMLResponse(html, status_code=401)


@app.get("/auth/logout", include_in_schema=False)
async def do_logout():
    resp = RedirectResponse(url="/login", status_code=302)
    resp.delete_cookie(_COOKIE)
    return resp


# ── main page ─────────────────────────────────────────────────────────────────

@app.get("/")
async def index():
    html = (BASE_DIR / "web" / "index.html").read_text(encoding="utf-8")
    return HTMLResponse(html)


@app.get("/health")
async def health():
    return {"status": "ok"}


# ── task API ──────────────────────────────────────────────────────────────────

class TaskCreate(BaseModel):
    url: str
    inpaint_mode: str = "lama"
    subtitle_areas: list = []
    preview_mode: bool = False
    detect_mode: str = "text"


@app.post("/api/tasks")
async def create_task(body: TaskCreate):
    task_id = uuid.uuid4().hex[:8]
    is_cache_file = body.url.startswith("cache_")
    if is_cache_file:
        cache_path = CACHE_DIR / body.url[6:]
        cached = cache_path.exists()
    else:
        cache_path = url_to_cache_path(body.url)
        cached = cache_path.exists()

    with task_lock:
        now = time.time()
        tasks[task_id] = {
            "id": task_id, "url": body.url,
            "is_cache_file": is_cache_file, "cache_path": str(cache_path),
            "inpaint_mode": body.inpaint_mode, "subtitle_areas": body.subtitle_areas,
            "preview_mode": body.preview_mode, "detect_mode": body.detect_mode,
            "status": "pending", "progress": 0,
            "phase": "pending", "detect_pct": 0, "process_pct": 0, "merge_pct": 0,
            "download_progress": 0, "cached": cached,
            "output_path": None, "error": None, "s3_failed": False, "platform_task_id": None,
            "logs": [], "created_at": now, "started_at": None, "completed_at": None,
        }
    task_queue.put(task_id)
    return {"task_id": task_id}


@app.get("/api/tasks")
async def list_tasks():
    with task_lock:
        return [_task_snapshot(t) for t in tasks.values()]


@app.get("/api/tasks/{task_id}")
async def get_task(task_id: str):
    with task_lock:
        task = tasks.get(task_id)
    if not task:
        raise HTTPException(404, "Not found")
    return task


@app.delete("/api/tasks/{task_id}")
async def delete_task(task_id: str):
    with task_lock:
        tasks.pop(task_id, None)
    conn = _db()
    conn.execute("DELETE FROM tasks WHERE id = ?", (task_id,))
    conn.execute("DELETE FROM task_logs WHERE task_id = ?", (task_id,))
    conn.commit()
    conn.close()
    return {"ok": True}


@app.get("/api/tasks/{task_id}/logs")
async def get_task_logs(task_id: str):
    conn = _db()
    rows = conn.execute(
        "SELECT message, created_at FROM task_logs WHERE task_id = ? ORDER BY id ASC",
        (task_id,)
    ).fetchall()
    conn.close()
    return {"logs": [{"message": r[0], "ts": r[1]} for r in rows]}


@app.get("/api/tasks/{task_id}/download")
async def download_result(task_id: str):
    with task_lock:
        task = tasks.get(task_id)
    if not task or not task.get("output_path"):
        raise HTTPException(404, "Result not ready")
    # 允许 done 或 S3 上传失败（本地文件完好）的任务下载
    if task["status"] not in ("done", "error") or (
        task["status"] == "error" and not task.get("s3_failed")
    ):
        raise HTTPException(404, "Result not ready")
    out = Path(task["output_path"])
    if not out.exists():
        raise HTTPException(404, "Output file missing")
    return FileResponse(str(out), filename=out.name, media_type="video/mp4")


@app.get("/api/tasks/{task_id}/stream")
async def stream_result(task_id: str, request: Request):
    with task_lock:
        task = tasks.get(task_id)
    if not task or not task.get("output_path"):
        raise HTTPException(404, "Result not ready")
    if task["status"] not in ("done", "error") or (
        task["status"] == "error" and not task.get("s3_failed")
    ):
        raise HTTPException(404, "Result not ready")
    out = Path(task["output_path"])
    if not out.exists():
        raise HTTPException(404, "Output file missing")
    file_size = out.stat().st_size
    range_header = request.headers.get("range")
    if range_header:
        try:
            range_value = range_header.replace("bytes=", "")
            start, end = range_value.split("-")
            start = int(start) if start else 0
            end = int(end) if end else file_size - 1
            if start >= file_size or end >= file_size or start > end:
                raise HTTPException(416, "Range Not Satisfiable")
            return FileResponse(str(out), media_type="video/mp4", status_code=206, headers={
                "Content-Range": f"bytes {start}-{end}/{file_size}",
                "Accept-Ranges": "bytes",
                "Content-Length": str(end - start + 1),
            })
        except (ValueError, IndexError):
            pass
    return FileResponse(str(out), filename=out.name, media_type="video/mp4",
                        headers={"Accept-Ranges": "bytes"})


@app.get("/api/tasks/{task_id}/source-stream")
async def stream_source(task_id: str, request: Request):
    """Stream the original (cached) video file."""
    with task_lock:
        task = tasks.get(task_id)
    if not task:
        raise HTTPException(404, "Not found")
    cache_path = Path(task.get("cache_path") or "")
    if not cache_path.exists():
        raise HTTPException(404, "Source file not available")
    file_size = cache_path.stat().st_size
    range_header = request.headers.get("range")
    if range_header:
        try:
            start_str, end_str = range_header.replace("bytes=", "").split("-")
            start = int(start_str) if start_str else 0
            end = int(end_str) if end_str else file_size - 1
            if start >= file_size or end >= file_size or start > end:
                raise HTTPException(416, "Range Not Satisfiable")
            return FileResponse(str(cache_path), media_type="video/mp4", status_code=206, headers={
                "Content-Range": f"bytes {start}-{end}/{file_size}",
                "Accept-Ranges": "bytes",
                "Content-Length": str(end - start + 1),
            })
        except (ValueError, IndexError):
            pass
    return FileResponse(str(cache_path), filename=cache_path.name, media_type="video/mp4",
                        headers={"Accept-Ranges": "bytes"})


# ── cache API ─────────────────────────────────────────────────────────────────

@app.get("/api/cache-check")
async def cache_check(url: str):
    return {"cached": url_to_cache_path(url).exists()}


@app.get("/api/cache-list")
async def cache_list():
    files = []
    if CACHE_DIR.exists():
        for f in sorted(CACHE_DIR.iterdir(), key=lambda x: x.stat().st_mtime, reverse=True):
            if f.is_file():
                files.append({"name": f.name, "size": f.stat().st_size, "mtime": f.stat().st_mtime})
    return {"files": files}


# ── settings API ──────────────────────────────────────────────────────────────

@app.get("/api/settings")
async def get_settings_api():
    saved = _settings_get_all()
    return {
        "platform_base_url": saved.get("platform_base_url", ""),
        "platform_app_id": saved.get("platform_app_id", ""),
        "platform_app_secret": saved.get("platform_app_secret", ""),
        "worker_concurrency": int(saved.get("worker_concurrency", "1")),
        "worker_poll_interval": float(saved.get("worker_poll_interval", "5")),
    }


@app.put("/api/settings")
async def put_settings(body: dict):
    allowed = set(WORKER_KEYS)
    to_save = {k: str(v) for k, v in body.items() if k in allowed}
    _settings_set_many(to_save)
    saved = _settings_get_all()
    _worker_manager.restart(
        base_url=saved.get("platform_base_url", "").strip(),
        app_id=saved.get("platform_app_id", "").strip(),
        app_secret=saved.get("platform_app_secret", "").strip(),
        concurrency=int(saved.get("worker_concurrency", "1")),
        poll_interval=float(saved.get("worker_poll_interval", "5")),
    )
    return {"ok": True, "worker_running": _worker_manager.running}


# ── worker log API ────────────────────────────────────────────────────────────

@app.get("/api/worker/logs")
async def get_worker_logs(since: int = 0):
    lines, next_seq = _worker_manager.log_buffer.since(since)
    return {"lines": lines, "next_seq": next_seq}


# ── status API ────────────────────────────────────────────────────────────────

@app.get("/api/status")
async def server_status():
    return {
        "vsr_available": VSR_AVAILABLE,
        "vsr_path": VSR_PATH,
        "queue_size": task_queue.qsize(),
        "worker_running": _worker_manager.running,
        "auth_enabled": bool(_admin_password()),
    }


# ── websocket ──────────────────────────────────────────────────────────────────

@app.websocket("/ws")
async def ws_endpoint(ws: WebSocket):
    await ws.accept()
    ws_connections.append(ws)
    try:
        with task_lock:
            init_tasks = [_task_snapshot(t) for t in tasks.values()]
        await ws.send_text(json.dumps({"type": "init", "tasks": init_tasks}))
        while True:
            await ws.receive_text()
    except WebSocketDisconnect:
        pass
    finally:
        if ws in ws_connections:
            ws_connections.remove(ws)


# ── entry point ───────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import uvicorn
    host = os.environ.get("HOST", "0.0.0.0")
    port = int(os.environ.get("PORT", "8001"))
    print(f"VSR_PATH = {VSR_PATH}")
    print(f"VSR backend available: {VSR_AVAILABLE}")
    if not VSR_AVAILABLE:
        print(f"  Import error: {_import_err_msg}")
    print(f"Auth enabled: {bool(_admin_password())}")
    print(f"Starting on http://{host}:{port}")
    uvicorn.run(app, host=host, port=port)