#!/usr/bin/env python3
"""
VideoResize - PT视频自动压缩服务
qBittorrent下载完成后自动压缩视频到目标目录,Web页面监控状态
"""

import json
import logging
import os
import re
import shutil
import subprocess
import threading
import time
from datetime import datetime
from pathlib import Path
from queue import Queue

from flask import Flask, jsonify, request

# ========== 配置 ==========
DOWNLOAD_DIR = "/data1/download"
OUTPUT_DIR = "/data1/download_resize"
TASKS_FILE = "/data1/download_resize/.tasks.json"
LOG_FILE = "/data1/download_resize/videoresize.log"
PORT = 5100
ENCODER = "hevc_nvenc"  # GPU: hevc_nvenc / CPU: libx265
CRF = 23               # libx265 用 CRF, nvenc 用 CQ (同数值)
PRESET = "medium"       # libx265 preset
NVENC_PRESET = "p7"     # NVENC 最高质量预设，适合离线压缩
WORKERS = 2             # 并发压缩数; 1080 Ti 消费级驱动最多 2 路 NVENC

# 归档模式: 这些目录下的文件走 CPU SVT-AV1 慢速高压(逼近奈飞式压缩率, 减 75-85%),
# 其余走 GPU 快压。填绝对路径前缀, 例:["/data1/download/收藏", "/data1/download/原盘"]
ARCHIVE_DIRS = []
ARCHIVE_CRF = 28        # SVT-AV1 CRF, 越小画质越高体积越大; 24-30 是视觉无损区间
ARCHIVE_PRESET = "5"    # SVT-AV1 preset 0-13, 越小越慢压得越狠; 4-6 平衡
MIN_BITRATE_KBPS = 5000
VIDEO_EXTS = {".mkv", ".mp4", ".avi", ".ts", ".rmvb", ".flv", ".wmv"}
FFMPEG = "/usr/lib/jellyfin-ffmpeg/ffmpeg"
FFPROBE = "/usr/lib/jellyfin-ffmpeg/ffprobe"
# ===========================

# 日志/任务文件都在 OUTPUT_DIR 下, 建目录必须早于 logging 打开日志文件
os.makedirs(OUTPUT_DIR, exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
    handlers=[
        logging.FileHandler(LOG_FILE, encoding="utf-8"),
        logging.StreamHandler(),
    ],
)
log = logging.getLogger("videoresize")

app = Flask(__name__)
tasks = []
tasks_lock = threading.Lock()
queue = Queue()


def save_tasks():
    with tasks_lock:
        # ponytail: atomic write — temp then rename, so an interrupted/concurrent
        # write can never leave a half-written (corrupt) tasks file on disk
        tmp = TASKS_FILE + ".tmp"
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump(tasks, f, ensure_ascii=False, indent=2)
        os.replace(tmp, TASKS_FILE)


def load_tasks():
    global tasks
    if os.path.exists(TASKS_FILE):
        try:
            with open(TASKS_FILE, encoding="utf-8") as f:
                tasks = json.load(f)
        except (json.JSONDecodeError, ValueError):
            bak = TASKS_FILE + ".corrupt"
            os.replace(TASKS_FILE, bak)
            log.error("Tasks file corrupt, backed up to %s, starting empty", bak)
            tasks = []
            return
        for i, t in enumerate(tasks):
            if t["status"] in ("queued", "compressing"):
                t["status"] = "queued"
                queue.put(i)


def ffprobe_info(filepath):
    cmd = [
        FFPROBE, "-v", "quiet", "-print_format", "json",
        "-show_streams", "-show_format", str(filepath),
    ]
    r = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
    return json.loads(r.stdout) if r.returncode == 0 else None


def get_duration_secs(info):
    try:
        return float(info["format"]["duration"])
    except (KeyError, ValueError, TypeError):
        return 0


