#!/usr/bin/env python3
"""Image editing pipeline (Pillow): crop, resize, pad, rotate, colour grade / LUT, sharpen, blur, vignette,
watermark (logo or text), border, round corners, format + size cap. Works on one image or a whole batch, with a
dry run and a before / after preview sheet.

Examples
  python imgedit.py --inputs photo.jpg --ops "crop=4:5@0.5,0.4; resize=1080x1350; look=warm-film:0.7; sharpen=0.8"
  python imgedit.py --inputs ./shoot/ --ops "resize=2000w; watermark=logo.png@br,0.14,0.7; text=@brew@bl" --format webp
  python imgedit.py --inputs a.png b.png --ops "pad=1080x1080:#FFFFFF; round=48" --format png --dry-run

Ops (separate with ';', applied in order):
  crop=16:9[@fx,fy]       crop to an aspect around a focus point (0..1, default centre)
  crop=x,y,w,h            crop a pixel box
  resize=WxH | Ww | Hh | 50%   resize (WxH fills and crops; add :contain to letterbox, :stretch to distort)
  pad=WxH[:#hex|:blur]    fit inside WxH on a colour or blurred fill
  rotate=DEG  flip=h|v    autocontrast
  exposure=+0.3 (stops)  contrast=1.1  saturation=1.2  warmth=-100..100  tint=-100..100  brightness=1.05
  look=NAME[:intensity]   built-in look (teal-orange, warm-film, cool-night, bleach-bypass, vintage-fade, mono,
                          punchy, golden-hour, matte-pastel, noir)
  lut=FILE.cube[:intensity]
  sharpen=AMOUNT (0.3-2)  blur=RADIUS  denoise=1|2  vignette=0..1  grayscale  sepia
  watermark=FILE[@pos,scale,opacity]   pos tl|tr|bl|br|c (default br), scale = logo width / image width
  text=TEXT[@pos,size]    text watermark, size = fraction of the short side (default 0.035)
  border=PX[:#hex]   round=RADIUS (makes PNG transparency)
"""
from __future__ import annotations

import argparse
import glob
import sys
from pathlib import Path

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

POS = {"tl": (0, 0), "tr": (1, 0), "bl": (0, 1), "br": (1, 1), "c": (0.5, 0.5), "t": (0.5, 0), "b": (0.5, 1)}


def parse_ops(spec: str):
    ops = []
    for part in [p.strip() for p in spec.replace("\n", ";").split(";") if p.strip()]:
        name, _, val = part.partition("=")
        ops.append((name.strip().lower(), val.strip()))
    return ops


KNOWN = {"crop", "resize", "pad", "rotate", "flip", "autocontrast", "exposure", "contrast", "saturation", "warmth", "tint",
         "brightness", "look", "lut", "sharpen", "blur", "denoise", "vignette", "grayscale", "sepia", "watermark", "text",
         "border", "round"}


def validate(ops, base_dir: Path):
    errs = []
    for n, v in ops:
        if n not in KNOWN:
            errs.append(f"unknown op '{n}'")
        if n == "watermark":
            path = v.rsplit("@", 1)[0] if "@" in v else v
            if not K.is_url(path) and not Path(path).exists():
                errs.append(f"watermark file not found: {path}")
        if n == "lut":
            path = split_intensity(v)[0]
            if not Path(path).exists():
                errs.append(f"LUT not found: {path}")
            else:
                try:
                    K.read_cube(Path(path))
                except ValueError as e:
                    errs.append(str(e))
        if n == "look":
            if split_intensity(v)[0] not in K.LOOKS:
                errs.append(f"unknown look '{v}' (looks: {', '.join(K.LOOKS)})")
    return errs


def split_intensity(v: str):
    """'name:0.7' -> ('name', 0.7); keeps Windows drive letters ('C:\\x.cube:0.5')."""
    head, sep, tail = v.rpartition(":")
    if sep and head and not (len(head) == 1 and head.isalpha()):
        try:
            return head, float(tail)
        except ValueError:
            pass
    return v, 1.0


