#!/usr/bin/env python3
"""Headline multiplier: one static ad visual + N headlines -> N finished ad images. Auto-fits the text, picks the
calmest area for it, chooses white or dark text by WCAG contrast and adds a scrim when contrast is too low.

Examples
  python headlines.py --image https://.../hero.png --headline "Cold brew, zero sugar" --headline "Your 3pm fix" \
      --sub "Oat milk. Real coffee." --cta "Shop now" --brand-color "#FF5A1F"
  python headlines.py --image hero.png --headlines headlines.txt --aspect 4:5 --layout bottom
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

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

DARK, LIGHT = (17, 17, 17), (255, 255, 255)


def bands(W: int, H: int):
    h = int(H * 0.36)
    return {"top": (0, int(H * 0.04), W, int(H * 0.04) + h), "center": (0, (H - h) // 2, W, (H + h) // 2),
            "bottom": (0, H - h - int(H * 0.04), W, H - int(H * 0.04))}


def tile_means(img, box, n: int = 10):
    """Average colours of an n x n/2 grid of tiles inside box (text must read on the worst one)."""
    from PIL import ImageStat
    x0, y0, x1, y1 = box
    small = img.convert("RGB").crop(box)
    small.thumbnail((240, 240))
    tw, th = max(1, small.width // n), max(1, small.height // max(1, n // 2))
    out = []
    for yy in range(0, small.height - th + 1, th):
        for xx in range(0, small.width - tw + 1, tw):
            out.append(tuple(int(v) for v in ImageStat.Stat(small.crop((xx, yy, xx + tw, yy + th))).mean[:3]))
    return out or [K.region_stats(img, box)[0]]


def calm_score(img, box) -> float:
    """Lower = calmer: luminance spread plus edge energy."""
    from PIL import ImageStat
    reg = img.convert("L").crop(box)
    reg.thumbnail((240, 240))
    return ImageStat.Stat(reg).stddev[0] / 64.0 + K.region_stats(img, box)[1]


def worst_contrast(col, tiles) -> float:
    return min(K.contrast(col, t) for t in tiles)


def scrim(img, box, color, strength: float, side: str):
    """Soft gradient towards the text's side so the photo still shows elsewhere."""
    from PIL import Image
    W, H = img.size
    x0, y0, x1, y1 = box
    pad = int(H * 0.08)
    y0, y1 = max(0, y0 - pad), min(H, y1 + pad)
    a = Image.new("L", (1, y1 - y0))
    for i in range(y1 - y0):
        t = i / max(y1 - y0 - 1, 1)
        if side == "top":
            v = 1 - t
        elif side == "bottom":
            v = t
        else:
            v = 1 - abs(t - 0.5) * 2
        a.putpixel((0, i), int(255 * strength * min(1.0, v * 1.6)))
    a = a.resize((W, y1 - y0))
    layer = Image.new("RGBA", (W, y1 - y0), tuple(color) + (255,))
    layer.putalpha(a)
    img.alpha_composite(layer, (0, y0))


