#!/usr/bin/env python3
"""Turn a local material folder into the plan a batch launch works from.

Two folder shapes are recognised, the ones the operator's editors deliver:

  video    <root>/<drama>/**/*.mp4              one campaign, one ad, every video
  carousel <root>/<drama>/轮播NN_*/images/*.png  one campaign, one ad PER SET

Prints one JSON document to stdout. Nothing is uploaded; the skill reads this
to know what it is about to build and to show the operator before anything
reaches the platform. Content hashes are included so identical bytes can be
spotted — for video they let the library reuse an asset, for carousels they are
a warning: the same still in two sets has to become two assets, or the second
set silently loses a frame.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import mimetypes
import re
import sys
from pathlib import Path

VIDEO_EXT = {".mp4", ".mov", ".webm"}
IMAGE_EXT = {".png", ".jpg", ".jpeg", ".webp"}
SET_DIR = re.compile(r"^轮播\d+")


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _file(path: Path, with_hash: bool) -> dict:
    content_type = mimetypes.guess_type(path.name)[0] or "application/octet-stream"
    if path.suffix.lower() == ".mp4":
        content_type = "video/mp4"
    entry = {
        "path": str(path),
        "name": path.name,
        "size_bytes": path.stat().st_size,
        "content_type": content_type,
    }
    if with_hash:
        entry["sha256"] = _sha256(path)
    return entry


def scan_drama(folder: Path, with_hash: bool) -> dict:
    sets = sorted(p for p in folder.iterdir() if p.is_dir() and SET_DIR.match(p.name))
    if sets:
        carousel_sets = []
        for set_dir in sets:
            images_dir = set_dir / "images"
            images = sorted(
                p for p in (images_dir.iterdir() if images_dir.is_dir() else [])
                if p.is_file() and p.suffix.lower() in IMAGE_EXT
            )
            caption = set_dir / "caption.txt"
            carousel_sets.append({
                "set": set_dir.name,
                "caption": caption.read_text(encoding="utf-8").strip() if caption.is_file() else None,
                "files": [_file(p, with_hash) for p in images],
            })
        all_hashes = [f.get("sha256") for s in carousel_sets for f in s["files"] if f.get("sha256")]
        return {
            "title": folder.name,
            "mode": "carousel",
            "sets": carousel_sets,
            "ad_count": len(carousel_sets),
            "images_per_ad": [len(s["files"]) for s in carousel_sets],
            "duplicate_frames_across_sets": len(all_hashes) - len(set(all_hashes)) if with_hash else None,
        }

    videos = sorted(p for p in folder.rglob("*") if p.is_file() and p.suffix.lower() in VIDEO_EXT)
    files = [_file(p, with_hash) for p in videos]
    unique = {f["sha256"]: f for f in files} if with_hash else None
    return {
        "title": folder.name,
        "mode": "video",
        "files": files,
        "unique_files": len(unique) if unique is not None else None,
        "ad_count": 1,
        "videos_per_ad": [len(unique) if unique is not None else len(files)],
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("root", help="folder whose subfolders are dramas")
    parser.add_argument("--only", action="append", default=[], help="drama folder name(s) to include; repeatable")
    parser.add_argument("--no-hash", action="store_true", help="skip SHA-256 (faster; loses duplicate detection)")
    args = parser.parse_args()

    root = Path(args.root).expanduser()
    if not root.is_dir():
        print(json.dumps({"error": f"not a directory: {root}"}), file=sys.stderr)
        return 2
    wanted = {name.strip() for name in args.only if name.strip()}
    dramas = []
    for folder in sorted(p for p in root.iterdir() if p.is_dir()):
        if wanted and folder.name not in wanted:
            continue
        drama = scan_drama(folder, with_hash=not args.no_hash)
        if drama["mode"] == "video" and not drama["files"]:
            continue
        if drama["mode"] == "carousel" and not any(s["files"] for s in drama["sets"]):
            continue
        dramas.append(drama)
    print(json.dumps({"root": str(root), "dramas": dramas}, ensure_ascii=False, indent=1))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())