#!/usr/bin/env python
"""Drumloop-first REAPER .rpp project generator for reggaeton instrumental.

The drumloop drives EVERYTHING: BPM, key, rhythm all come from analysis.
Bass, chords, lead, and pad are built to sync with the drumloop's rhythm.
NO vocals — this is an instrumental-only generator.

Usage:
    python scripts/compose.py --output output/drumloop_v2.rpp
    python scripts/compose.py --bpm 95 --key Am --output output/song.rpp
    python scripts/compose.py --bpm 95 --key Am --seed 42 --output output/song.rpp
"""
from __future__ import annotations

import argparse
import random
import sys
from pathlib import Path

_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(_ROOT))

from src.core.schema import (
    SongDefinition, SongMeta, TrackDef, ClipDef, MidiNote,
    PluginDef, SectionDef,
)
from src.composer.drum_analyzer import DrumLoopAnalyzer
from src.selector import SampleSelector
from src.reaper_builder import RPPBuilder, PLUGIN_REGISTRY, PLUGIN_PRESETS


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

NOTE_NAMES = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]
NOTE_TO_MIDI = {n: i for i, n in enumerate(NOTE_NAMES)}

ROLE_COLORS = {
    "drumloop": 3,
    "clap": 4,
    "bass": 5,
    "chords": 9,
    "lead": 11,
    "pad": 13,
}

# Section structure: (name, bars, energy, has_clap)
# Clap ONLY on chorus and verse sections
SECTIONS = [
    ("intro",  4, 0.4, False),
    ("verse",  8, 0.6, True),
    ("build",  4, 0.7, False),
    ("chorus", 8, 1.0, True),
    ("break",  4, 0.5, False),
    ("chorus", 8, 1.0, True),
    ("outro",  4, 0.3, False),
]

# Tresillo rhythm positions in beats (within a bar)
TRESILLO_POSITIONS = [0.0, 0.75, 1.5, 2.0, 2.75, 3.5]

# Clap positions in beats (within a bar)
CLAP_POSITIONS = [1.0, 3.5]

# i-VI-III-VII chord progression in semitones from root (minor key)
CHORD_PROGRESSION = [
    (0, "minor"),   # i
    (8, "major"),   # VI
    (3, "major"),   # III
    (10, "major"),  # VII
]

# FX chains per track role (before return sends)
FX_CHAINS = {
    "drumloop": ["Decapitator", "Radiator"],
    "bass":     ["Serum_2", "Decapitator", "Gullfoss_Master"],
    "chords":   ["Omnisphere", "PhaseMistress", "EchoBoy"],
    "lead":     ["Serum_2", "Tremolator"],
    "clap":     ["Decapitator"],
    "pad":      ["Omnisphere", "ValhallaDelay"],
}

# Send levels (reverb, delay) per track role
SEND_LEVELS = {
    "bass":   (0.05, 0.02),
    "chords": (0.15, 0.08),
    "lead":   (0.10, 0.05),
    "clap":   (0.05, 0.02),
    "pad":    (0.25, 0.15),
}

# Track volume levels
VOLUME_LEVELS = {
    "drumloop": 0.85,
    "bass":     0.82,
    "chords":   0.70,
    "lead":     0.75,
    "clap":     0.80,
    "pad":      0.65,
}

# Master volume
MASTER_VOLUME = 0.85


# ---------------------------------------------------------------------------
# Phase 1: Infrastructure
# ---------------------------------------------------------------------------

def score_drumloop(sample: dict, analysis) -> float:
    """Score a drumloop candidate for selection quality.

    Formula: key_confidence*0.4 + onset_density_normalized*0.3 + duration_score*0.2 + balance_score*0.1

    Args:
        sample: sample dict from index (used for duration)
        analysis: DrumLoopAnalysis result

    Returns:
        Composite score 0.0–1.0 (higher = better)
    """
    # key_confidence: already 0-1 from analysis
    kc = analysis.key_confidence

    # onset_density_normalized: normalize against typical max (15.0)
    transients = analysis.transients
    duration = analysis.duration
    onset_density = len(transients) / duration if duration > 0 else 0.0
    onset_density_normalized = min(1.0, onset_density / 15.0)

    # duration_score: prefer >= 8 second loops for clean looping
    dur = sample.get("signal", {}).get("duration", 0.0)
    duration_score = 1.0 if dur >= 8.0 else dur / 8.0

    # balance_score: penalize if kick/snare ratio is lopsided
    kick_count = len(analysis.transients_of_type("kick"))
    snare_count = len(analysis.transients_of_type("snare"))
    total = len(transients) if transients else 1
    kick_ratio = kick_count / total
    snare_ratio = snare_count / total
    balance_score = 2.0 * min(kick_ratio, snare_ratio)
    balance_score = min(1.0, balance_score)

    # format_score: prefer WAV over MP3 (lossless > lossy)
    ext = sample.get("original_path", "").rsplit(".", 1)[-1].lower()
    format_score = 1.0 if ext == "wav" else 0.85

    return kc * 0.3 + onset_density_normalized * 0.25 + duration_score * 0.15 + balance_score * 0.1 + format_score * 0.2


