#!/usr/bin/env python3
"""moon_composite.py — paint a high-res lunar texture into a sky-cam frame.

Goal: from a frame east captured of a small white blob (~38 px), produce a
1920x1080 image that looks like east took it through a 65x telephoto.

What's real, from east:
  - Sky color, atmospheric halo, any clouds drifting past
  - Time, parallactic angle (orientation of "up" on the moon)
  - Position of the moon in the frame at that instant

What's borrowed:
  - The lunar surface texture (one cached high-res reference image)

Library entry point:

  from moon_composite import composite_full_moon
  composite_full_moon(
      source_jpg='/data/east/2026-04-29/21-07-00.jpg',
      detection=detect_moon(...),
      when_utc=datetime(2026, 4, 29, 21, 7, 0, tzinfo=timezone.utc),
      ref_moon_path='/path/to/full-moon.jpg',
      out_path='/movies/east/full-moons/2026-04-full-moon.jpg',
  )

CLI:

  python3 moon_composite.py SOURCE.jpg WHEN_UTC REF_MOON.jpg OUT.jpg
"""
from __future__ import annotations

import argparse
import math
import sys
from datetime import datetime, timezone
from pathlib import Path

import numpy as np
from PIL import Image, ImageDraw, ImageFilter, ImageFont

# Pillow ≥ 10 renamed resampling constants
try:
    LANCZOS = Image.Resampling.LANCZOS
    BICUBIC = Image.Resampling.BICUBIC
except AttributeError:  # Pillow < 10
    LANCZOS = Image.LANCZOS
    BICUBIC = Image.BICUBIC

# Common DejaVu Sans paths across distros; first hit wins.
_DEJAVU_CANDIDATES = [
    '/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf',
    '/usr/share/fonts/dejavu/DejaVuSans.ttf',
    '/usr/share/fonts/TTF/DejaVuSans.ttf',
    '/Library/Fonts/DejaVuSans.ttf',  # macOS homebrew
]


def _load_font(size: int = 18) -> ImageFont.ImageFont:
    """Return a Unicode-capable font at `size` pt, or the PIL default."""
    for path in _DEJAVU_CANDIDATES:
        try:
            return ImageFont.truetype(path, size)
        except (OSError, IOError):
            pass
    return ImageFont.load_default()