def apply(img, n: str, v: str, font, cache: dict):
    from PIL import Image, ImageDraw, ImageEnhance, ImageFilter, ImageOps
    W, H = img.size
    if n == "crop":
        if "," in v and ":" not in v.split("@")[0]:
            x, y, w, h = (int(float(t)) for t in v.split(","))
            return img.crop((x, y, x + w, y + h))
        asp, _, foc = v.partition("@")
        fx, fy = (float(t) for t in foc.split(",")) if foc else (0.5, 0.5)
        a, b = (float(t) for t in asp.split(":"))
        if W / H > a / b:
            nw, nh = int(H * a / b), H
        else:
            nw, nh = W, int(W * b / a)
        x, y = int((W - nw) * fx), int((H - nh) * fy)
        return img.crop((x, y, x + nw, y + nh))
    if n == "resize":
        size, _, mode = v.partition(":")
        if size.endswith("%"):
            k = float(size[:-1]) / 100
            return img.resize((max(1, int(W * k)), max(1, int(H * k))), Image.LANCZOS)
        if size.endswith("w"):
            w = int(size[:-1])
            return img.resize((w, max(1, round(H * w / W))), Image.LANCZOS)
        if size.endswith("h"):
            h = int(size[:-1])
            return img.resize((max(1, round(W * h / H)), h), Image.LANCZOS)
        w, h = (int(t) for t in size.lower().split("x"))
        if mode == "stretch":
            return img.resize((w, h), Image.LANCZOS)
        if mode == "contain":
            return K.contain(img.convert("RGBA"), w, h, bg=(0, 0, 0, 0) if img.mode == "RGBA" else (0, 0, 0, 255))
        return K.cover(img, w, h)
    if n == "pad":
        size, _, fill = v.partition(":")
        w, h = (int(t) for t in size.lower().split("x"))
        if fill == "blur":
            return K.contain(img.convert("RGB"), w, h, blur_bg=True)
        col = K.hex_rgb(fill) + (255,) if fill else (255, 255, 255, 255)
        return K.contain(img.convert("RGBA"), w, h, bg=col)
    if n == "rotate":
        return img.rotate(-float(v), expand=True, resample=Image.BICUBIC, fillcolor=(0, 0, 0, 0) if img.mode == "RGBA" else (0, 0, 0))
    if n == "flip":
        return ImageOps.mirror(img) if v.startswith("h") else ImageOps.flip(img)
    if n == "autocontrast":
        return _keep_alpha(img, lambda im: ImageOps.autocontrast(im, cutoff=0.5))
    if n == "brightness":
        return _keep_alpha(img, lambda im: ImageEnhance.Brightness(im).enhance(float(v)))
    if n in ("exposure", "contrast", "saturation", "warmth", "tint", "look", "lut"):
        kw = {"exposure": 0.0, "contrast_k": 1.0, "saturation": 1.0, "warmth": 0.0, "tint": 0.0}
        look, inten = None, 1.0
        if n == "exposure":
            kw["exposure"] = float(v)
        elif n == "contrast":
            kw["contrast_k"] = float(v)
        elif n in ("saturation", "warmth", "tint"):
            kw[n] = float(v)
        else:
            look, inten = split_intensity(v)
        key = (n, v)
        if key not in cache:
            cache[key] = K.write_cube(K.grade_fn(look, inten, **kw), cache["_tmp"] / f"op{len(cache)}.cube", size=33)
        return K.apply_look_image(img, str(cache[key]), 1.0)
    if n == "sharpen":
        amt = float(v or 1)
        return _keep_alpha(img, lambda im: im.filter(ImageFilter.UnsharpMask(radius=max(1.0, min(W, H) / 900), percent=int(120 * amt), threshold=2)))
    if n == "blur":
        return img.filter(ImageFilter.GaussianBlur(float(v or 2)))
    if n == "denoise":
        return _keep_alpha(img, lambda im: im.filter(ImageFilter.MedianFilter(3 if (v or "1") == "1" else 5)))
    if n == "vignette":
        k = float(v or 0.4)
        mask = Image.new("L", (256, 256), 0)
        d = ImageDraw.Draw(mask)
        for i in range(128):
            d.ellipse((i - 40, i - 40, 296 - i, 296 - i), fill=int(255 * (i / 127) ** 0.6))
        mask = mask.resize((W, H), Image.BILINEAR).filter(ImageFilter.GaussianBlur(min(W, H) / 30))
        dark = Image.new("RGB", (W, H), (0, 0, 0))
        rgb = img.convert("RGB")
        out = Image.composite(rgb, Image.blend(rgb, dark, k), mask)
        return _restore_alpha(out, img)
    if n == "grayscale":
        return _keep_alpha(img, lambda im: ImageOps.grayscale(im).convert("RGB"))
    if n == "sepia":
        return _keep_alpha(img, lambda im: ImageOps.colorize(ImageOps.grayscale(im), (40, 26, 13), (255, 240, 205)))
    if n == "watermark":
        path, _, opts = v.rpartition("@") if "@" in v else (v, "", "")
        o = (opts.split(",") + ["", "", ""])[:3]
        pos, scale, op = (o[0] or "br"), float(o[1] or 0.14), float(o[2] or 0.7)
        if "wm" not in cache or cache["wm"][0] != path:
            cache["wm"] = (path, K.open_image(K.fetch(path, cache["_tmp"], "wm")).convert("RGBA"))
        lg = cache["wm"][1].copy()
        w = max(8, int(W * scale))
        lg = lg.resize((w, max(1, round(lg.height * w / lg.width))), Image.LANCZOS)
        lg = K.with_opacity(lg, op)
        m = int(min(W, H) * 0.035)
        px, py = POS.get(pos, (1, 1))
        base = img.convert("RGBA")
        base.alpha_composite(lg, (int(m + (W - 2 * m - lg.width) * px), int(m + (H - 2 * m - lg.height) * py)))
        return base if img.mode == "RGBA" else base.convert("RGB")
    if n == "text":
        txt, _, opts = v.rpartition("@") if "@" in v else (v, "", "")
        o = (opts.split(",") + ["", ""])[:2]
        pos, size = (o[0] or "br"), float(o[1] or 0.035)
        spr = K.text_img(txt, max(10, int(min(W, H) * size)), (255, 255, 255, 230), "bold", font, stroke=max(1, int(min(W, H) * size * 0.08)),
                         stroke_fill=(0, 0, 0, 160))
        m = int(min(W, H) * 0.035)
        px, py = POS.get(pos, (1, 1))
        base = img.convert("RGBA")
        base.alpha_composite(spr, (int(m + (W - 2 * m - spr.width) * px), int(m + (H - 2 * m - spr.height) * py)))
        return base if img.mode == "RGBA" else base.convert("RGB")
    if n == "border":
        px_, _, col = v.partition(":")
        b = int(px_)
        c = K.hex_rgb(col) if col else (255, 255, 255)
        return ImageOps.expand(img, border=b, fill=c + ((255,) if img.mode == "RGBA" else ()))
    if n == "round":
        r = int(v or 32)
        mask = Image.new("L", (W * 2, H * 2), 0)
        ImageDraw.Draw(mask).rounded_rectangle((0, 0, W * 2 - 1, H * 2 - 1), radius=r * 2, fill=255)
        mask = mask.resize((W, H), Image.LANCZOS)
        out = img.convert("RGBA")
        a = out.getchannel("A")
        from PIL import ImageChops
        out.putalpha(ImageChops.multiply(a, mask))
        return out
    raise ValueError(f"unknown op '{n}'")


