"""
place_logo.py

Takes a mockup image, a zone definition JSON (from the zone-definition tool),
and one or two logo files (full lockup + optional mark-only), and composites
the logo into each defined zone -- resizing for rect zones, perspective-warping
for quad zones.

Usage:
    python place_logo.py \
        --mockup spray-bottle-01.png \
        --zones spray-bottle-01.json \
        --logo-full client-logo-full.png \
        --logo-mark client-logo-mark.png \
        --out output.png
"""

import argparse
import json
import numpy as np
from PIL import Image, ImageChops


def find_coeffs(source_corners, target_corners):
    """
    Compute the 8 coefficients PIL's Image.transform(..., Image.PERSPECTIVE, ...)
    needs to map `target_corners` (in the OUTPUT canvas) back to `source_corners`
    (in the INPUT image), so that source pixels land on the target quad.
    """
    matrix = []
    for (x, y), (X, Y) in zip(target_corners, source_corners):
        matrix.append([x, y, 1, 0, 0, 0, -X * x, -X * y])
        matrix.append([0, 0, 0, x, y, 1, -Y * x, -Y * y])
    A = np.array(matrix, dtype=float)
    B = np.array(source_corners, dtype=float).reshape(8)
    coeffs = np.linalg.solve(A, B)
    return coeffs


def fit_logo_to_width(logo, target_width_px):
    """Resize logo preserving aspect ratio to a target pixel width."""
    w, h = logo.size
    scale = target_width_px / w
    new_size = (max(1, int(w * scale)), max(1, int(h * scale)))
    return logo.resize(new_size, Image.LANCZOS)


def bbox_of_points(points):
    xs = [p[0] for p in points]
    ys = [p[1] for p in points]
    return min(xs), min(ys), max(xs), max(ys)


def composite_normal(canvas, overlay, position):
    """Paste an RGBA overlay onto canvas at position, respecting alpha."""
    canvas.paste(overlay, position, overlay)


def composite_multiply(canvas, overlay, position):
    """
    Multiply-blend overlay's RGB into canvas at position, using overlay's
    alpha channel as the blend mask, so the logo picks up shading/texture
    from the mockup underneath (useful for embossed/fabric surfaces).
    """
    x, y = position
    w, h = overlay.size
    region = canvas.crop((x, y, x + w, y + h)).convert("RGB")
    overlay_rgb = overlay.convert("RGB")
    blended = ImageChops.multiply(region, overlay_rgb)
    alpha = overlay.split()[3]
    canvas.paste(blended, (x, y), alpha)


def place_rect_zone(canvas, logo, zone):
    x0, y0, x1, y1 = bbox_of_points(zone["points"])
    zone_w = x1 - x0
    target_w = int(zone_w * (zone.get("max_logo_width_pct", 40) / 100))
    resized = fit_logo_to_width(logo, target_w)
    lw, lh = resized.size
    zone_h = y1 - y0
    px = x0 + (zone_w - lw) // 2
    py = y0 + (zone_h - lh) // 2

    if zone.get("blend") == "multiply":
        composite_multiply(canvas, resized, (px, py))
    else:
        composite_normal(canvas, resized, (px, py))


def place_quad_zone(canvas, logo, zone):
    x0, y0, x1, y1 = bbox_of_points(zone["points"])
    zone_w = x1 - x0
    target_w = int(zone_w * (zone.get("max_logo_width_pct", 40) / 100))
    resized = fit_logo_to_width(logo, target_w)
    lw, lh = resized.size

    quad = zone["points"]
    cx = sum(p[0] for p in quad) / 4
    cy = sum(p[1] for p in quad) / 4
    scaled_quad = []
    for px, py in quad:
        dx = (px - cx) * (target_w / zone_w)
        dy = (py - cy) * (target_w / zone_w)
        scaled_quad.append((cx + dx, cy + dy))

    qx0, qy0, qx1, qy1 = bbox_of_points(scaled_quad)
    pad = 4
    ox0, oy0 = int(qx0 - pad), int(qy0 - pad)
    ox1, oy1 = int(qx1 + pad), int(qy1 + pad)
    out_w, out_h = ox1 - ox0, oy1 - oy0

    local_quad = [(px - ox0, py - oy0) for px, py in scaled_quad]
    logo_corners = [(0, 0), (lw, 0), (lw, lh), (0, lh)]

    coeffs = find_coeffs(logo_corners, local_quad)
    warped = resized.transform(
        (out_w, out_h), Image.PERSPECTIVE, coeffs, Image.BICUBIC
    )

    if zone.get("blend") == "multiply":
        composite_multiply(canvas, warped, (ox0, oy0))
    else:
        composite_normal(canvas, warped, (ox0, oy0))


def place_zones(mockup_path, zones_path, logo_full_path, logo_mark_path, out_path):
    with open(zones_path) as f:
        spec = json.load(f)

    canvas = Image.open(mockup_path).convert("RGBA")
    logo_full = Image.open(logo_full_path).convert("RGBA") if logo_full_path else None
    logo_mark = Image.open(logo_mark_path).convert("RGBA") if logo_mark_path else logo_full

    for zone in spec["zones"]:
        variant = zone.get("logo_variant", "full-lockup")
        logo = logo_mark if variant == "mark-only" and logo_mark is not None else logo_full
        if logo is None:
            print(f"Skipping zone {zone['zone_id']}: no logo provided for variant '{variant}'")
            continue

        if zone["zone_type"] == "rect":
            place_rect_zone(canvas, logo, zone)
        elif zone["zone_type"] == "quad":
            place_quad_zone(canvas, logo, zone)
        else:
            print(f"Unknown zone_type '{zone['zone_type']}' for zone {zone['zone_id']}, skipping.")

    canvas.save(out_path)
    print(f"Saved: {out_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Composite a logo into mockup zones.")
    parser.add_argument("--mockup", required=True, help="Path to blank mockup image")
    parser.add_argument("--zones", required=True, help="Path to zone definition JSON")
    parser.add_argument("--logo-full", help="Path to full logo lockup (PNG, transparent bg)")
    parser.add_argument("--logo-mark", help="Path to standalone mark/icon (PNG, transparent bg)")
    parser.add_argument("--out", required=True, help="Output image path")
    args = parser.parse_args()

    place_zones(args.mockup, args.zones, args.logo_full, args.logo_mark, args.out)