"""
Audio Prompt Guide node for ComfyUI.

Turns an audio clip + your free-text motion notes into two things the (vision-only)
VLM can use to write an audio-aligned LTX beat timeline:

  1. waveform_image (IMAGE) - a rendered energy envelope with beat markers and
     suggested segment boundaries, so the model can *see* the audio's shape.
  2. audio_summary (STRING) - duration, tempo/beats (if librosa is installed),
     a per-segment energy breakdown with snapped frame counts, and your notes,
     formatted as guidance the model can map onto beats.

Wire waveform_image -> the judge node's `reference_image` and audio_summary ->
its `user_prompt` (mode=chat, json_output=true), with your LTX system prompt.

librosa is optional: without it you still get the energy envelope + segments;
with it you also get tempo (BPM) and beat times.
"""

from __future__ import annotations

import numpy as np
import torch
from PIL import Image, ImageDraw


def _to_mono(audio) -> tuple[np.ndarray, int]:
    """ComfyUI AUDIO dict -> (mono float32 samples, sample_rate)."""
    wf = audio["waveform"]
    sr = int(audio["sample_rate"])
    arr = wf.detach().cpu().numpy() if hasattr(wf, "detach") else np.asarray(wf)
    arr = np.asarray(arr, dtype=np.float32)
    while arr.ndim > 2:      # [B, C, N] -> [C, N]
        arr = arr[0]
    if arr.ndim == 2:        # [C, N] -> mono
        arr = arr.mean(axis=0)
    return arr, sr