def should_compress(filepath):
    info = ffprobe_info(filepath)
    if not info:
        return False, "ffprobe失败", None

    video_stream = next(
        (s for s in info.get("streams", []) if s.get("codec_type") == "video"), None
    )
    if not video_stream:
        return False, "无视频流", None

    codec = video_stream.get("codec_name", "").lower()
    if codec in ("hevc", "h265", "av1"):
        return False, f"已是{codec}编码", info

    bitrate = int(video_stream.get("bit_rate", 0)) / 1000
    if 0 < bitrate < MIN_BITRATE_KBPS:
        return False, f"视频码率过低({bitrate:.0f}kbps)", info

    return True, codec, info


def find_videos(path):
    p = Path(path)
    if p.is_file() and p.suffix.lower() in VIDEO_EXTS:
        return [p]
    if p.is_dir():
        return sorted(v for v in p.rglob("*") if v.suffix.lower() in VIDEO_EXTS)
    return []


def format_size(n):
    if n is None:
        return "-"
    for u in ("B", "KB", "MB", "GB", "TB"):
        if abs(n) < 1024:
            return f"{n:.1f}{u}"
        n /= 1024
    return f"{n:.1f}PB"


def rel_under_download(v):
    """v 相对 DOWNLOAD_DIR 的路径; 不在其下时退回文件名。保留完整子目录结构。"""
    try:
        return Path(v).relative_to(DOWNLOAD_DIR)
    except ValueError:
        return Path(Path(v).name)


def is_archive(src):
    """src 是否落在归档目录下 (走 CPU SVT-AV1 高压)。"""
    s = os.path.abspath(str(src))
    for d in ARCHIVE_DIRS:
        d = os.path.abspath(d)
        if s == d or s.startswith(d + os.sep):
            return True
    return False


def output_path(src, task=None):
    if task and task.get("out"):
        return Path(task["out"])
    return Path(OUTPUT_DIR) / rel_under_download(src)


def nvenc_video_options():
    return [
        "-c:v", ENCODER, "-preset", NVENC_PRESET, "-tune", "hq",
        "-rc", "vbr", "-cq", str(CRF), "-b:v", "0",
        "-spatial_aq", "1", "-aq-strength", "8", "-rc-lookahead", "20",
    ]


def find_bdmv_discs(root):
    """找出所有蓝光原盘目录 (含 BDMV/STREAM 的文件夹)。"""
    discs = []
    for bdmv in Path(root).rglob("BDMV"):
        if bdmv.is_dir() and (bdmv / "STREAM").is_dir():
            discs.append(bdmv.parent)
    return discs


def main_title_m2ts(disc_dir):
    """选最大的 m2ts 作为正片。
    ponytail: 最大文件启发式 — 绝大多数碟成立; 个别分段/多版本碟可能选错,
    需要更准就得解析 BDMV/PLAYLIST/*.mpls 播放列表。"""
    streams = list((Path(disc_dir) / "BDMV" / "STREAM").glob("*.m2ts"))
    return max(streams, key=lambda f: f.stat().st_size) if streams else None


def bdmv_task_for(disc):
    """给蓝光原盘生成压缩任务 (只压正片, 输出单个 mkv); 无流/已存在则 None。"""
    main = main_title_m2ts(disc)
    if not main:
        return None
    disc_rel = rel_under_download(disc)
    out = Path(OUTPUT_DIR) / disc_rel / (Path(disc).name + ".mkv")
    if out.exists():
        return None
    return {
        "src": str(main),
        "name": Path(disc).name + ".mkv",
        "torrent_name": disc_rel.parts[0] if disc_rel.parts else Path(disc).name,
        "out": str(out),
        "status": "queued",
        "added_at": datetime.now().isoformat(),
    }