def _draw_caption(draw: ImageDraw.ImageDraw, caption: str,
                  xy: tuple[int, int], output_w: int) -> None:
    """Draw a drop-shadow caption, sanitising characters the default font
    cannot encode when no TTF is available."""
    font = _load_font(max(14, output_w // 100))
    # If we fell back to the bitmap default font it only handles latin-1; strip
    # anything outside that range rather than crashing.
    if not hasattr(font, 'path'):
        caption = caption.encode('latin-1', errors='replace').decode('latin-1')
        caption = caption.replace('\x3f', '-')  # '?' placeholder → dash
    x, y = xy
    draw.text((x + 2, y + 2), caption, fill=(0, 0, 0), font=font)
    draw.text((x,     y),     caption, fill=(220, 220, 220), font=font)


def _square_crop_to_disk(ref: Image.Image, threshold: int = 25) -> Image.Image:
    """Tight-crop a reference moon image to the disk's bounding square.

    Most lunar reference photos have the disk on a large black field with
    significant padding.  We threshold on luminance and crop to bbox + small
    margin so the disk fills our target square evenly.
    """
    g = np.asarray(ref.convert('L'))
    mask = g >= threshold
    ys, xs = np.where(mask)
    if len(xs) == 0:
        return ref
    x0, x1 = int(xs.min()), int(xs.max())
    y0, y1 = int(ys.min()), int(ys.max())
    cx = (x0 + x1) // 2
    cy = (y0 + y1) // 2
    half = max(x1 - x0, y1 - y0) // 2 + 4  # tiny margin
    left = max(0, cx - half)
    top = max(0, cy - half)
    right = min(ref.width, cx + half)
    bottom = min(ref.height, cy + half)
    return ref.crop((left, top, right, bottom))


def _atmospheric_sky_background(
    sky_img: Image.Image,
    moon_cx: float,
    moon_cy: float,
    moon_diam_px: float,
) -> Image.Image:
    """Smooth atmospheric glow on a deep night-sky backdrop.

    Samples the real halo colour from a ring just outside the moon disk, then
    builds a smooth radial gradient from that colour at the disk edge to deep
    night-sky beyond ~2.5× the radius.  No sky-cam pixels are used directly,
    so there is no cloud texture or foil artefact regardless of how bright or
    hazy the original sky was.
    """
    arr = np.asarray(sky_img).astype(np.float32)
    h, w = arr.shape[:2]
    r = moon_diam_px / 2.0

    yy, xx = np.ogrid[0:h, 0:w]
    dist = np.sqrt((xx - moon_cx) ** 2 + (yy - moon_cy) ** 2)

    # Sample mean colour of the real halo in a ring from 1.15× to 1.8× radius.
    ring = (dist >= 1.15 * r) & (dist < 1.8 * r)
    if ring.any():
        halo_rgb = arr[ring].mean(axis=0)          # real atmospheric colour
    else:
        halo_rgb = np.array([80, 85, 95], dtype=np.float32)
    # Scale to a tasteful glow level (original is often overexposed)
    halo_rgb = np.clip(halo_rgb * 0.35, 0, 180).astype(np.float32)

    # Deep night-sky colour — very dark desaturated blue
    night = np.array([4, 6, 14], dtype=np.float32)

    # Radial alpha: 1.0 at the moon disk edge (dist_norm=1), 0 at 3× radius.
    dist_norm = dist / r
    glow_alpha = np.clip(1.0 - (dist_norm - 1.0) / 2.0, 0.0, 1.0) ** 2
    glow_alpha = glow_alpha[..., None]

    bg_arr = glow_alpha * halo_rgb + (1.0 - glow_alpha) * night
    return Image.fromarray(np.clip(bg_arr, 0, 255).astype(np.uint8))


def _disk_mask(size: int, feather_px: int = 6) -> Image.Image:
    """Soft circular alpha mask the size of the reference moon image."""
    m = Image.new('L', (size, size), 0)
    draw = ImageDraw.Draw(m)
    # Inset slightly so the feather sits inside the disk edge
    draw.ellipse((feather_px, feather_px, size - feather_px, size - feather_px),
                 fill=255)
    if feather_px > 0:
        m = m.filter(ImageFilter.GaussianBlur(radius=feather_px))
    return m


def _apply_phase_shadow(disk: Image.Image, phase_angle_deg: float,
                        waxing: bool) -> Image.Image:
    """Darken the un-lit portion of the disk based on phase.

    phase_angle_deg: 0 = full, 90 = quarter, 180 = new.
    waxing: True = lit on right, False = lit on left.

    Implementation: the terminator is an ellipse whose semi-minor axis is
    cos(phase_angle).  Pixels on the un-lit side are multiplied by a small
    factor (not zero, so the un-lit limb stays visible like real earthshine).
    """
    if phase_angle_deg < 1.0:
        return disk  # full enough that shadow would be a single-pixel sliver
    w, h = disk.size
    # Build a mask: 1.0 in lit area, 0.04 in un-lit area, soft transition near
    # the terminator.
    cx, cy = w / 2.0, h / 2.0
    r = min(w, h) / 2.0
    yy, xx = np.mgrid[0:h, 0:w].astype(np.float32)
    # Normalise to disk coords (-1..1)
    nx = (xx - cx) / r
    ny = (yy - cy) / r
    # Distance from disk centre (we still want to clip to the disk)
    in_disk = (nx * nx + ny * ny) <= 1.0
    # Terminator equation: x_norm = cos(phase)  on the appropriate side.
    # For waxing moon, the lit portion is right of the terminator (nx > x_t).
    cos_p = math.cos(math.radians(phase_angle_deg))
    # When the moon is more than half lit (cos_p > 0), terminator is on the
    # un-lit side and the lit portion is broader.  When less than half
    # (cos_p < 0), terminator is on the lit side.
    # Distance from terminator (positive = lit side)
    if waxing:
        d = nx - (-cos_p)
    else:
        d = -(nx - cos_p)
    # Smooth step around the terminator (~2% of radius)
    soft_px = max(1.5 / r, 0.01)
    lit = np.clip(0.5 + d / (2 * soft_px), 0.04, 1.0)
    lit = np.where(in_disk, lit, 1.0)  # leave outside-disk untouched
    arr = np.asarray(disk).astype(np.float32)
    arr = arr * lit[..., None]
    return Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8), disk.mode)