def build_section_structure():
    """Build section list and compute cumulative bar offsets.

    Returns:
        sections: list of SectionDef
        offsets: list of bar offsets (cumulative, in bars)
    """
    sections = [SectionDef(name=n, bars=b, energy=e) for n, b, e, _ in SECTIONS]
    offsets = []
    off = 0.0
    for sec in sections:
        offsets.append(off)
        off += sec.bars
    return sections, offsets


def root_to_midi(root: str, octave: int) -> int:
    """Backward compat: convert note name (e.g. 'C', 'A') to MIDI number."""
    return NOTE_TO_MIDI[root] + (octave + 1) * 12


def key_to_midi_root(key_str: str, octave: int = 2) -> int:
    """Convert key string (e.g. "Am") to MIDI root note number.

    Args:
        key_str: Key like "Am", "Dm", "Gm", "C", "F#m"
        octave: MIDI octave (2 = bass, 3 = chords/pad)

    Returns:
        MIDI note number (e.g. 45 for A2, 57 for A3)
    """
    root = key_str.rstrip("m")
    return NOTE_TO_MIDI[root] + (octave + 1) * 12


# ---------------------------------------------------------------------------
# Music theory helpers
# ---------------------------------------------------------------------------

def parse_key(key_str: str) -> tuple[str, bool]:
    """Parse key string into root and minor flag."""
    if key_str.endswith("m"):
        return key_str[:-1], True
    return key_str, False


def get_pentatonic(root: str, is_minor: bool, octave: int) -> list[int]:
    """Get pentatonic scale pitches for root in given octave."""
    root_midi = key_to_midi_root(root, octave)
    if is_minor:
        intervals = [0, 3, 5, 7, 10]  # minor pentatonic
    else:
        intervals = [0, 2, 4, 7, 9]    # major pentatonic
    return [root_midi + i for i in intervals]


def build_chord(root_midi: int, quality: str) -> list[int]:
    """Build a triad chord from root MIDI note and quality."""
    if quality == "minor":
        return [root_midi, root_midi + 3, root_midi + 7]
    return [root_midi, root_midi + 4, root_midi + 7]


# ---------------------------------------------------------------------------
# Plugin builder
# ---------------------------------------------------------------------------

def make_plugin(registry_key: str, index: int) -> PluginDef:
    """Create a PluginDef from the PLUGIN_REGISTRY."""
    if registry_key in PLUGIN_REGISTRY:
        display, path, uid = PLUGIN_REGISTRY[registry_key]
        preset = PLUGIN_PRESETS.get(registry_key)
        return PluginDef(name=registry_key, path=path, index=index, preset_data=preset)
    return PluginDef(name=registry_key, path=registry_key, index=index)


# ---------------------------------------------------------------------------
# Phase 2: Track Generation
# ---------------------------------------------------------------------------

def build_drumloop_track(drumloop_path: str, total_beats: float) -> TrackDef:
    """Build the drumloop track — single audio clip spanning entire song, looping."""
    clips = [
        ClipDef(
            position=0.0,
            length=total_beats,
            name="Drumloop Full",
            audio_path=drumloop_path,
            loop=True,
        )
    ]
    plugins = [make_plugin(fx, i) for i, fx in enumerate(FX_CHAINS.get("drumloop", []))]
    return TrackDef(
        name="Drumloop",
        volume=VOLUME_LEVELS["drumloop"],
        pan=0.0,
        color=ROLE_COLORS["drumloop"],
        clips=clips,
        plugins=plugins,
    )