def mirror_non_videos(root):
    """把 root 下所有非视频文件原样复制到输出目录, 保留结构, 已存在则跳过。"""
    p = Path(root)
    files = [p] if p.is_file() else p.rglob("*")
    for f in files:
        try:
            if not f.is_file():
                continue
            low = f.name.lower()
            if f.suffix.lower() in VIDEO_EXTS:
                continue  # 视频交给压缩队列处理
            if low.endswith(".!qb") or low.endswith(".part"):
                continue  # 未下完的临时文件, 跳过
            if "BDMV" in f.parts or "CERTIFICATE" in f.parts:
                continue  # 蓝光原盘目录, 由正片 mkv 替代, 不复制光盘结构
            dst = Path(OUTPUT_DIR) / rel_under_download(f)
            if dst.exists():
                continue
            dst.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(f, dst)
        except Exception as e:
            log.warning("mirror skip %s: %s", f, e)


def mirror_async(root):
    """后台镜像, 不阻塞启动/请求。"""
    threading.Thread(target=mirror_non_videos, args=(root,), daemon=True).start()


def copy_original(src, task):
    dst = output_path(src, task)
    dst.parent.mkdir(parents=True, exist_ok=True)
    if not dst.exists():
        shutil.copy2(src, dst)
    task["src_size"] = Path(src).stat().st_size
    task["dst_size"] = task["src_size"]


def compress_video(src, task):
    dst = output_path(src, task)
    dst.parent.mkdir(parents=True, exist_ok=True)

    if dst.exists():
        task["status"] = "skipped"
        task["reason"] = "输出文件已存在"
        return

    src_size = Path(src).stat().st_size
    task["src_size"] = src_size
    task["status"] = "compressing"
    save_tasks()

    tmp = dst.with_suffix(".tmp" + dst.suffix)
    duration = task.get("duration", 0)

    if is_archive(src):
        # 归档: CPU SVT-AV1 慢速高压, 逼近奈飞式压缩率
        in_opts = []
        v_opts = ["-c:v", "libsvtav1", "-crf", str(ARCHIVE_CRF), "-preset", ARCHIVE_PRESET]
        task["enc"] = "av1"
    elif ENCODER == "hevc_nvenc":
        # 全 GPU 管线: NVDEC 解码 + 帧留显存 + NVENC 编码, 不占 CPU
        # ponytail: nvenc uses -cq not -crf, -rc vbr for quality mode
        in_opts = ["-hwaccel", "cuda", "-hwaccel_output_format", "cuda"]
        v_opts = nvenc_video_options()
        task["enc"] = "nvenc"
    else:
        in_opts = []
        v_opts = ["-c:v", ENCODER, "-crf", str(CRF), "-preset", PRESET, "-threads", "2"]
        task["enc"] = "x265"

    if Path(src).suffix.lower() == ".m2ts":
        # 原盘: 正片视频流 + 全部音轨 + 全部字幕, 丢掉数据/菜单流 (否则 mkv 封装易失败)
        map_opts = ["-map", "0:v:0", "-map", "0:a?", "-map", "0:s?"]
    else:
        map_opts = ["-map", "0"]

    cmd = [
        "nice", "-n", "19",
        FFMPEG, "-y", *in_opts, "-i", str(src),
        *v_opts,
        "-c:a", "copy", "-c:s", "copy", *map_opts,
        "-progress", "pipe:1",
        str(tmp),
    ]

    log.info("Starting: %s", Path(src).name)

    # ponytail: stderr to DEVNULL prevents pipe buffer deadlock
    proc = subprocess.Popen(
        cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True
    )

    for line in proc.stdout:
        m = re.match(r"out_time_us=(\d+)", line)
        if m and duration > 0:
            progress = int(m.group(1)) / 1_000_000 / duration
            task["progress"] = min(round(progress * 100, 1), 100)

    proc.wait()

    if proc.returncode != 0:
        task["status"] = "failed"
        task["reason"] = f"ffmpeg exit code {proc.returncode}"
        log.error("Failed: %s (exit %d)", Path(src).name, proc.returncode)
        tmp.unlink(missing_ok=True)
        return

    dst_size = tmp.stat().st_size
    if dst_size >= src_size:
        tmp.unlink()
        copy_original(src, task)
        task["status"] = "copied"
        task["reason"] = f"压缩后更大({dst_size / src_size:.0%})"
        log.info("Copied (larger after compress): %s", Path(src).name)
        return

    tmp.rename(dst)
    task["dst_size"] = dst_size
    task["ratio"] = round((1 - dst_size / src_size) * 100, 1)
    task["status"] = "done"
    log.info("Done: %s (-%s%%)", Path(src).name, task["ratio"])