def _keep_alpha(img, fn):
    out = fn(img.convert("RGB"))
    return _restore_alpha(out, img)


def _restore_alpha(out, src):
    if src.mode == "RGBA":
        out = out.convert("RGBA")
        out.putalpha(src.getchannel("A") if src.size == out.size else out.getchannel("A"))
    return out


def gather(inputs):
    out = []
    for s in inputs:
        if K.is_url(s):
            out.append(s)
        elif Path(s).is_dir():
            out += sorted(str(p) for p in Path(s).iterdir() if p.suffix.lower() in K.IMAGE_EXT)
        elif any(c in s for c in "*?["):
            out += sorted(p for p in glob.glob(s) if Path(p).suffix.lower() in K.IMAGE_EXT)
        else:
            out.append(s)
    return out


def save(img, dst: Path, fmt: str, quality: int, max_kb: float):
    if fmt in ("jpg", "jpeg"):
        im = img.convert("RGB") if img.mode != "RGB" else img
        if img.mode == "RGBA":
            from PIL import Image
            bg = Image.new("RGB", img.size, (255, 255, 255))
            bg.paste(img, (0, 0), img)
            im = bg
        q = quality
        while True:
            im.save(dst, quality=q, optimize=True, progressive=True)
            if not max_kb or dst.stat().st_size / 1024 <= max_kb or q <= 40:
                return q
            q -= 6
    if fmt == "webp":
        q = quality
        while True:
            img.save(dst, quality=q, method=5)
            if not max_kb or dst.stat().st_size / 1024 <= max_kb or q <= 35:
                return q
            q -= 7
    img.save(dst, optimize=True)
    return None


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--inputs", nargs="+", required=True, help="images, URLs, folders or globs")
    ap.add_argument("--ops", default="", help="op pipeline, e.g. \"crop=4:5; resize=1080x1350; look=warm-film:0.6\"")
    ap.add_argument("--ops-file", help="text file with one op per line (same syntax)")
    ap.add_argument("--format", choices=["keep", "jpg", "png", "webp"], default="keep")
    ap.add_argument("--quality", type=int, default=90)
    ap.add_argument("--max-kb", type=float, default=0, help="shrink jpg / webp quality until each file is under this size")
    ap.add_argument("--suffix", default="_edit", help="added to each output name")
    K.add_font_args(ap, "TTF/OTF for text= watermarks")
    ap.add_argument("--preview", action="store_true", help="before / after contact sheet (first 6 images)")
    ap.add_argument("--dry-run", action="store_true", help="validate the ops and list the outputs without writing anything")
    ap.add_argument("--out-dir", default="kenerate-out/image-editing")
    a = K.parse_args(ap)

    st = K.Status(4, "image-editing")
    K.require(ffmpeg=False, status=st)
    spec = a.ops + (";" + Path(a.ops_file).read_text(encoding="utf-8") if a.ops_file else "")
    ops = parse_ops(spec)
    st.step(f"Checking {len(ops)} op(s)")
    if not ops:
        K.fail("No ops given.", "Pass --ops \"resize=1080w; sharpen=0.8\" (see --help for the list).")
    errs = validate(ops, Path.cwd())
    if errs:
        K.fail("; ".join(errs), "Fix the op list (python imgedit.py --help shows every op).")
    st.info(" -> ".join(f"{n}{'=' + v if v else ''}" for n, v in ops))
    srcs = gather(a.inputs)
    if not srcs:
        K.fail("No images found in --inputs.")
    out_dir = Path(a.out_dir)
    plan = []
    for i, s in enumerate(srcs, 1):
        stem = K.safe_name(Path(s.split("?")[0]).stem) if not K.is_url(s) else f"img{i:02d}"
        ext = Path(s.split("?")[0]).suffix.lower().lstrip(".") or "png"
        fmt = (ext if ext in ("jpg", "jpeg", "png", "webp") else "png") if a.format == "keep" else a.format
        if any(n == "round" for n, _ in ops) and fmt in ("jpg", "jpeg"):
            fmt = "png"
        plan.append((s, out_dir / f"{stem}{a.suffix}.{'jpg' if fmt == 'jpeg' else fmt}", fmt))
    if a.dry_run:
        for s, d, _ in plan:
            st.info(f"{s} -> {d}")
        K.finish({"dry_run": True, "ops": [f"{n}={v}" if v else n for n, v in ops],
                  "outputs": [str(d) for _, d, _ in plan], "count": len(plan)})
        return

    tmp = K.workdir()
    out_dir.mkdir(parents=True, exist_ok=True)
    st.step(f"Fetching {len(srcs)} image(s)")
    paths = K.fetch_all(srcs, tmp / "in", st, prefix="img")
    st.step("Editing")
    cache: dict = {"_tmp": tmp}
    results, before, after = [], [], []
    for i, (p, (_, dst, fmt)) in enumerate(zip(paths, plan), 1):
        img = K.open_image(p)
        w0, h0 = img.size
        for n, v in ops:
            img = apply(img, n, v, a.font, cache)
        q = save(img, dst, fmt, a.quality, a.max_kb)
        results.append({"input": str(p), "output": str(dst.resolve()), "from": f"{w0}x{h0}", "to": f"{img.width}x{img.height}",
                        "kb": round(dst.stat().st_size / 1024, 1), "quality": q})
        st.progress(i * 100 / len(paths), f"{i}/{len(paths)} {dst.name}")
        if len(before) < 6:
            before.append(p)
            after.append(dst)
    st.step("Preview" if a.preview else "Done")
    sheet = None
    if a.preview:
        pairs, labels = [], []
        for b, f in zip(before, after):
            pairs += [b, f]
            labels += [f"before {Path(b).name[:24]}", "after"]
        sheet = K.contact_sheet(pairs, out_dir / "_preview.jpg", cols=4, cell=320, labels=labels)
    over = [r for r in results if a.max_kb and r["kb"] > a.max_kb]
    if over:
        st.warn(f"{len(over)} file(s) still over {a.max_kb} KB - resize smaller or use webp")
    K.finish({"outputs": results, "count": len(results), "preview": str(sheet.resolve()) if sheet else None, "elapsed_s": st.elapsed()})


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