def build_bass_track(
    analysis,
    sections: list[SectionDef],
    offsets: list[float],
    key_root: str,
    key_minor: bool,
) -> TrackDef:
    """Build the bass track — MIDI tresillo, filtered by kick_free_zones."""
    root_midi = key_to_midi_root(key_root, 2)
    beat_dur = 60.0 / analysis.bpm
    kfz = analysis.kick_free_zones(margin_beats=0.25)

    def in_kfz(abs_beat: float) -> bool:
        """Check if absolute beat position is in a kick-free zone."""
        s = abs_beat * beat_dur
        return any(zs <= s <= ze for zs, ze in kfz)

    clips = []
    for section, sec_off in zip(sections, offsets):
        vm = section.energy
        notes = []
        for bar in range(section.bars):
            for pos in TRESILLO_POSITIONS:
                abs_beat = sec_off * 4.0 + bar * 4.0 + pos
                if in_kfz(abs_beat):
                    # Note: position within clip is relative to clip start (bar * 4.0)
                    notes.append(MidiNote(
                        pitch=root_midi,
                        start=bar * 4.0 + pos,
                        duration=0.5,
                        velocity=int(100 * vm),
                    ))
        if notes:
            clips.append(ClipDef(
                position=sec_off * 4.0,
                length=section.bars * 4.0,
                name=f"{section.name.capitalize()} Bass",
                midi_notes=notes,
            ))

    plugins = [make_plugin(fx, i) for i, fx in enumerate(FX_CHAINS.get("bass", []))]
    return TrackDef(
        name="Bass",
        volume=VOLUME_LEVELS["bass"],
        pan=0.0,
        color=ROLE_COLORS["bass"],
        clips=clips,
        plugins=plugins,
    )


def build_chords_track(
    analysis,
    sections: list[SectionDef],
    offsets: list[float],
    key_root: str,
    key_minor: bool,
) -> TrackDef:
    """Build the chords track — i-VI-III-VII on downbeats, one clip per section."""
    root_midi = key_to_midi_root(key_root, 3)
    clips = []
    for section, sec_off in zip(sections, offsets):
        vm = section.energy
        notes = []
        for bar in range(section.bars):
            ci = bar % len(CHORD_PROGRESSION)
            interval, quality = CHORD_PROGRESSION[ci]
            for pitch in build_chord(root_midi + interval, quality):
                notes.append(MidiNote(
                    pitch=pitch,
                    start=bar * 4.0,
                    duration=4.0,
                    velocity=int(80 * vm),
                ))
        if notes:
            clips.append(ClipDef(
                position=sec_off * 4.0,
                length=section.bars * 4.0,
                name=f"{section.name.capitalize()} Chords",
                midi_notes=notes,
            ))

    plugins = [make_plugin(fx, i) for i, fx in enumerate(FX_CHAINS.get("chords", []))]
    return TrackDef(
        name="Chords",
        volume=VOLUME_LEVELS["chords"],
        pan=0.0,
        color=ROLE_COLORS["chords"],
        clips=clips,
        plugins=plugins,
    )


def build_lead_track(
    analysis,
    sections: list[SectionDef],
    offsets: list[float],
    key_root: str,
    key_minor: bool,
    seed: int = 42,
) -> TrackDef:
    """Build the lead track — pentatonic melody, avoid transients, chord tones on strong beats."""
    penta_low  = get_pentatonic(key_root, key_minor, 4)
    penta_high = get_pentatonic(key_root, key_minor, 5)
    penta = penta_low + penta_high

    transient_times = [t.time for t in analysis.transients]
    beat_dur = 60.0 / analysis.bpm

    def near_transient(beat: float, margin_beats: float = 0.2) -> bool:
        """Return True if beat position is near a transient."""
        s = beat * beat_dur
        return any(abs(s - tt) < margin_beats * beat_dur for tt in transient_times)

    rng = random.Random(seed)
    clips = []
    for section, sec_off in zip(sections, offsets):
        vm = section.energy
        # Density by section name
        density_map = {"chorus": 0.6, "verse": 0.35, "build": 0.35, "intro": 0.2, "break": 0.2, "outro": 0.15}
        density = density_map.get(section.name, 0.3)

        notes = []
        for bar in range(section.bars):
            for sixteenth in range(16):
                bp = bar * 4.0 + sixteenth * 0.25
                abs_bp = sec_off * 4.0 + bp

                if rng.random() > density:
                    continue
                if near_transient(abs_bp, margin_beats=0.2):
                    continue

                strong = sixteenth in (0, 8)  # beat 0 and beat 2 of bar
                # On strong beats: emphasize chord tones (1st, 3rd, 5th of pentatonic)
                pool = [penta[0], penta[2], penta[4]] if strong else penta
                notes.append(MidiNote(
                    pitch=rng.choice(pool),
                    start=bp,
                    duration=0.5 if strong else 0.25,
                    velocity=int((90 if strong else 70) * vm),
                ))

        if notes:
            clips.append(ClipDef(
                position=sec_off * 4.0,
                length=section.bars * 4.0,
                name=f"{section.name.capitalize()} Lead",
                midi_notes=notes,
            ))

    plugins = [make_plugin(fx, i) for i, fx in enumerate(FX_CHAINS.get("lead", []))]
    return TrackDef(
        name="Lead",
        volume=VOLUME_LEVELS["lead"],
        pan=0.0,
        color=ROLE_COLORS["lead"],
        clips=clips,
        plugins=plugins,
    )