def worker():
    log.info("Worker started")
    while True:
        idx = queue.get()
        try:
            with tasks_lock:
                task = tasks[idx]

            src = Path(task["src"])
            log.info("Processing [%d]: %s", idx, src.name)

            if not src.exists():
                task["status"] = "failed"
                task["reason"] = "源文件不存在"
                save_tasks()
                continue

            worth, reason, info = should_compress(src)
            if not worth:
                task["status"] = "copied"
                task["reason"] = reason
                copy_original(str(src), task)
                log.info("Copied: %s (%s)", src.name, reason)
                save_tasks()
                continue

            task["codec"] = reason
            task["duration"] = get_duration_secs(info) if info else 0
            task["src_size"] = src.stat().st_size
            compress_video(str(src), task)
            task["finished_at"] = datetime.now().isoformat()
            save_tasks()
        except Exception as e:
            log.exception("Worker error on task %d", idx)
            try:
                task["status"] = "failed"
                task["reason"] = str(e)[:300]
                save_tasks()
            except Exception:
                pass
        finally:
            queue.task_done()


def scan_download_dir():
    known = {t["src"] for t in tasks}
    videos = find_videos(DOWNLOAD_DIR)
    discs = find_bdmv_discs(DOWNLOAD_DIR)
    added = []
    with tasks_lock:
        for v in videos:
            sv = str(v)
            if sv in known:
                continue
            rel = v.relative_to(DOWNLOAD_DIR)
            torrent_name = rel.parts[0] if len(rel.parts) > 1 else ""
            # 输出已存在 = 压过了, 跳过 (即使任务记录丢失也不会重跑)
            if (Path(OUTPUT_DIR) / rel).exists():
                continue
            task = {
                "src": sv,
                "name": v.name,
                "torrent_name": torrent_name,
                "status": "queued",
                "added_at": datetime.now().isoformat(),
            }
            tasks.append(task)
            queue.put(len(tasks) - 1)
            added.append(v.name)
        # 蓝光原盘: 只压正片, 输出单个 mkv
        for disc in discs:
            t = bdmv_task_for(disc)
            if not t or t["src"] in known:
                continue
            tasks.append(t)
            queue.put(len(tasks) - 1)
            added.append(t["name"])
    if added:
        save_tasks()
    # 视频入队后, 后台镜像非视频文件 (复制慢, 不阻塞启动/请求)
    mirror_async(DOWNLOAD_DIR)
    return added


@app.route("/api/retry", methods=["POST"])
def api_retry():
    count = 0
    with tasks_lock:
        for i, t in enumerate(tasks):
            if t["status"] == "failed":
                t["status"] = "queued"
                t.pop("reason", None)
                t.pop("progress", None)
                queue.put(i)
                count += 1
    if count:
        save_tasks()
    return jsonify({"retried": count})


@app.route("/api/scan", methods=["POST"])
def api_scan():
    added = scan_download_dir()
    return jsonify({"scanned": str(DOWNLOAD_DIR), "added": added, "count": len(added)})


