#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
get.py — Tải phụ đề (transcript) cho danh sách video YouTube trong urls.txt.

Cách dùng (chạy trên máy của bạn, KHÔNG cần cài gì trước):

    Windows :  py -3 get.py
    macOS   :  python3 get.py
    Linux   :  python3 get.py

Lần chạy đầu, script tự tạo môi trường Python riêng trong out/venv và tự cài
yt-dlp + youtube-transcript-api vào đó (không đụng vào Python hệ thống).

Chạy lại bất cứ lúc nào: video nào đã có out/<id>.txt dài hơn 200 ký tự sẽ
được bỏ qua, nên dừng giữa chừng rồi chạy lại chỉ tải phần còn thiếu.

Tuỳ chọn:
    --urls PATH     dùng tệp danh sách khác (mặc định: urls.txt cạnh script)
    --out DIR       thư mục đầu ra (mặc định: out/ cạnh script)
    --limit N       chỉ làm N video đầu tiên (để thử)
    --only ID[,ID]  chỉ làm các id này
    --force         tải lại cả video đã có .txt
    --no-venv       dùng Python hiện tại thay vì tạo out/venv
    --no-zip        không tạo holetex-transcripts.zip ở cuối
"""

import argparse
import html
import json
import os
import re
import subprocess
import sys
import time
import zipfile
from pathlib import Path

HERE = Path(__file__).resolve().parent

# Thứ tự ưu tiên ngôn ngữ phụ đề: vi-orig > vi > en-orig > en
LANG_PRIORITY = ["vi-orig", "vi", "en-orig", "en"]

MIN_CHARS = 200          # .txt ngắn hơn mức này coi như chưa có
YTDLP_TIMEOUT = 300      # giây / video
SUB_LANGS = "vi.*,en.*"

TIMESTAMP_RE = re.compile(r"^\d{1,2}:\d{2}(?::\d{2})?[.,]\d{3}\s*-->")
TAG_RE = re.compile(r"<[^>]*>")
INLINE_TS_RE = re.compile(r"<\d{1,2}:\d{2}(?::\d{2})?[.,]\d{3}>")


# --------------------------------------------------------------------------
# tiện ích
# --------------------------------------------------------------------------
def say(msg=""):
    try:
        print(msg, flush=True)
    except UnicodeEncodeError:                       # cmd.exe cũ / cp1252
        enc = (sys.stdout.encoding or "utf-8")
        print(msg.encode(enc, "replace").decode(enc, "replace"), flush=True)


def venv_python(venv_dir: Path):
    """Dò Python trong venv — Windows: Scripts\\python.exe, *nix: bin/python3."""
    candidates = [
        venv_dir / "Scripts" / "python.exe",   # Windows
        venv_dir / "Scripts" / "python3.exe",
        venv_dir / "bin" / "python3",          # Linux / macOS
        venv_dir / "bin" / "python",
    ]
    for c in candidates:
        if c.exists():
            return c
    return None


def run(cmd, timeout=None, capture=True):
    try:
        return subprocess.run(
            [str(c) for c in cmd],
            stdout=subprocess.PIPE if capture else None,
            stderr=subprocess.STDOUT if capture else None,
            timeout=timeout,
            text=True,
            encoding="utf-8",
            errors="replace",
        )
    except subprocess.TimeoutExpired:
        return None
    except FileNotFoundError as exc:
        say("  ! không chạy được lệnh: %s" % exc)
        return None


def ensure_venv(out_dir: Path, use_venv: bool):
    """Tạo (nếu chưa có) out/venv và trả về đường dẫn python của venv."""
    if not use_venv:
        return Path(sys.executable)

    venv_dir = out_dir / "venv"
    py = venv_python(venv_dir)
    if py:
        return py

    say("· Tạo môi trường Python riêng trong %s ..." % venv_dir)
    res = run([sys.executable, "-m", "venv", str(venv_dir)])
    if res is None or res.returncode != 0:
        detail = (res.stdout or "").strip()[-800:] if res else "timeout"
        say("  ! Không tạo được venv. Chi tiết:\n%s" % detail)
        say("  ! Chạy lại bằng:  python3 get.py --no-venv   (hoặc cài python3-venv)")
        sys.exit(2)

    py = venv_python(venv_dir)
    if py is None:
        say("  ! Tạo venv xong nhưng không tìm thấy python bên trong %s" % venv_dir)
        sys.exit(2)
    return py


def ensure_deps(py: Path):
    """Cài yt-dlp + youtube-transcript-api vào đúng interpreter đang dùng."""
    probe = run([py, "-c", "import yt_dlp, youtube_transcript_api"])
    if probe is not None and probe.returncode == 0:
        return True

    say("· Cài yt-dlp + youtube-transcript-api (lần đầu, ~30 giây) ...")
    res = run([py, "-m", "pip", "install", "--upgrade", "--disable-pip-version-check",
               "--quiet", "yt-dlp", "youtube-transcript-api"], timeout=900)
    if res is None or res.returncode != 0:
        detail = (res.stdout or "").strip()[-1200:] if res else "timeout"
        say("  ! Cài thư viện thất bại. Chi tiết:\n%s" % detail)
        say("  ! Thử:  <python> -m pip install yt-dlp youtube-transcript-api")
        sys.exit(3)

    probe = run([py, "-c", "import yt_dlp, youtube_transcript_api"])
    return probe is not None and probe.returncode == 0


# --------------------------------------------------------------------------
# đọc danh sách
# --------------------------------------------------------------------------
def load_urls(path: Path):
    """Đọc urls.txt dạng  id|tiêu đề|video|short  (tiêu đề có thể chứa '|')."""
    items = []
    with open(path, encoding="utf-8-sig") as fh:
        for lineno, raw in enumerate(fh, 1):
            line = raw.strip()
            if not line or line.startswith("#"):
                continue
            parts = line.split("|")
            if len(parts) < 2:
                say("  ! bỏ qua dòng %d (không đúng định dạng): %s" % (lineno, line[:60]))
                continue
            vid = parts[0].strip()
            kind = parts[-1].strip().lower() if len(parts) >= 3 else "video"
            title = "|".join(parts[1:-1]).strip() if len(parts) >= 3 else parts[1].strip()
            if kind not in ("video", "short"):
                kind = "video"
            if not re.fullmatch(r"[A-Za-z0-9_-]{5,20}", vid):
                say("  ! bỏ qua dòng %d (video id lạ): %s" % (lineno, vid))
                continue
            items.append({"id": vid, "title": title, "type": kind})
    return items


# --------------------------------------------------------------------------
# VTT -> text
# --------------------------------------------------------------------------
def vtt_to_text(raw: str) -> str:
    """Bỏ header/cue number/timestamp/tag, gộp các dòng trùng liên tiếp."""
    lines = []
    seen_cue = False          # mọi thứ trước cue đầu tiên là header (WEBVTT, Kind:, Language:)
    for line in raw.replace("\r\n", "\n").replace("\r", "\n").split("\n"):
        s = line.strip()
        if not s:
            continue
        if not seen_cue:
            if "-->" in s:
                seen_cue = True
            continue
        upper = s.upper()
        if upper.startswith("WEBVTT") or upper.startswith("NOTE") \
                or upper.startswith("STYLE") or upper.startswith("REGION"):
            continue
        if "-->" in s or TIMESTAMP_RE.match(s):
            continue
        if s.isdigit():                     # số thứ tự cue
            continue
        s = INLINE_TS_RE.sub("", s)         # <00:00:01.234>
        s = TAG_RE.sub("", s)               # <c>, </c>, <v Tên>, <b>...
        s = html.unescape(s)
        s = re.sub(r"\s+", " ", s).strip()
        if s:
            lines.append(s)

    out = []
    for s in lines:
        if out:
            prev = out[-1]
            if s == prev:                        # trùng y hệt
                continue
            if prev.endswith(s):                 # phụ đề cuộn: lặp lại đuôi câu
                continue
            if s.startswith(prev):               # phụ đề cuộn: câu dài dần
                out[-1] = s
                continue
        out.append(s)

    text = " ".join(out)
    return re.sub(r"\s+", " ", text).strip()


def pick_vtt(out_dir: Path, vid: str):
    """Chọn tệp .vtt theo đúng thứ tự vi-orig > vi > en-orig > en.

    Không dùng glob-sort: một video có thể sinh <id>.vi-orig.vtt + <id>.vi.vtt
    + <id>.en.vtt cùng lúc, sort theo tên sẽ chọn sai.
    """
    cands = []
    for p in out_dir.glob(vid + ".*.vtt"):
        name = p.name
        prefix = vid + "."
        if not name.startswith(prefix) or not name.endswith(".vtt"):
            continue
        lang = name[len(prefix):-len(".vtt")]
        if lang == "vi-orig":
            rank = 0
        elif lang == "vi":
            rank = 1
        elif lang.startswith("vi"):
            rank = 2
        elif lang == "en-orig":
            rank = 3
        elif lang == "en":
            rank = 4
        elif lang.startswith("en"):
            rank = 5
        else:
            rank = 9
        cands.append((rank, lang, p))
    if not cands:
        return None, None
    cands.sort(key=lambda c: (c[0], c[1]))
    rank, lang, path = cands[0]
    return path, lang


def existing_chars(txt_path: Path) -> int:
    try:
        return len(txt_path.read_text(encoding="utf-8", errors="replace").strip())
    except OSError:
        return 0


# --------------------------------------------------------------------------
# tải 1 video
# --------------------------------------------------------------------------
def ytdlp_subs(py: Path, out_dir: Path, vid: str):
    """Gọi yt-dlp chỉ để lấy phụ đề (KHÔNG tải audio/video)."""
    cmd = [
        py, "-m", "yt_dlp",
        "--skip-download",
        "--write-subs", "--write-auto-subs",
        "--sub-langs", SUB_LANGS,
        "--sub-format", "vtt",
        "--no-playlist",
        "--retries", "2",
        "--extractor-retries", "2",
        "--socket-timeout", "20",
        "--no-mtime",
        "-P", str(out_dir),
        "-o", "%(id)s.%(ext)s",
        "https://www.youtube.com/watch?v=%s" % vid,
    ]
    res = run(cmd, timeout=YTDLP_TIMEOUT)
    if res is None:
        return "timeout", "yt-dlp quá %ds" % YTDLP_TIMEOUT
    blob = (res.stdout or "")
    low = blob.lower()
    if "sign in to confirm" in low or "confirm you’re not a bot" in low \
            or "confirm you're not a bot" in low:
        return "blocked", "YouTube chặn IP (Sign in to confirm you're not a bot)"
    if res.returncode != 0:
        tail = " / ".join([l for l in blob.strip().splitlines() if l.strip()][-3:])
        return "error", tail[:300] or ("yt-dlp exit %d" % res.returncode)
    return "ok", ""


def api_fallback(vid: str):
    """Dự phòng khi yt-dlp không ra .vtt: youtube-transcript-api."""
    try:
        from youtube_transcript_api import YouTubeTranscriptApi  # noqa
    except Exception:
        return None, None
    try:
        api = YouTubeTranscriptApi()
        try:
            fetched = api.fetch(vid, languages=["vi", "en"])
        except Exception:
            lst = api.list(vid)
            fetched = lst.find_transcript(["vi", "en"]).fetch()
        lang = getattr(fetched, "language_code", None) or "unknown"
        chunks = []
        for snip in fetched:
            txt = getattr(snip, "text", None)
            if txt is None and isinstance(snip, dict):
                txt = snip.get("text", "")
            chunks.append((txt or "").replace("\n", " ").strip())
        return lang, chunks
    except Exception:
        return None, None


def process_one(py: Path, out_dir: Path, item, force=False, prev=None):
    vid = item["id"]
    txt_path = out_dir / (vid + ".txt")

    if not force:
        chars = existing_chars(txt_path)
        if chars > MIN_CHARS:
            lang = (prev or {}).get("lang")
            generated = (prev or {}).get("generated")
            if not lang:                              # thử đọc lại từ .vtt còn sót
                _p, detected = pick_vtt(out_dir, vid)
                lang = detected
                generated = bool(lang and lang.endswith("-orig"))
            return {"id": vid, "title": item["title"], "type": item["type"],
                    "lang": lang, "generated": generated, "chars": chars,
                    "status": "skipped"}

    for old in out_dir.glob(vid + ".*.vtt"):
        try:
            old.unlink()
        except OSError:
            pass

    status, note = ytdlp_subs(py, out_dir, vid)
    vtt_path, lang = pick_vtt(out_dir, vid)

    if vtt_path is not None:
        try:
            raw = vtt_path.read_text(encoding="utf-8", errors="replace")
        except OSError as exc:
            return {"id": vid, "title": item["title"], "type": item["type"],
                    "lang": None, "generated": None, "chars": 0,
                    "status": "error", "note": "không đọc được vtt: %s" % exc}
        text = vtt_to_text(raw)
        if len(text) > MIN_CHARS:
            txt_path.write_text(text, encoding="utf-8")
            return {"id": vid, "title": item["title"], "type": item["type"],
                    "lang": lang, "generated": bool(lang and lang.endswith("-orig")),
                    "chars": len(text), "status": "ok"}
        status = "no_subs" if not text else "too_short"
        note = "VTT rỗng/ngắn (%d ký tự)" % len(text)

    # dự phòng bằng youtube-transcript-api
    api_lang, chunks = api_fallback(vid)
    if chunks:
        text = vtt_to_text("\n".join(chunks))
        if len(text) > MIN_CHARS:
            txt_path.write_text(text, encoding="utf-8")
            return {"id": vid, "title": item["title"], "type": item["type"],
                    "lang": api_lang, "generated": bool(api_lang and "orig" in api_lang),
                    "chars": len(text), "status": "ok", "note": "qua youtube-transcript-api"}
        status, note = "no_subs", "API trả về quá ngắn"

    if status not in ("blocked", "timeout"):
        status = "no_subs"
    return {"id": vid, "title": item["title"], "type": item["type"],
            "lang": None, "generated": None, "chars": 0, "status": status, "note": note}


# --------------------------------------------------------------------------
# manifest + zip
# --------------------------------------------------------------------------
def write_manifest(out_dir: Path, entries):
    payload = {
        "channel": "holetex",
        "generated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "total": len(entries),
        "by_status": count_by(entries, "status"),
        "by_lang": count_by([e for e in entries
                             if e.get("status") in ("ok", "skipped") and e.get("lang")], "lang"),
        "videos": entries,
    }
    (out_dir / "MANIFEST.json").write_text(
        json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
    return payload


def count_by(entries, field):
    counts = {}
    for e in entries:
        k = e.get(field) or "không rõ"
        counts[str(k)] = counts.get(str(k), 0) + 1
    return counts


def make_zip(out_dir: Path, zip_path: Path):
    n = 0
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
        for p in sorted(out_dir.iterdir()):
            if p.is_dir() or p.suffix.lower() == ".vtt":
                continue                      # không đóng gói venv/ và .vtt
            if p.name == "MANIFEST.json" or p.suffix.lower() == ".txt":
                zf.write(p, arcname=p.name)
                n += 1
    return n


# --------------------------------------------------------------------------
# main
# --------------------------------------------------------------------------
def build_parser():
    ap = argparse.ArgumentParser(
        prog="get.py",
        description="Tải phụ đề YouTube cho danh sách trong urls.txt (không tải video).")
    ap.add_argument("--urls", default=None, help="tệp danh sách (mặc định urls.txt)")
    ap.add_argument("--out", default=None, help="thư mục đầu ra (mặc định out/)")
    ap.add_argument("--limit", type=int, default=0, help="chỉ làm N video đầu")
    ap.add_argument("--only", default=None, help="chỉ làm các id này, cách nhau dấu phẩy")
    ap.add_argument("--force", action="store_true", help="tải lại cả video đã có .txt")
    ap.add_argument("--no-venv", action="store_true", help="dùng Python hiện tại")
    ap.add_argument("--no-zip", action="store_true", help="không tạo zip cuối")
    return ap


def main(argv=None):
    args = build_parser().parse_args(argv)

    urls_path = Path(args.urls).expanduser() if args.urls else HERE / "urls.txt"
    out_dir = Path(args.out).expanduser() if args.out else HERE / "out"
    out_dir.mkdir(parents=True, exist_ok=True)

    if not urls_path.exists():
        say("! Không thấy tệp danh sách: %s" % urls_path)
        say("  Giải nén đầy đủ gói rồi chạy lại (get.py và urls.txt phải cùng thư mục).")
        return 1

    items = load_urls(urls_path)
    if args.only:
        wanted = [s.strip() for s in args.only.split(",") if s.strip()]
        items = [i for i in items if i["id"] in wanted]
    if args.limit and args.limit > 0:
        items = items[:args.limit]

    say("=" * 68)
    say(" HoleTex transcripts — %d video trong danh sách" % len(items))
    say(" Thư mục ra: %s" % out_dir)
    say("=" * 68)
    if not items:
        say("! Danh sách rỗng, không có gì để làm.")
        return 1

    py = ensure_venv(out_dir, use_venv=not args.no_venv)
    say("· Python dùng để tải: %s" % py)
    if not ensure_deps(py):
        say("! Thiếu thư viện, dừng.")
        return 3

    entries = []
    manifest_path = out_dir / "MANIFEST.json"
    if manifest_path.exists() and not args.force:
        try:
            old = json.loads(manifest_path.read_text(encoding="utf-8"))
            entries = [e for e in old.get("videos", []) if isinstance(e, dict)]
        except Exception:
            entries = []
    done_ids = {e["id"] for e in entries if e.get("status") == "ok"}
    prev_by_id = {e["id"]: e for e in entries if isinstance(e.get("id"), str)}

    started = time.time()
    for idx, item in enumerate(items, 1):
        say("[%d/%d] %s  %s" % (idx, len(items), item["id"], item["title"][:60]))
        try:
            entry = process_one(py, out_dir, item, force=args.force,
                                prev=prev_by_id.get(item["id"]))
        except KeyboardInterrupt:
            say("\n! Bạn đã dừng. Chạy lại lệnh cũ để tải tiếp phần còn thiếu.")
            write_manifest(out_dir, entries)
            return 130

        entries = [e for e in entries if e.get("id") != entry["id"]] + [entry]
        ent = entry
        if ent["status"] == "ok":
            say("        OK  %-8s %6d ký tự%s" % (
                ent.get("lang") or "?", ent.get("chars") or 0,
                "  (auto)" if ent.get("generated") else ""))
            if ent["id"] not in done_ids:
                done_ids.add(ent["id"])
        elif ent["status"] == "skipped":
            say("        BỎ QUA  đã có %d ký tự" % (ent.get("chars") or 0))
        else:
            say("        %s  %s" % (ent["status"].upper(), (ent.get("note") or "")[:120]))

        write_manifest(out_dir, entries)          # ghi ngay -> resume an toàn
        time.sleep(0.4)

    manifest = write_manifest(out_dir, entries)
    elapsed = time.time() - started
    item_ids = {i["id"] for i in items}
    have = [e for e in entries if e.get("id") in item_ids
            and e.get("status") in ("ok", "skipped")]
    fresh = [e for e in have if e.get("status") == "ok"]
    say("-" * 68)
    say("Xong sau %.0f giây. Có transcript: %d/%d video (mới tải %d, đã có từ trước %d)."
        % (elapsed, len(have), len(items), len(fresh), len(have) - len(fresh)))
    for k, v in sorted(manifest["by_status"].items()):
        say("   trong MANIFEST: %-10s %d" % (k, v))
    if manifest["by_lang"]:
        say("   trong MANIFEST — ngôn ngữ: " + ", ".join("%s=%d" % (k, v)
                                        for k, v in sorted(manifest["by_lang"].items())))

    if not args.no_zip:
        zip_path = HERE / "holetex-transcripts.zip"
        n = make_zip(out_dir, zip_path)
        say("Đã đóng gói %d tệp vào %s (%.1f KB)." % (
            n, zip_path, zip_path.stat().st_size / 1024.0))
        say("→ Gửi tệp holetex-transcripts.zip lên trang upload.")
    return 0


if __name__ == "__main__":
    sys.exit(main())