def _rms_envelope(y: np.ndarray, sr: int, fps_env: int = 100) -> tuple[np.ndarray, np.ndarray]:
    """RMS energy per ~1/fps_env-second frame, normalized 0..1, with frame times."""
    hop = max(1, sr // fps_env)
    n = (len(y) // hop) * hop
    if n < hop:
        return np.array([0.0]), np.array([0.0])
    frames = y[:n].reshape(-1, hop)
    rms = np.sqrt((frames ** 2).mean(axis=1) + 1e-9)
    rms_n = rms / (rms.max() + 1e-9)
    times = np.arange(len(rms_n)) * hop / sr
    return rms_n, times


def _tempo_beats(y: np.ndarray, sr: int):
    """(bpm, beat_times) via librosa if available, else (None, [])."""
    try:
        import librosa
        tempo, beat_frames = librosa.beat.beat_track(y=y.astype(np.float32), sr=sr)
        beats = librosa.frames_to_time(beat_frames, sr=sr).tolist()
        return float(np.atleast_1d(tempo)[0]), beats
    except Exception:
        return None, []


def _snap8(frames: int) -> int:
    """Snap a frame count to LTX's 8n+1 grid (e.g. 120 -> 121)."""
    n = max(0, round((frames - 1) / 8))
    return int(8 * n + 1)


def _segments(rms_n, times, duration, fps, max_segments, beats):
    """Split into ~5s segments; label each by mean energy + snapped frame count."""
    seg_count = int(min(max(3, round(duration / 5.0)), max(3, max_segments)))
    bounds = np.linspace(0.0, duration, seg_count + 1)
    stages = ["establish", "build", "peak", "settle"]
    segs = []
    for i in range(seg_count):
        t0, t1 = float(bounds[i]), float(bounds[i + 1])
        mask = (times >= t0) & (times < t1)
        e = float(rms_n[mask].mean()) if mask.any() else 0.0
        peak_t = float(times[mask][np.argmax(rms_n[mask])]) if mask.any() else t0
        dur = round(t1 - t0, 2)
        label = "high" if e > 0.66 else ("medium" if e > 0.33 else "low")
        stage = stages[i] if i < len(stages) else ("peak" if e > 0.6 else "settle")
        segs.append({
            "segment": i + 1, "start_s": round(t0, 2), "duration_s": dur,
            "frames": _snap8(round(dur * fps)), "energy": label,
            "energy_val": round(e, 3), "peak_s": round(peak_t, 2),
            "stage_hint": stage,
            "beats_in": [round(b, 2) for b in beats if t0 <= b < t1],
        })
    return segs, bounds


def _render(rms_n, times, duration, beats, bounds, segs):
    """Render the envelope + beats + segment boundaries to a ComfyUI IMAGE tensor."""
    W, H = 1024, 256
    img = Image.new("RGB", (W, H), (18, 18, 22))
    d = ImageDraw.Draw(img)
    dur = max(duration, 1e-6)

    def X(t):
        return int(max(0, min(W - 1, t / dur * (W - 1))))

    pts = [(0, H)] + [(X(times[i]), H - int(rms_n[i] * (H - 26))) for i in range(len(rms_n))] + [(W - 1, H)]
    d.polygon(pts, fill=(60, 140, 220))
    for b in beats:                                   # beat markers (orange)
        d.line([(X(b), 0), (X(b), H)], fill=(230, 110, 60), width=1)
    for bd in bounds:                                 # segment boundaries (white)
        d.line([(X(bd), 0), (X(bd), H)], fill=(240, 240, 240), width=1)
    for s in segs:
        d.text((X(s["start_s"]) + 4, 4), f"S{s['segment']} {s['energy']}", fill=(255, 255, 255))
    arr = np.asarray(img, dtype=np.float32) / 255.0
    return torch.from_numpy(arr)[None, ...]           # [1, H, W, 3]


def _summary(duration, sr, bpm, beats, segs, notes):
    lines = ["AUDIO GUIDE — align the video beats to this audio:",
             f"- duration: {duration:.2f}s  |  sample_rate: {sr} Hz"]
    if bpm:
        lines.append(f"- tempo: ~{bpm:.0f} BPM")
    if beats:
        preview = ", ".join(f"{b:.2f}" for b in beats[:24])
        lines.append(f"- beat times (s): {preview}{' ...' if len(beats) > 24 else ''}")
    lines.append("- suggested segments (use these durations/frames and energy):")
    for s in segs:
        peak = f", energy peak @ {s['peak_s']}s" if s["energy"] == "high" else ""
        lines.append(
            f"    seg{s['segment']}: {s['start_s']}–{s['start_s'] + s['duration_s']:.2f}s, "
            f"{s['duration_s']}s, {s['frames']} frames, energy {s['energy'].upper()} "
            f"({s['stage_hint']}){peak}")
    lines.append("")
    lines.append("Guidance: higher energy -> faster motion and bigger camera moves; "
                 "lower energy -> slower, settle. Put escalation on rising energy and the "
                 "release on the final segment. Snap each beat's frames to 8n+1.")
    if notes.strip():
        lines.append("")
        lines.append("USER NOTES (map these onto the matching segments/times):")
        lines.append(notes.strip())
    return "\n".join(lines)


class AudioPromptGuide:
    CATEGORY = "prompt_calibrator"
    FUNCTION = "guide"
    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("waveform_image", "audio_summary")

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "audio": ("AUDIO",),
                "notes": ("STRING", {"default": "", "multiline": True}),
                "fps": ("INT", {"default": 24, "min": 1, "max": 120}),
                "max_segments": ("INT", {"default": 6, "min": 3, "max": 12}),
            },
        }

    def guide(self, audio, notes, fps, max_segments):
        y, sr = _to_mono(audio)
        duration = len(y) / sr if sr else 0.0
        rms_n, times = _rms_envelope(y, sr)
        bpm, beats = _tempo_beats(y, sr)
        segs, bounds = _segments(rms_n, times, duration, fps, max_segments, beats)
        image = _render(rms_n, times, duration, beats, bounds, segs)
        summary = _summary(duration, sr, bpm, beats, segs, notes)
        return (image, summary)


NODE_CLASS_MAPPINGS = {"AudioPromptGuide": AudioPromptGuide}
NODE_DISPLAY_NAME_MAPPINGS = {"AudioPromptGuide": "SxCP Audio Prompt Guide"}