@app.route("/api/add", methods=["POST"])
def add_task():
    path = request.form.get("path") or request.json.get("path", "")
    name = request.form.get("name", "") or request.json.get("name", "")
    if not path:
        return jsonify({"error": "missing path"}), 400

    videos = find_videos(path)
    discs = find_bdmv_discs(path)
    if not videos and not discs:
        return jsonify({"msg": "no video files found"}), 200

    known = {t["src"] for t in tasks}
    added = []
    with tasks_lock:
        for disc in discs:
            t = bdmv_task_for(disc)
            if not t or t["src"] in known:
                continue
            tasks.append(t)
            queue.put(len(tasks) - 1)
            added.append(t["name"])
        for v in videos:
            if str(v) in known:
                continue
            if (Path(OUTPUT_DIR) / rel_under_download(v)).exists():
                continue
            task = {
                "src": str(v),
                "name": v.name,
                "torrent_name": name,
                "status": "queued",
                "added_at": datetime.now().isoformat(),
            }
            tasks.append(task)
            queue.put(len(tasks) - 1)
            added.append(v.name)
    if added:
        save_tasks()
    # 把这个种子的非视频文件也镜像过去 (后台)
    mirror_async(path)
    return jsonify({"added": added}), 200


@app.route("/api/tasks")
def get_tasks():
    with tasks_lock:
        return jsonify(tasks)


@app.route("/api/stats")
def get_stats():
    with tasks_lock:
        done = [t for t in tasks if t["status"] == "done"]
        total_src = sum(t.get("src_size", 0) for t in done)
        total_dst = sum(t.get("dst_size", 0) for t in done)
        return jsonify({
            "total": len(tasks),
            "done": len(done),
            "copied": sum(1 for t in tasks if t["status"] == "copied"),
            "failed": sum(1 for t in tasks if t["status"] == "failed"),
            "queued": sum(1 for t in tasks if t["status"] == "queued"),
            "compressing": sum(1 for t in tasks if t["status"] == "compressing"),
            "total_saved": format_size(total_src - total_dst) if done else "0B",
            "total_src": format_size(total_src),
            "avg_ratio": round((1 - total_dst / total_src) * 100, 1) if total_src else 0,
        })


@app.route("/api/log")
def get_log():
    try:
        with open(LOG_FILE, encoding="utf-8") as f:
            lines = f.readlines()
        return "<br>".join(lines[-100:])
    except Exception:
        return "no log"


@app.route("/")
def index():
    return HTML_PAGE


