#!/usr/bin/env python3
"""Hook multiplier: put N different ~3 s hooks (Kenerate stills or clips + a bold text overlay) in front of one ad
body -> N ad variants to A/B test, plus a contact sheet and a CSV of what each variant says.

Examples
  python hooks.py --body https://.../ad_body.mp4 \
      --hook "https://.../hook1.png::Stop scrolling if you hate bitter coffee" \
      --hook "https://.../hook2.mp4::POV: your 3pm slump" \
      --hook "::I tried every cold brew so you don't have to" --aspect 9:16
  python hooks.py --body body.mp4 --hooks-json hooks.json   # [{"media": "...", "text": "...", "seconds": 3}]

A hook is MEDIA::TEXT. MEDIA can be an image (gets a slow push-in), a clip (first N seconds), or empty (the text
sits over a blurred freeze of the body's first frame). TEXT can be empty (no overlay).
"""
from __future__ import annotations

import argparse
import csv
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import kmedia as K  # noqa: E402


def text_png(text: str, W: int, H: int, dst: Path, pos: str, style: str, font, accent) -> Path:
    from PIL import Image
    img = Image.new("RGBA", (W, H), (0, 0, 0, 0))
    short = min(W, H)
    bw = int(W * 0.86)
    bh = int(H * (0.22 if H > W else 0.3))
    cy = {"top": 0.24 if H > W else 0.22, "center": 0.5, "bottom": 0.68 if H > W else 0.78}[pos]
    box = ((W - bw) // 2, int(H * cy - bh / 2), (W + bw) // 2, int(H * cy + bh / 2))
    if style == "box":
        K.draw_block(img, text, box, fill=(17, 17, 17), max_size=int(short * 0.085), min_size=22, font_override=font,
                     highlight={"fill": (255, 255, 255, 240)}, max_lines=4)
    elif style == "accent":
        K.draw_block(img, text, box, fill=(17, 17, 17), max_size=int(short * 0.09), min_size=22, font_override=font,
                     highlight={"fill": tuple(accent) + (245,)}, upper=True, max_lines=4)
    else:
        K.draw_block(img, text, box, fill=(255, 255, 255), max_size=int(short * 0.1), min_size=22, font_override=font,
                     stroke_ratio=0.08, upper=True, shadow=True, max_lines=4)
    img.save(dst)
    return dst


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--body", required=True, help="the ad body clip (URL or file)")
    ap.add_argument("--hook", action="append", default=[], help="MEDIA::TEXT (repeat)")
    ap.add_argument("--hooks-json", help='[{"media": "...", "text": "...", "seconds": 3}]')
    ap.add_argument("--hook-seconds", type=float, default=3.0)
    ap.add_argument("--aspect", default="9:16")
    ap.add_argument("--res", type=int, default=1080)
    ap.add_argument("--fps", type=int, default=30)
    ap.add_argument("--text-pos", choices=["top", "center", "bottom"], default="top")
    ap.add_argument("--text-style", choices=["bold", "box", "accent"], default="bold")
    ap.add_argument("--accent", default="#FFD60A")
    K.add_font_args(ap)
    ap.add_argument("--transition", default="cut", help="cut (default, punchier) or an xfade name like fade")
    ap.add_argument("--xfade", type=float, default=0.2)
    ap.add_argument("--name", default="ad", help="output file prefix")
    ap.add_argument("--out-dir", default="kenerate-out/hook-multiplier")
    ap.add_argument("--keep-temp", action="store_true")
    a = K.parse_args(ap)

    hooks = []
    if a.hooks_json:
        for h in json.loads(Path(a.hooks_json).read_text(encoding="utf-8")):
            hooks.append({"media": h.get("media") or "", "text": h.get("text") or "", "seconds": float(h.get("seconds") or a.hook_seconds)})
    for spec in a.hook:
        media, _, text = spec.partition("::")
        hooks.append({"media": media.strip(), "text": text.strip(), "seconds": a.hook_seconds})
    if not hooks:
        K.fail("No hooks given.", "Pass --hook MEDIA::TEXT (repeat), e.g. --hook 'hook1.png::Stop scrolling'.")
    if len(hooks) > 20:
        K.fail("Too many hooks (max 20 per run).", "Split the batch.")

    st = K.Status(5, "hook-multiplier")
    K.require(status=st)
    W, H = K.size_for(a.aspect, a.res)
    accent = K.hex_rgb(a.accent)
    tmp = K.workdir(a.keep_temp)
    out = Path(a.out_dir)
    out.mkdir(parents=True, exist_ok=True)

    st.step(f"Fetching body + {sum(1 for h in hooks if h['media'])} hook media")
    body_src = K.fetch(a.body, tmp, "body", st)
    for i, h in enumerate(hooks, 1):
        h["path"] = K.fetch(h["media"], tmp, f"hook{i:02d}", st) if h["media"] else None

    st.step(f"Normalizing the body to {W}x{H}")
    body = tmp / "body_norm.mp4"
    body_d = K.normalize_clip(body_src, body, W, H, a.fps, status=st)
    freeze = K.extract_frame(body, 0.05, tmp / "body_first.png")

    st.step(f"Building {len(hooks)} hook(s)")
    built = []
    for i, h in enumerate(hooks, 1):
        st.info(f"hook {i}/{len(hooks)}: {h['text'][:60] or '(no text)'}")
        raw = tmp / f"hook{i:02d}_raw.mp4"
        src = h["path"]
        secs = h["seconds"]
        if src is None:  # blurred freeze of the body
            from PIL import ImageFilter
            im = K.open_image(freeze).convert("RGB").filter(ImageFilter.GaussianBlur(min(W, H) / 45))
            p = tmp / f"hook{i:02d}_freeze.png"
            im.save(p)
            K.still_to_clip(p, raw, W, H, secs, a.fps, move=1)
        elif K.kind_of(str(src)) == "image":
            K.still_to_clip(src, raw, W, H, secs, a.fps, move=i - 1)
        else:
            K.normalize_clip(src, raw, W, H, a.fps, 0.0, secs, status=None)
        hd = K.probe(raw)["duration"]
        if h["text"]:
            png = text_png(h["text"], W, H, tmp / f"hook{i:02d}_text.png", a.text_pos, a.text_style, a.font, accent)
            withtext = tmp / f"hook{i:02d}.mp4"
            K.overlay_png(raw, png, withtext, start=0.08, end=None, fade=0.15)
            raw = withtext
        built.append((raw, hd))
        st.progress(i * 100 / len(hooks), f"hook {i}/{len(hooks)}")

    st.step(f"Rendering {len(built)} variant(s)")
    rows, files, thumbs = [], [], []
    for i, ((hp, hd), h) in enumerate(zip(built, hooks), 1):
        dst = out / f"{K.safe_name(a.name)}_hook{i:02d}.mp4"
        st.info(f"variant {i}/{len(built)} -> {dst.name}")
        total = K.join_clips([(hp, hd), (body, body_d)], dst, a.transition, a.xfade)
        th = K.extract_frame(dst, min(1.0, hd / 2), tmp / f"thumb{i:02d}.jpg")
        files.append(dst)
        thumbs.append(th)
        rows.append({"file": dst.name, "hook_text": h["text"], "hook_media": h["media"] or "(body freeze)",
                     "hook_seconds": round(hd, 2), "total_seconds": round(total, 2)})
        st.progress(i * 100 / len(built), f"variant {i}/{len(built)}")

    st.step("Contact sheet + CSV")
    sheet = K.contact_sheet(thumbs, out / f"{K.safe_name(a.name)}_hooks_sheet.jpg", cols=min(5, len(thumbs)), cell=320,
                            labels=[f"#{i} {r['hook_text'][:28]}" for i, r in enumerate(rows, 1)])
    csv_path = out / f"{K.safe_name(a.name)}_hooks.csv"
    with open(csv_path, "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0]))
        w.writeheader()
        w.writerows(rows)
    K.finish({"variants": [K.file_info(p) for p in files], "sheet": str(sheet.resolve()), "csv": str(csv_path.resolve()),
              "body_seconds": round(body_d, 2), "frame": f"{W}x{H}", "elapsed_s": st.elapsed()})


if __name__ == "__main__":
    K.main_guard(main)