def build_clap_track(
    selector: SampleSelector,
    sections: list[SectionDef],
    offsets: list[float],
) -> TrackDef:
    """Build the clap track — audio snare samples at beats 1.0 and 3.5 ONLY in chorus/verse."""
    # Get clap (snare) samples — select best one
    snare_results = selector.select(role="snare", limit=5)
    clap_path = snare_results[0].sample["original_path"] if snare_results else None

    clips = []
    if clap_path:
        for section, sec_off in zip(sections, offsets):
            if section.name not in ("chorus", "verse"):
                continue
            for bar in range(section.bars):
                for cb in CLAP_POSITIONS:
                    clips.append(ClipDef(
                        position=sec_off * 4.0 + bar * 4.0 + cb,
                        length=0.5,
                        name=f"{section.name.capitalize()} Clap",
                        audio_path=clap_path,
                    ))

    plugins = [make_plugin(fx, i) for i, fx in enumerate(FX_CHAINS.get("clap", []))]
    return TrackDef(
        name="Clap",
        volume=VOLUME_LEVELS["clap"],
        pan=0.0,
        color=ROLE_COLORS["clap"],
        clips=clips,
        plugins=plugins,
    )


def build_pad_track(
    sections: list[SectionDef],
    offsets: list[float],
    key_root: str,
    key_minor: bool,
) -> TrackDef:
    """Build the pad track — sustained root chord, one clip per section."""
    root_midi = key_to_midi_root(key_root, 3)
    quality = "minor" if key_minor else "major"
    chord = build_chord(root_midi, quality)

    clips = []
    for section, sec_off in zip(sections, offsets):
        vm = section.energy
        notes = [
            MidiNote(pitch=p, start=0.0, duration=section.bars * 4.0, velocity=int(60 * vm))
            for p in chord
        ]
        clips.append(ClipDef(
            position=sec_off * 4.0,
            length=section.bars * 4.0,
            name=f"{section.name.capitalize()} Pad",
            midi_notes=notes,
        ))

    plugins = [make_plugin(fx, i) for i, fx in enumerate(FX_CHAINS.get("pad", []))]
    return TrackDef(
        name="Pad",
        volume=VOLUME_LEVELS["pad"],
        pan=0.0,
        color=ROLE_COLORS["pad"],
        clips=clips,
        plugins=plugins,
    )


# ---------------------------------------------------------------------------
# Phase 3: Mixing — Return tracks and sends
# ---------------------------------------------------------------------------