def composite_full_moon(
    source_jpg: str,
    detection,                                  # MoonDetection from moon_detect
    when_utc: datetime,
    ref_moon_path: str,
    out_path: str,
    output_size: tuple[int, int] = (1920, 1080),
    moon_height_pct: float = 0.70,
    caption: str | None = None,
):
    """Build the full-moon close-up composite and write it to out_path."""
    # Lazy import — avoids loading skyfield when caller doesn't need it
    import moon_phase

    src = Image.open(source_jpg).convert('RGB')
    cx, cy = detection.centroid_xy
    src_diam = detection.diameter_px

    # ── Crop east around the moon, sized so the moon fills moon_height_pct ──
    crop_h = int(round(src_diam / moon_height_pct))
    crop_w = int(round(crop_h * output_size[0] / output_size[1]))
    sw, sh = src.size
    crop_w = min(crop_w, sw)
    crop_h = min(crop_h, sh)
    left = int(round(cx - crop_w / 2))
    top = int(round(cy - crop_h / 2))
    left = max(0, min(left, sw - crop_w))
    top = max(0, min(top, sh - crop_h))
    crop = src.crop((left, top, left + crop_w, top + crop_h))
    bg = crop.resize(output_size, LANCZOS)

    # ── Moon centre in output space ──
    moon_x_in_crop = cx - left
    moon_y_in_crop = cy - top
    scale = output_size[1] / crop_h
    out_moon_cx = moon_x_in_crop * scale
    out_moon_cy = moon_y_in_crop * scale

    # ── Replace raw upscaled sky with atmospheric halo on dark background ──
    out_moon_diam = output_size[1] * moon_height_pct
    bg = _atmospheric_sky_background(bg, out_moon_cx, out_moon_cy, out_moon_diam)

    # ── Load reference texture, tight-crop to disk ──
    ref = Image.open(ref_moon_path).convert('RGB')
    ref = _square_crop_to_disk(ref)
    target_size = int(round(output_size[1] * moon_height_pct))
    target_size += target_size % 2  # even
    ref_resized = ref.resize((target_size, target_size), LANCZOS)
    ref_resized = ref_resized.filter(
        ImageFilter.UnsharpMask(radius=2, percent=160, threshold=2)
    )

    # Phase shadow is NOT applied here: the NASA SVS Dial-a-Moon source already
    # renders the correct terminator, libration and earthshine for the exact
    # hour.  Applying _apply_phase_shadow on top would double-darken the unlit
    # limb.  (_apply_phase_shadow remains available for callers that supply a
    # full-moon reference photo rather than a phase-correct NASA render.)

    # ── Parallactic-angle rotation ──
    par = moon_phase.parallactic_angle(when_utc)
    ref_rot = ref_resized.rotate(-par, resample=BICUBIC, expand=False)

    # ── Composite with feathered circular mask ──
    feather = max(4, target_size // 200)
    mask = _disk_mask(target_size, feather_px=feather)

    paste_x = int(round(out_moon_cx - target_size / 2))
    paste_y = int(round(out_moon_cy - target_size / 2))
    # Clamp so the disk stays fully on canvas (recenter if needed)
    paste_x = max(0, min(paste_x, output_size[0] - target_size))
    paste_y = max(0, min(paste_y, output_size[1] - target_size))
    bg.paste(ref_rot, (paste_x, paste_y), mask)

    # ── Caption — below the moon disk, not overlapping it ──
    if caption:
        draw = ImageDraw.Draw(bg)
        caption_y = min(paste_y + target_size + 12, output_size[1] - 30)
        _draw_caption(draw, caption, (22, caption_y), output_size[0])

    Path(out_path).parent.mkdir(parents=True, exist_ok=True)
    bg.save(out_path, quality=92)
    return out_path


def render_phase_closeup(
    nasa_render_path: str,
    out_path: str,
    output_size: tuple[int, int] = (1920, 1080),
    moon_height_pct: float = 0.92,
    caption: str | None = None,
    background: tuple[int, int, int] = (0, 0, 0),
    when_utc: datetime | None = None,
):
    """Full-screen close-up rendering using a NASA SVS Dial-a-Moon image.

    Renders the NASA moon disk centred on a black background, rotated by
    the parallactic angle so its orientation matches what east's camera sees
    from its geographic location at the given UTC time.
    """
    import moon_phase

    bg = Image.new('RGB', output_size, background)
    moon = Image.open(nasa_render_path).convert('RGB')
    moon = _square_crop_to_disk(moon)

    target = int(round(output_size[1] * moon_height_pct))
    target += target % 2
    moon_resized = moon.resize((target, target), LANCZOS)
    moon_resized = moon_resized.filter(
        ImageFilter.UnsharpMask(radius=2, percent=160, threshold=2)
    )

    # Rotate by parallactic angle so "up" on the moon matches east's sky
    if when_utc is not None:
        par = moon_phase.parallactic_angle(when_utc)
        moon_resized = moon_resized.rotate(-par, resample=BICUBIC, expand=False)

    feather = max(3, target // 240)
    mask = _disk_mask(target, feather_px=feather)

    px = (output_size[0] - target) // 2
    py = (output_size[1] - target) // 2
    bg.paste(moon_resized, (px, py), mask)

    if caption:
        draw = ImageDraw.Draw(bg)
        _draw_caption(draw, caption, (22, output_size[1] - 48), output_size[0])

    Path(out_path).parent.mkdir(parents=True, exist_ok=True)
    bg.save(out_path, quality=92)
    return out_path


def _cli():
    p = argparse.ArgumentParser()
    p.add_argument('source')
    p.add_argument('when_utc', help='ISO 8601 UTC, e.g. 2026-04-29T21:07:00Z')
    p.add_argument('ref_moon')
    p.add_argument('out')
    p.add_argument('--caption', default=None)
    p.add_argument('--moon-pct', type=float, default=0.70)
    args = p.parse_args()

    from moon_detect import detect_moon
    det = detect_moon(args.source)
    if det is None:
        print('ERROR: no moon detected in source', file=sys.stderr)
        return 2
    when = datetime.fromisoformat(args.when_utc.replace('Z', '+00:00'))
    composite_full_moon(
        args.source, det, when, args.ref_moon, args.out,
        moon_height_pct=args.moon_pct, caption=args.caption,
    )
    print(f'wrote {args.out}')
    return 0


if __name__ == '__main__':
    sys.exit(_cli())