def render(base, headline: str, sub: str, cta: str, layout: str, brand, min_contrast: float, font, upper: bool):
    from PIL import Image, ImageDraw
    img = base.convert("RGBA").copy()
    W, H = img.size
    bs = bands(W, H)
    if layout == "auto":
        layout = min(bs, key=lambda k: calm_score(img, bs[k]) + (0.05 if k == "center" else 0))
    box = bs[layout]
    tiles = tile_means(img, box)
    col = max((LIGHT, DARK), key=lambda c: worst_contrast(c, tiles))
    ratio = worst_contrast(col, tiles)
    scrimmed = False
    if ratio < min_contrast:
        sc = (0, 0, 0) if col == LIGHT else (255, 255, 255)
        for strength in (0.45, 0.6, 0.75, 0.9):
            trial = img.copy()
            scrim(trial, box, sc, strength, layout)
            r2 = worst_contrast(col, tile_means(trial, box))
            if r2 >= min_contrast or strength >= 0.9:
                img, ratio, scrimmed = trial, r2, True
                break
    short = min(W, H)
    margin = int(W * 0.07)
    x0, y0, x1, y1 = margin, box[1], W - margin, box[3]
    has_sub, has_cta = bool(sub), bool(cta)
    head_h = int((y1 - y0) * (0.62 if (has_sub or has_cta) else 0.95))
    r = K.draw_block(img, headline, (x0, y0, x1, y0 + head_h), fill=col, max_size=int(short * 0.13), min_size=18,
                     kind="bold", font_override=font, upper=upper, shadow=col == LIGHT,
                     valign="bottom" if (has_sub or has_cta) else "center", max_lines=4)
    y = r["bbox"][3] + int(short * 0.025)
    if has_sub:
        r2 = K.draw_block(img, sub, (x0, y, x1, y + int(short * 0.09)), fill=col, max_size=int(short * 0.045), min_size=14,
                          kind="regular", valign="top", max_lines=2)
        y = r2["bbox"][3] + int(short * 0.03)
    cta_info = None
    if has_cta:
        f = K.load_font(int(short * 0.04), "bold")
        d = ImageDraw.Draw(img)
        tw = K.text_w(d, cta.upper(), f)
        ph = int(short * 0.075)
        pw = tw + ph
        px = (W - pw) // 2
        py = min(y, H - ph - int(H * 0.03))
        tcol, tratio = K.best_text_color(brand, (LIGHT, DARK))
        d.rounded_rectangle((px, py, px + pw, py + ph), radius=ph // 2, fill=tuple(brand) + (255,))
        bb = d.textbbox((0, 0), cta.upper(), font=f)
        d.text((px + (pw - (bb[2] - bb[0])) // 2 - bb[0], py + (ph - (bb[3] - bb[1])) // 2 - bb[1]), cta.upper(), font=f, fill=tcol)
        cta_info = {"text_color": K.rgb_hex(tcol), "contrast": tratio}
    return img, {"layout": layout, "text_color": K.rgb_hex(col), "contrast": ratio, "passes": ratio >= min_contrast,
                 "scrim": scrimmed, "font_size": r["font_size"], "lines": r["lines"], "cta": cta_info}


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--image", required=True, help="ad visual (URL or file); generate it without text")
    ap.add_argument("--headline", action="append", default=[])
    ap.add_argument("--headlines", help="text file, one headline per line")
    ap.add_argument("--sub", default="", help="sub-headline shared by all variants")
    ap.add_argument("--cta", default="", help="button text, e.g. 'Shop now'")
    ap.add_argument("--brand-color", default="#FFD60A", help="CTA button colour")
    ap.add_argument("--layout", choices=["auto", "top", "center", "bottom"], default="auto")
    ap.add_argument("--aspect", default="", help="crop the visual to 1:1, 4:5, 9:16, 16:9 first (default: keep)")
    ap.add_argument("--min-contrast", type=float, default=4.5, help="WCAG ratio to reach (default 4.5)")
    ap.add_argument("--no-upper", action="store_true", help="keep headline case (default: UPPERCASE)")
    K.add_font_args(ap)
    ap.add_argument("--format", choices=["png", "jpg"], default="jpg")
    ap.add_argument("--name", default="headline")
    ap.add_argument("--out-dir", default="kenerate-out/headline-multiplier")
    a = K.parse_args(ap)

    heads = list(a.headline)
    if a.headlines:
        heads += [ln.strip() for ln in Path(a.headlines).read_text(encoding="utf-8").splitlines() if ln.strip()]
    if not heads:
        K.fail("No headlines given.", "Pass --headline TEXT (repeat) or --headlines file.txt.")
    st = K.Status(3, "headline-multiplier")
    K.require(ffmpeg=False, status=st)
    out = Path(a.out_dir)
    out.mkdir(parents=True, exist_ok=True)
    tmp = K.workdir()

    st.step("Fetching the visual")
    base = K.open_image(K.fetch(a.image, tmp, "visual", st)).convert("RGB")
    if a.aspect:
        W, H = K.size_for(a.aspect, min(base.size))
        base = K.cover(base, W, H)
    st.info(f"visual {base.width}x{base.height}")

    st.step(f"Rendering {len(heads)} variant(s)")
    files, rows = [], []
    brand = K.hex_rgb(a.brand_color)
    for i, h in enumerate(heads, 1):
        img, meta = render(base, h, a.sub, a.cta, a.layout, brand, a.min_contrast, a.font, not a.no_upper)
        dst = out / f"{K.safe_name(a.name)}_{i:02d}.{a.format}"
        img.convert("RGB").save(dst, quality=92) if a.format == "jpg" else img.save(dst)
        files.append(dst)
        rows.append({"file": str(dst.resolve()), "headline": h, **meta})
        if not meta["passes"]:
            st.warn(f"#{i} contrast {meta['contrast']} < {a.min_contrast} even with a scrim: try --layout or a calmer visual")
        st.progress(i * 100 / len(heads), f"{i}/{len(heads)} {meta['layout']} {meta['text_color']} {meta['contrast']}:1")

    st.step("Contact sheet")
    sheet = K.contact_sheet(files, out / f"{K.safe_name(a.name)}_sheet.jpg", cols=min(4, len(files)), cell=360,
                            labels=[f"#{i} {r['contrast']}:1{' scrim' if r['scrim'] else ''}" for i, r in enumerate(rows, 1)])
    (out / f"{K.safe_name(a.name)}_report.json").write_text(json.dumps(rows, indent=2, ensure_ascii=False), encoding="utf-8")
    K.finish({"variants": rows, "sheet": str(sheet.resolve()), "elapsed_s": st.elapsed()})


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