import os
import json
import shutil
import tempfile
import subprocess
import time
import re
import uuid
import threading
from concurrent.futures import ThreadPoolExecutor

import gradio as gr

# ═══════════════════════════════════════════════
# 环境变量配置
# ═══════════════════════════════════════════════
HF_TOKEN       = os.environ.get("HF_TOKEN", "")
DATASET_ID     = os.environ.get("DATASET_ID", "")
MS_TOKEN       = os.environ.get("MS_TOKEN", "")
MS_DATASET_ID  = os.environ.get("MS_DATASET_ID", "")
ACCESS_PASSWORD = os.environ.get("ACCESS_PASSWORD", "")

RESULT_KEEP_HOURS = float(os.environ.get("RESULT_KEEP_HOURS", "24"))   # 成品保留小时
STATE_KEEP_DAYS   = float(os.environ.get("STATE_KEEP_DAYS", "7"))      # 状态保留天数


MAX_WORKERS    = int(os.environ.get("MAX_WORKERS", "2"))
UPLOAD_RETRY   = int(os.environ.get("UPLOAD_RETRY", "3"))
FFMPEG_TIMEOUT = int(os.environ.get("FFMPEG_TIMEOUT", "1800"))
MIN_FREE_GB    = float(os.environ.get("MIN_FREE_GB", "2"))

# SRT 烧录时强制套用的样式（对应你的 Default 样式）
# SRT 烧录时强制套用的样式（对应你的 Default 样式）
SRT_FORCE_STYLE = os.environ.get(
    "SRT_FORCE_STYLE",
    "FontName=微软雅黑,FontSize=64,Bold=1,"
    "PrimaryColour=&H0008d76c,OutlineColour=&H00000000,"
    "BorderStyle=1,Outline=2.2,Shadow=0,Alignment=2,"
    "MarginV=8,Spacing=3.4"
)


FONTS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "fonts")

BASE_DIR = "/data" if os.path.isdir("/data") and os.access("/data", os.W_OK) \
    else tempfile.gettempdir()
JOBS_DIR  = os.path.join(BASE_DIR, "jobs")
STATE_DIR = os.path.join(BASE_DIR, "job_state")
os.makedirs(JOBS_DIR, exist_ok=True)
os.makedirs(STATE_DIR, exist_ok=True)

VIDEO_EXTS = {".mp4", ".avi", ".mov", ".mkv", ".wmv", ".flv",
              ".webm", ".m4v", ".3gp", ".ts", ".rmvb"}
SUB_EXTS = {".ass", ".ssa", ".srt"}

# ═══════════════════════════════════════════════
# 任务表 + 常驻线程池
# ═══════════════════════════════════════════════
EXECUTOR = ThreadPoolExecutor(max_workers=MAX_WORKERS)
TASKS = {}
_tasks_lock = threading.Lock()
_ms_login_lock = threading.Lock()
_ms_logged_in = {"done": False}


def _state_path(tid):
    return os.path.join(STATE_DIR, f"{tid}.json")