def create_return_tracks() -> list[TrackDef]:
    """Create Reverb and Delay return tracks."""
    return [
        TrackDef(
            name="Reverb",
            volume=0.7,
            pan=0.0,
            clips=[],
            plugins=[make_plugin("FabFilter_Pro-R_2", 0)],
        ),
        TrackDef(
            name="Delay",
            volume=0.7,
            pan=0.0,
            clips=[],
            plugins=[make_plugin("ValhallaDelay", 0)],
        ),
    ]


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Compose a REAPER .rpp project from drumloop analysis — instrumental only."
    )
    parser.add_argument("--bpm", type=float, default=None, help="BPM override")
    parser.add_argument("--key", default=None, help="Key override (e.g. Am)")
    parser.add_argument("--output", default="output/drumloop_v2.rpp", help="Output path")
    parser.add_argument("--seed", type=int, default=None, help="Random seed")
    args = parser.parse_args()

    if args.seed is not None:
        random.seed(args.seed)

    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    # ===== Step 1: Select BEST drumloop (scored, not random) =====
    index_path = _ROOT / "data" / "sample_index.json"
    if not index_path.exists():
        print(f"ERROR: sample index not found at {index_path}", file=sys.stderr)
        sys.exit(1)

    selector = SampleSelector(str(index_path))
    selector._load()

    # Filter drumloops in reggaeton tempo range (85-105 BPM)
    candidates = [
        s for s in selector._samples
        if s["role"] == "drumloop" and 85 <= s["perceptual"]["tempo"] <= 105
    ]
    if not candidates:
        # Fallback: wider range
        candidates = [
            s for s in selector._samples
            if s["role"] == "drumloop" and 80 <= s["perceptual"]["tempo"] <= 120
        ]
    if not candidates:
        print("ERROR: No suitable drumloops found", file=sys.stderr)
        sys.exit(1)

    # Score each candidate and pick the best
    scored_candidates = []
    for c in candidates:
        analysis = DrumLoopAnalyzer(c["original_path"]).analyze()
        c["_score"] = score_drumloop(c, analysis)
        c["_analysis"] = analysis
        scored_candidates.append(c)

    best = max(scored_candidates, key=lambda x: x["_score"])
    drumloop_path = best["original_path"]
    analysis = best["_analysis"]

    print(f"Selected drumloop: {best.get('original_name', drumloop_path)}")
    print(f"  Score: {best['_score']:.3f}")
    print(f"  BPM: {best['perceptual']['tempo']:.1f}, Key: {best['musical']['key']}")
    print(f"  Transients: {len(analysis.transients)} "
          f"(kicks={len(analysis.transients_of_type('kick'))}, "
          f"snares={len(analysis.transients_of_type('snare'))})")

    # ===== Step 2: Project parameters (overrides win) =====
    bpm = args.bpm if args.bpm is not None else analysis.bpm
    key = args.key if args.key is not None else (analysis.key or "Am")
    if bpm <= 0:
        raise ValueError(f"bpm must be > 0, got {bpm}")
    key_root, key_minor = parse_key(key)

    print(f"\nProject: {bpm:.1f} BPM, Key: {key}")

    # ===== Step 3: Build section structure =====
    sections, offsets = build_section_structure()

    # ===== Step 4: Build all tracks =====
    total_beats = sum(s.bars for s in sections) * 4.0

    tracks = [
        build_drumloop_track(drumloop_path, total_beats),
        build_bass_track(analysis, sections, offsets, key_root, key_minor),
        build_chords_track(analysis, sections, offsets, key_root, key_minor),
        build_lead_track(analysis, sections, offsets, key_root, key_minor, seed=args.seed or 42),
        build_clap_track(selector, sections, offsets),
        build_pad_track(sections, offsets, key_root, key_minor),
    ]

    return_tracks = create_return_tracks()
    all_tracks = tracks + return_tracks

    # ===== Step 5: Wire sends =====
    reverb_idx = len(tracks)       # first return track
    delay_idx = len(tracks) + 1    # second return track
    for track in all_tracks:
        if track.name in ("Reverb", "Delay"):
            continue
        role = track.name.lower()
        sends = SEND_LEVELS.get(role, (0.0, 0.0))
        track.send_level = {reverb_idx: sends[0], delay_idx: sends[1]}

    # ===== Step 6: Assemble SongDefinition =====
    meta = SongMeta(bpm=bpm, key=key, title="Reggaeton Drumloop Instrumental")
    song = SongDefinition(
        meta=meta,
        tracks=all_tracks,
        sections=sections,
        master_plugins=["Pro-Q_3", "Pro-C_2", "Pro-L_2"],
    )

    errors = song.validate()
    if errors:
        print("WARNING: validation errors:", file=sys.stderr)
        for e in errors:
            print(f"  - {e}", file=sys.stderr)

    # ===== Step 7: Write RPP =====
    builder = RPPBuilder(song, seed=args.seed)
    builder.write(str(output_path))
    print(f"\nWritten: {output_path.resolve()}")


# Backward compat stubs (used by tests)
EFFECT_ALIASES: dict = {}

def build_section_tracks(*args, **kwargs):
    return [], []


def build_fx_chain(*args, **kwargs):
    return []


def build_sampler_plugin(*args, **kwargs):
    return None


# Alias for renamed function
def build_melody_track(*args, **kwargs):
    """Backward compat alias — use build_lead_track instead."""
    return build_lead_track(*args, **kwargs)


if __name__ == "__main__":
    main()