HTML_PAGE = r"""<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1">
<title>VideoResize</title>
<style>
:root { --bg:#0f1117; --card:#1a1d27; --border:#2a2d3a; --text:#e1e4ed; --dim:#6b7194; --green:#34d399; --red:#f87171; --yellow:#fbbf24; --blue:#60a5fa; }
* { margin:0; padding:0; box-sizing:border-box; }
body { font-family:-apple-system,system-ui,sans-serif; background:var(--bg); color:var(--text); padding:20px; }
h1 { font-size:1.4em; margin-bottom:16px; }
.stats { display:flex; gap:12px; flex-wrap:wrap; margin-bottom:20px; }
.stat { background:var(--card); border:1px solid var(--border); border-radius:8px; padding:12px 18px; min-width:120px; }
.stat .val { font-size:1.5em; font-weight:700; }
.stat .label { font-size:.8em; color:var(--dim); margin-top:2px; }
table { width:100%; border-collapse:collapse; background:var(--card); border-radius:8px; overflow:hidden; }
th,td { padding:10px 14px; text-align:left; border-bottom:1px solid var(--border); font-size:.85em; }
th { color:var(--dim); font-weight:500; font-size:.75em; text-transform:uppercase; letter-spacing:.5px; }
tr:last-child td { border-bottom:none; }
.s-done { color:var(--green); } .s-failed { color:var(--red); } .s-copied { color:var(--dim); }
.s-compressing { color:var(--blue); } .s-queued { color:var(--yellow); }
.bar { width:60px; height:6px; background:var(--border); border-radius:3px; display:inline-block; vertical-align:middle; margin-left:6px; }
.bar-fill { height:100%; background:var(--blue); border-radius:3px; transition:width .3s; }
.ratio { color:var(--green); font-weight:600; }
.empty { text-align:center; padding:40px; color:var(--dim); }
@media(max-width:700px) { th:nth-child(3),td:nth-child(3),th:nth-child(4),td:nth-child(4) { display:none; } }
</style>
</head>
<body>
<h1>VideoResize <button onclick="doScan()" style="font-size:.5em;padding:4px 12px;border-radius:4px;border:1px solid var(--border);background:var(--card);color:var(--text);cursor:pointer;vertical-align:middle">扫描目录</button> <button onclick="doRetry()" style="font-size:.5em;padding:4px 12px;border-radius:4px;border:1px solid var(--border);background:var(--card);color:var(--text);cursor:pointer;vertical-align:middle">重试失败</button></h1>
<div class="stats" id="stats"></div>
<table>
<thead><tr><th>文件</th><th>状态</th><th>原始编码</th><th>大小</th><th>压缩率</th></tr></thead>
<tbody id="tbody"><tr><td colspan="5" class="empty">暂无任务</td></tr></tbody>
</table>
<script>
function fmt(s){return s||'-'}
function statusHtml(t){
  let cls='s-'+t.status, txt=t.status;
  if(t.status==='done') txt='完成';
  else if(t.status==='copied') txt='原样复制: '+(t.reason||'');
  else if(t.status==='failed') txt='失败';
  else if(t.status==='queued') txt='排队中';
  else if(t.status==='compressing'){
    let p=t.progress||0;
    txt='压缩中 '+p.toFixed(1)+'%';
    txt+=' <div class="bar"><div class="bar-fill" style="width:'+p+'%"></div></div>';
  }
  return '<span class="'+cls+'">'+txt+'</span>';
}
function sizeCell(t){
  if(!t.src_size) return '-';
  let s=(t.src_size/1048576).toFixed(0)+'MB';
  if(t.dst_size) s+=' → '+(t.dst_size/1048576).toFixed(0)+'MB';
  return s;
}
function ratioCell(t){
  if(t.status==='done'&&t.ratio!=null) return '<span class="ratio">-'+t.ratio+'%</span>';
  return '-';
}
async function refresh(){
  try{
    let [st,ts]=await Promise.all([fetch('/api/stats').then(r=>r.json()),fetch('/api/tasks').then(r=>r.json())]);
    document.getElementById('stats').innerHTML=
      `<div class="stat"><div class="val">${st.done}</div><div class="label">已完成</div></div>`+
      `<div class="stat"><div class="val">${st.compressing+st.queued}</div><div class="label">队列中</div></div>`+
      `<div class="stat"><div class="val">${st.copied}</div><div class="label">原样复制</div></div>`+
      `<div class="stat"><div class="val">${st.total_saved}</div><div class="label">总节省</div></div>`+
      `<div class="stat"><div class="val">${st.avg_ratio}%</div><div class="label">平均压缩率</div></div>`;
    let tb=document.getElementById('tbody');
    if(!ts.length){tb.innerHTML='<tr><td colspan="5" class="empty">暂无任务</td></tr>';return;}
    tb.innerHTML=ts.slice().reverse().map(t=>
      `<tr><td title="${t.src}">${t.name}</td><td>${statusHtml(t)}</td><td>${fmt(t.codec)}</td><td>${sizeCell(t)}</td><td>${ratioCell(t)}</td></tr>`
    ).join('');
  }catch(e){}
}
async function doRetry(){
  let r=await fetch('/api/retry',{method:'POST'});
  let d=await r.json();
  alert('已重新入队 '+d.retried+' 个失败任务');
  refresh();
}
async function doScan(){
  let r=await fetch('/api/scan',{method:'POST'});
  let d=await r.json();
  alert('扫描完成，新增 '+d.count+' 个文件');
  refresh();
}
refresh();
setInterval(refresh,2000);
</script>
</body>
</html>"""


if __name__ == "__main__":
    Path(OUTPUT_DIR).mkdir(parents=True, exist_ok=True)
    load_tasks()
    added = scan_download_dir()
    if added:
        log.info("Startup scan: found %d new videos", len(added))
    for _ in range(WORKERS):
        threading.Thread(target=worker, daemon=True).start()
    log.info("Server starting on port %d (workers=%d)", PORT, WORKERS)
    app.run(host="0.0.0.0", port=PORT)