def _save_state(tid):
    try:
        with _tasks_lock:
            data = dict(TASKS.get(tid, {}))
        with open(_state_path(tid), "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False)
    except Exception:
        pass


def _load_states():
    for fn in os.listdir(STATE_DIR):
        if not fn.endswith(".json"):
            continue
        try:
            with open(os.path.join(STATE_DIR, fn), encoding="utf-8") as f:
                d = json.load(f)
            if d.get("status") in ("queued", "running"):
                d["status"] = "interrupted"
                d["logs"] = d.get("logs", []) + ["⚠ Space 重启，任务被中断"]
            TASKS[d["id"]] = d
        except Exception:
            pass


def new_task():
    tid = uuid.uuid4().hex[:12]
    with _tasks_lock:
        TASKS[tid] = {
            "id": tid, "status": "queued", "logs": [],
            "total": 0, "done": 0, "success": 0,
            "outputs": [], "created": time.time(),
        }
    _save_state(tid)
    return tid


def t_log(tid, msg, save=True):
    ts = time.strftime("%H:%M:%S")
    with _tasks_lock:
        if tid in TASKS:
            TASKS[tid]["logs"].append(f"[{ts}] {msg}")
            if len(TASKS[tid]["logs"]) > 1000:
                TASKS[tid]["logs"] = TASKS[tid]["logs"][-800:]
    if save:
        _save_state(tid)


def t_set(tid, **kw):
    with _tasks_lock:
        if tid in TASKS:
            TASKS[tid].update(kw)
    _save_state(tid)


def t_get(tid):
    with _tasks_lock:
        return dict(TASKS[tid]) if tid in TASKS else None


# ═══════════════════════════════════════════════
# 工具函数
# ═══════════════════════════════════════════════
def sanitize_name(name, max_len=80):
    if not name:
        return ""
    name = os.path.basename(name)
    name = re.sub(r'[\\/:*?"<>|\r\n\t]', "_", name)
    name = name.strip().strip(".").strip()
    name = re.sub(r"\s+", " ", name)
    return name[:max_len].rstrip() if len(name) > max_len else name

def strip_english_from_srt(src_path, dst_path):
    """读取 SRT，每条字幕只保留含中文的行，去掉纯英文行"""
    def has_chinese(s):
        return any('\u4e00' <= ch <= '\u9fff' for ch in s)

    with open(src_path, "r", encoding="utf-8-sig", errors="ignore") as f:
        content = f.read()

    # 按空行分块（每条字幕一块）
    blocks = re.split(r"\n\s*\n", content.strip())
    out_blocks = []
    for block in blocks:
        lines = block.splitlines()
        if len(lines) < 2:
            continue
        idx = lines[0].strip()          # 序号
        timeline = lines[1].strip()     # 时间轴
        text_lines = lines[2:]          # 字幕文字（可能多行）
        # 只保留含中文的文字行
        kept = [ln for ln in text_lines if has_chinese(ln)]
        if not kept:
            continue                    # 整条都没中文则丢弃
        out_blocks.append(f"{idx}\n{timeline}\n" + "\n".join(kept))

    with open(dst_path, "w", encoding="utf-8") as f:
        f.write("\n\n".join(out_blocks) + "\n")

def cleanup_once():
    """清理过期的成品目录与状态文件"""
    now = time.time()
    # 清理过期成品目录 JOBS_DIR/<tid>/
    try:
        for tid in os.listdir(JOBS_DIR):
            d = os.path.join(JOBS_DIR, tid)
            if not os.path.isdir(d):
                continue
            age_h = (now - os.path.getmtime(d)) / 3600
            if age_h > RESULT_KEEP_HOURS:
                shutil.rmtree(d, ignore_errors=True)
                # 内存里把 outputs 清掉，避免查询页给出失效路径
                with _tasks_lock:
                    if tid in TASKS:
                        TASKS[tid]["outputs"] = []
    except Exception:
        pass
    # 清理过期状态文件
    try:
        for fn in os.listdir(STATE_DIR):
            if not fn.endswith(".json"):
                continue
            p = os.path.join(STATE_DIR, fn)
            age_d = (now - os.path.getmtime(p)) / 86400
            if age_d > STATE_KEEP_DAYS:
                os.remove(p)
                with _tasks_lock:
                    TASKS.pop(fn[:-5], None)
    except Exception:
        pass


def cleanup_loop():
    while True:
        time.sleep(3600)      # 每小时清理一次
        cleanup_once()


def start_cleanup():
    cleanup_once()            # 启动时先清一次
    t = threading.Thread(target=cleanup_loop, daemon=True)
    t.start()



def free_space_gb(path=None):
    try:
        return shutil.disk_usage(path or BASE_DIR)[2] / (1024 ** 3)
    except Exception:
        return 999


def _run_ffmpeg_live(cmd, tid, tag, output_path, action):
    t_log(tid, f"[{tag}] {action}…")
    try:
        proc = subprocess.Popen(
            cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
            universal_newlines=True, bufsize=1,
        )
    except Exception as e:
        t_log(tid, f"[{tag}] ❌ 启动 ffmpeg 失败: {e}")
        return False

    start = time.time()
    last_log = 0
    try:
        for line in proc.stdout:
            line = line.strip()
            if line and any(k in line for k in ("frame=", "time=", "size=")):
                now = time.time()
                if now - last_log >= 2:
                    t_log(tid, f"[{tag}] {line}", save=False)
                    last_log = now
            elif line and any(k in line.lower() for k in ("error", "invalid", "failed")):
                t_log(tid, f"[{tag}] {line}", save=False)
            if time.time() - start > FFMPEG_TIMEOUT:
                proc.kill()
                t_log(tid, f"[{tag}] ❌ {action}超时，已终止")
                return False
        proc.wait()
    except Exception as e:
        t_log(tid, f"[{tag}] ❌ {action}异常: {e}")
        return False
    finally:
        _save_state(tid)

    if proc.returncode != 0:
        t_log(tid, f"[{tag}] ❌ {action}失败（码 {proc.returncode}）")
        return False
    if not (os.path.exists(output_path) and os.path.getsize(output_path) > 0):
        t_log(tid, f"[{tag}] ❌ {action}产物为空")
        return False
    t_log(tid, f"[{tag}] ✅ {action}完成 ({os.path.getsize(output_path)/1048576:.1f} MB)")
    return True


def burn_subtitle(video_path, sub_path, output_path, fast_mode, tid, tag):
    ext = os.path.splitext(sub_path)[1].lower()
    if ext == ".srt":
        vf = (f"subtitles='{sub_path}':fontsdir='{FONTS_DIR}'"
              f":force_style='{SRT_FORCE_STYLE}'")
    else:
        vf = f"subtitles='{sub_path}':fontsdir='{FONTS_DIR}'"

    preset, crf = ("veryfast", "23") if fast_mode else ("medium", "20")
    cmd = ["ffmpeg", "-y", "-i", video_path, "-vf", vf,
           "-c:v", "libx264", "-preset", preset, "-crf", crf,
           "-c:a", "copy", output_path]
    return _run_ffmpeg_live(
        cmd, tid, tag, output_path,
        f"烧录字幕（{'快速' if fast_mode else '标准'}模式）")


def convert_to_amv(input_path, output_path, tid, tag):
    cmd = ["ffmpeg", "-y", "-i", input_path,
           "-s", "320x240", "-r", "15",
           "-c:v", "amv", "-c:a", "adpcm_ima_amv",
           "-ar", "22050", "-ac", "1", "-block_size", "1470",
           output_path]
    return _run_ffmpeg_live(cmd, tid, tag, output_path, "转换 AMV")


def _retry(fn, tid, tag, what):
    for attempt in range(1, UPLOAD_RETRY + 1):
        try:
            fn()
            return True
        except Exception as e:
            if attempt < UPLOAD_RETRY:
                wait = 2 ** attempt
                t_log(tid, f"[{tag}] ⚠ {what}第{attempt}次失败: {e}，{wait}s后重试")
                time.sleep(wait)
            else:
                t_log(tid, f"[{tag}] ❌ {what}最终失败: {e}")
    return False


def upload_to_hf(fp, repo_path, tid, tag):
    if not (HF_TOKEN and DATASET_ID):
        t_log(tid, f"[{tag}] ⏭ 未配置 HF，跳过")
        return
    def _do():
        from huggingface_hub import HfApi
        HfApi(token=HF_TOKEN).upload_file(
            path_or_fileobj=fp, path_in_repo=repo_path,
            repo_id=DATASET_ID, repo_type="dataset")
    if _retry(_do, tid, tag, f"HF上传({os.path.basename(repo_path)})"):
        t_log(tid, f"[{tag}] ✅ 已上传 HF: {repo_path}")


def upload_to_modelscope(fp, repo_path, tid, tag):
    if not (MS_TOKEN and MS_DATASET_ID):
        t_log(tid, f"[{tag}] ⏭ 未配置魔搭，跳过")
        return
    def _do():
        from modelscope.hub.api import HubApi
        api = HubApi()
        with _ms_login_lock:
            if not _ms_logged_in["done"]:
                api.login(MS_TOKEN)
                _ms_logged_in["done"] = True
        api.upload_file(path_or_fileobj=fp, path_in_repo=repo_path,
                        repo_id=MS_DATASET_ID, repo_type="dataset",
                        commit_message=f"上传: {repo_path}")
    if _retry(_do, tid, tag, f"魔搭上传({os.path.basename(repo_path)})"):
        t_log(tid, f"[{tag}] ✅ 已上传魔搭: {repo_path}")


def pair_files(files):
    videos, subs = {}, {}
    for f in files:
        ext = os.path.splitext(f)[1].lower()
        stem = os.path.splitext(os.path.basename(f))[0]
        if ext in VIDEO_EXTS:
            videos[stem] = f
        elif ext in SUB_EXTS:
            subs[stem] = f
    pairs, unmatched = [], []
    for stem, v in videos.items():
        if stem in subs:
            pairs.append((stem, v, subs[stem]))
        else:
            unmatched.append(f"视频「{stem}」缺少同名字幕")
    for stem in subs:
        if stem not in videos:
            unmatched.append(f"字幕「{stem}」缺少同名视频")
    return pairs, unmatched


def handle_one(stem, video, subtitle, fast_mode, folder, date_dir, tid, out_dir):
    tag = stem if len(stem) <= 20 else stem[:18] + "…"
    if free_space_gb() < MIN_FREE_GB:
        t_log(tid, f"[{tag}] ❌ 磁盘不足，跳过")
        return None
    work = tempfile.mkdtemp()
    results = []
    try:
        lv = os.path.join(work, "input" + os.path.splitext(video)[1])
        shutil.copy(video, lv)
        sub_ext = os.path.splitext(subtitle)[1].lower()
        ls = os.path.join(work, "sub" + sub_ext)
        if sub_ext == ".srt":
            # SRT：去掉英文行，只留中文，再烧录
            strip_english_from_srt(subtitle, ls)
            t_log(tid, f"[{tag}] 已清洗 SRT，仅保留中文行")
        else:
            shutil.copy(subtitle, ls)

        clean = sanitize_name(stem)
        mp4_name = f"{clean}_subbed.mp4"
        amv_name = f"{clean}_subbed.amv"
        mp4_tmp = os.path.join(work, mp4_name)
        amv_tmp = os.path.join(work, amv_name)

        # ① 烧录字幕 -> MP4
        if not burn_subtitle(lv, ls, mp4_tmp, fast_mode, tid, tag):
            return results

        mp4_repo = f"{folder}/mp4/{date_dir}/{mp4_name}"
        upload_to_hf(mp4_tmp, mp4_repo, tid, tag)
        upload_to_modelscope(mp4_tmp, mp4_repo, tid, tag)
        mp4_keep = os.path.join(out_dir, mp4_name)
        shutil.copy(mp4_tmp, mp4_keep)
        results.append(mp4_keep)

        # ② MP4 -> AMV
        if convert_to_amv(mp4_tmp, amv_tmp, tid, tag):
            amv_repo = f"{folder}/amv/{date_dir}/{amv_name}"
            upload_to_hf(amv_tmp, amv_repo, tid, tag)
            upload_to_modelscope(amv_tmp, amv_repo, tid, tag)
            amv_keep = os.path.join(out_dir, amv_name)
            shutil.copy(amv_tmp, amv_keep)
            results.append(amv_keep)

        return results
    except Exception as e:
        t_log(tid, f"[{tag}] ❌ 处理异常: {e}")
        return results
    finally:
        shutil.rmtree(work, ignore_errors=True)


def run_job(tid, saved_files, fast_mode, folder):
    try:
        t_set(tid, status="running")
        folder = sanitize_name(folder) or "video-srt"
        pairs, unmatched = pair_files(saved_files)
        for u in unmatched:
            t_log(tid, f"⚠ {u}")
        if not pairs:
            t_log(tid, "❌ 无同名配对，任务结束")
            t_set(tid, status="error")
            return

        out_dir = os.path.join(JOBS_DIR, tid)
        os.makedirs(out_dir, exist_ok=True)
        date_dir = time.strftime("%Y%m%d")
        t_set(tid, total=len(pairs))
        t_log(tid, f"共 {len(pairs)} 组 | 目标 {folder}/(mp4,amv) | 磁盘 {free_space_gb():.1f}GB")

        inner = ThreadPoolExecutor(max_workers=MAX_WORKERS)
        futs = [inner.submit(handle_one, stem, v, s, fast_mode,
                             folder, date_dir, tid, out_dir)
                for stem, v, s in pairs]
        for fut in futs:
            files = fut.result() or []
            with _tasks_lock:
                TASKS[tid]["done"] += 1
                if files:
                    TASKS[tid]["success"] += 1
                    TASKS[tid]["outputs"].extend(files)
            _save_state(tid)
        inner.shutdown(wait=True)

        with _tasks_lock:
            s = TASKS[tid]["success"]
        t_log(tid, f"✅ 全部完成，成功 {s}/{len(pairs)}")
        t_set(tid, status="done")
    except Exception as e:
        t_log(tid, f"❌ 任务异常: {e}")
        t_set(tid, status="error")
    finally:
        # 删除整个上传临时目录（含空壳）
        dirs = {os.path.dirname(f) for f in saved_files}
        for d in dirs:
            shutil.rmtree(d, ignore_errors=True)



# ═══════════════════════════════════════════════
# 前端回调
# ═══════════════════════════════════════════════
def submit(files, fast_mode, folder):
    if not files:
        raise gr.Error("请上传视频和字幕文件")
    persist = tempfile.mkdtemp(prefix="upload_", dir=BASE_DIR)
    saved = []
    for f in files:
        dst = os.path.join(persist, os.path.basename(f))
        shutil.copy(f, dst)
        saved.append(dst)
    tid = new_task()
    EXECUTOR.submit(run_job, tid, saved, fast_mode, folder)
    return (tid,
            f"✅ 任务已提交！任务ID：{tid}\n"
            f"可关闭页面，后台继续执行。到「查询进度」页粘贴此 ID 并勾选自动刷新即可看实时日志。")


def query(tid):
    tid = (tid or "").strip()
    if not tid:
        return "请输入任务ID", None
    t = t_get(tid)
    if not t:
        return f"未找到任务 {tid}", None
    status_map = {"queued": "排队中", "running": "处理中", "done": "已完成",
                  "error": "出错", "interrupted": "被中断(Space重启)"}
    head = (f"任务 {tid} | 状态：{status_map.get(t['status'], t['status'])} | "
            f"进度 {t['done']}/{t['total']}（成功 {t['success']}）\n"
            + "─" * 40 + "\n")
    body = "\n".join(t["logs"][-300:])
    outs = None
    if t["status"] in ("done", "running"):
        outs = [p for p in t["outputs"] if os.path.exists(p)] or None
    return head + body, outs


_load_states()
start_cleanup()          # 启动后台定期清理

# ═══════════════════════════════════════════════
# 界面
# ═══════════════════════════════════════════════
CSS = """
.gradio-container {max-width: 920px !important; margin: auto;}
#title-box {text-align:center; padding:18px;
  background:linear-gradient(135deg,#6366f1,#8b5cf6);
  border-radius:14px; color:#fff; margin-bottom:14px;}
#title-box h1 {margin:0; font-size:26px;}
#title-box p {margin:6px 0 0; opacity:.92; font-size:14px;}
"""

with gr.Blocks(title="ASS/SRT 硬字幕批量烧录", theme=gr.themes.Soft(
        primary_hue="indigo", secondary_hue="violet"), css=CSS) as demo:

    gr.HTML('<div id="title-box"><h1>🎬 ASS / SRT 硬字幕批量烧录工具</h1>'
            '<p>实时日志 · 后台任务 · MP4+AMV 双格式 · 多线程 · 失败重试 · 上传 HF / 魔搭</p></div>')

    with gr.Tab("📤 提交任务"):
        files_in = gr.File(label="📁 视频 + 字幕(.srt/.ass)（可多选，同名自动配对）",
                           file_count="multiple")
        with gr.Row():
            folder_in = gr.Textbox(label="📂 数据集文件夹名", value="video-srt", scale=3)
            fast_in = gr.Checkbox(label="⚡ 快速模式", value=False, scale=1)
        gr.Markdown(
            "> SRT 会统一套用预设样式（微软雅黑/亮黄绿/大字）；ASS 用文件自带样式。\n"
            "> 产物上传到：`文件夹名/mp4/日期/` 和 `文件夹名/amv/日期/`")
        submit_btn = gr.Button("🚀 提交后台任务", variant="primary")
        tid_box = gr.Textbox(label="任务ID（请复制保存）", interactive=False)
        submit_msg = gr.Textbox(label="提交结果", lines=3, interactive=False)
        submit_btn.click(submit, [files_in, fast_in, folder_in],
                         [tid_box, submit_msg])

    with gr.Tab("🔍 查询进度 / 实时日志"):
        gr.Markdown("输入任务ID，**勾选自动刷新**即可每 3 秒实时更新日志与进度。")
        with gr.Row():
            q_tid = gr.Textbox(label="任务ID", scale=3)
            auto = gr.Checkbox(label="🔄 自动刷新(3s)", value=True, scale=1)
        q_btn = gr.Button("查询", variant="primary")
        q_log = gr.Textbox(label="实时状态与日志", lines=20, show_copy_button=True)
        q_files = gr.File(label="📥 结果文件（MP4 / AMV）", file_count="multiple")
        q_btn.click(query, [q_tid], [q_log, q_files])
        timer = gr.Timer(3, active=True)
        auto.change(lambda a: gr.Timer(active=a), auto, timer)
        timer.tick(query, [q_tid], [q_log, q_files])

# Docker 必需：绑定 0.0.0.0:7860
# 访问认证：设置了密码则启用登录（用户名固定 user，密码取环境变量）
_auth = ("user", ACCESS_PASSWORD) if ACCESS_PASSWORD else None

demo.launch(
    server_name="0.0.0.0",
    server_port=7860,
    auth=_auth,
    auth_message="请输入访问密码（用户名：user）",
)