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

The drumloop drives EVERYTHING: BPM, key, rhythm all come from analysis.
Bass, chords, melody, and vocals are built to sync with the drumloop's rhythm.

Usage:
    python scripts/compose.py --output output/song.rpp
    python scripts/compose.py --bpm 95 --key Am --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,
    "melody": 11,
    "pad": 13,
    "vocal": 15,
}

SECTIONS = [
    ("intro", 4, 0.4),
    ("verse", 8, 0.6),
    ("build", 4, 0.7),
    ("chorus", 8, 1.0),
    ("break", 4, 0.5),
    ("chorus", 8, 1.0),
    ("outro", 4, 0.3),
]

TRESILLO_POSITIONS = [0.0, 0.75, 1.5, 2.0, 2.75, 3.5]
CLAP_POSITIONS = [1.0, 3.5]

CHORD_PROGRESSION = [
    (0, "minor"),
    (8, "major"),
    (3, "major"),
    (10, "major"),
]

FX_CHAINS = {
    "drumloop": ["Decapitator", "Radiator"],
    "bass": ["Decapitator", "Gullfoss_Master"],
    "chords": ["PhaseMistress", "EchoBoy"],
    "melody": ["Tremolator"],
    "vocal": ["VC_76", "Radiator", "EchoBoy"],
    "pad": ["ValhallaDelay"],
}

SEND_LEVELS = {
    "bass": (0.05, 0.02),
    "chords": (0.15, 0.08),
    "melody": (0.10, 0.05),
    "vocal": (0.20, 0.10),
    "pad": (0.25, 0.15),
}

VOLUME_LEVELS = {
    "bass": 0.82,
    "drumloop": 0.85,
    "chords": 0.70,
    "melody": 0.75,
    "vocal": 0.80,
    "pad": 0.65,
    "clap": 0.80,
}

# Backward compat stubs for test imports
EFFECT_ALIASES: dict[str, str] = {}


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

def parse_key(key_str: str) -> tuple[str, bool]:
    if key_str.endswith("m"):
        return key_str[:-1], True
    return key_str, False


def root_to_midi(root: str, octave: int) -> int:
    return NOTE_TO_MIDI[root] + (octave + 1) * 12


def get_pentatonic(root: str, is_minor: bool, octave: int) -> list[int]:
    root_midi = root_to_midi(root, octave)
    if is_minor:
        intervals = [0, 3, 5, 7, 10]
    else:
        intervals = [0, 2, 4, 7, 9]
    return [root_midi + i for i in intervals]


def build_chord(root_midi: int, quality: str) -> list[int]:
    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:
    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)


# ---------------------------------------------------------------------------
# Track builders
# ---------------------------------------------------------------------------

def build_drumloop_track(drumloop_path: str, total_beats: float) -> TrackDef:
    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, offsets, key_root, key_minor,
) -> TrackDef:
    root_midi = root_to_midi(key_root, 2)
    beat_dur = 60.0 / analysis.bpm
    kfz = analysis.kick_free_zones(margin_beats=0.25)

    def in_kfz(beat: float) -> bool:
        s = 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):
                    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("Serum_2", 0)]
    plugins += [make_plugin(fx, i + 1) 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, offsets, key_root, key_minor,
) -> TrackDef:
    root_midi = root_to_midi(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("Omnisphere", 0)]
    plugins += [make_plugin(fx, i + 1) 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_melody_track(
    analysis, sections, offsets, key_root, key_minor, seed=42,
) -> TrackDef:
    penta = get_pentatonic(key_root, key_minor, 4) + get_pentatonic(key_root, key_minor, 5)
    transient_times = [t.time for t in analysis.transients]
    beat_dur = 60.0 / analysis.bpm

    def near_transient(beat: float, margin: float = 0.2) -> bool:
        s = beat * beat_dur
        return any(abs(s - tt) < margin * beat_dur for tt in transient_times)

    rng = random.Random(seed)
    clips = []
    for section, sec_off in zip(sections, offsets):
        vm = section.energy
        notes = []
        density = {"chorus": 0.6, "verse": 0.35, "build": 0.35}.get(section.name, 0.2)
        for bar in range(section.bars):
            for sixteenth in range(16):
                bp = bar * 4.0 + sixteenth * 0.25
                if rng.random() > density:
                    continue
                if near_transient(sec_off * 4.0 + bp):
                    continue
                strong = sixteenth in (0, 8)
                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()} Melody", midi_notes=notes,
            ))

    plugins = [make_plugin("Serum_2", 0)]
    plugins += [make_plugin(fx, i + 1) for i, fx in enumerate(FX_CHAINS.get("melody", []))]
    return TrackDef(
        name="Melody", volume=VOLUME_LEVELS["melody"], pan=0.0,
        color=ROLE_COLORS["melody"], clips=clips, plugins=plugins,
    )


def build_vocal_track(
    selector, sections, offsets, key, bpm, analysis,
) -> TrackDef:
    beat_dur = 60.0 / analysis.bpm
    transient_times = sorted(t.time for t in analysis.transients)
    used_ids: list[str] = []
    clips = []

    for section, sec_off in zip(sections, offsets):
        char = "powerful" if section.name == "chorus" else "melodic"
        vs = selector.select_diverse(
            role="vocal", n=1, exclude=used_ids, key=key, bpm=bpm, character=char,
        )
        if not vs:
            continue
        vpath = vs[0]["original_path"]
        sid = vs[0].get("file_hash", "")
        if sid:
            used_ids.append(sid)

        if section.name == "chorus":
            for bar in range(section.bars):
                bar_start = (sec_off * 4.0 + bar * 4.0) * beat_dur
                bar_end = bar_start + 4.0 * beat_dur
                gap_start = bar_start
                for tt in transient_times:
                    if tt < bar_start:
                        continue
                    if tt > bar_end:
                        break
                    if tt - gap_start > 0.08:
                        bp = sec_off * 4.0 + bar * 4.0 + (gap_start - bar_start) / beat_dur
                        lb = min((tt - gap_start) / beat_dur, 2.0)
                        clips.append(ClipDef(
                            position=bp, length=max(lb, 0.5),
                            name=f"{section.name.capitalize()} Vocal",
                            audio_path=vpath,
                        ))
                    gap_start = tt
                if bar_end - gap_start > 0.08:
                    bp = sec_off * 4.0 + bar * 4.0 + (gap_start - bar_start) / beat_dur
                    clips.append(ClipDef(
                        position=bp, length=max((bar_end - gap_start) / beat_dur, 0.5),
                        name=f"{section.name.capitalize()} Vocal", audio_path=vpath,
                    ))
        else:
            for bar in range(0, section.bars, 4):
                clips.append(ClipDef(
                    position=sec_off * 4.0 + bar * 4.0,
                    length=4.0 * min(4, section.bars - bar),
                    name=f"{section.name.capitalize()} Vocal",
                    audio_path=vpath, loop=True,
                ))

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


def build_clap_track(selector, sections, offsets) -> TrackDef:
    clap_results = selector.select(role="snare", limit=5)
    clap_path = clap_results[0].sample["original_path"] if clap_results else None
    clips = []
    if clap_path:
        for section, sec_off in zip(sections, offsets):
            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,
                    ))
    return TrackDef(
        name="Clap", volume=VOLUME_LEVELS["clap"], pan=0.0,
        color=ROLE_COLORS["clap"], clips=clips,
    )


def build_pad_track(sections, offsets, key_root, key_minor) -> TrackDef:
    root_midi = root_to_midi(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("Omnisphere", 0)]
    plugins += [make_plugin(fx, i + 1) 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,
    )


# ---------------------------------------------------------------------------
# Return tracks + backward compat
# ---------------------------------------------------------------------------

def create_return_tracks() -> list[TrackDef]:
    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)],
        ),
    ]


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


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


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


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

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Compose a REAPER .rpp project from drumloop analysis."
    )
    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_song.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 drumloop
    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()
    drumloops = [
        s for s in selector._samples
        if s["role"] == "drumloop" and 85 <= s["perceptual"]["tempo"] <= 105
    ]
    if not drumloops:
        drumloops = [
            s for s in selector._samples
            if s["role"] == "drumloop" and 80 <= s["perceptual"]["tempo"] <= 120
        ]
    if not drumloops:
        print("ERROR: No suitable drumloops found", file=sys.stderr)
        sys.exit(1)

    drumloop = random.choice(drumloops)
    drumloop_path = drumloop["original_path"]
    print(f"Selected drumloop: {drumloop.get('original_name', drumloop_path)}")
    print(f"  BPM: {drumloop['perceptual']['tempo']:.1f}, Key: {drumloop['musical']['key']}")

    # Step 2: Analyze drumloop
    print("Analyzing drumloop...")
    analyzer = DrumLoopAnalyzer(drumloop_path)
    analysis = analyzer.analyze()
    print(f"  Detected BPM: {analysis.bpm:.1f}")
    print(f"  Detected Key: {analysis.key}")
    print(f"  Transients: {len(analysis.transients)} "
          f"(kicks={len(analysis.transients_of_type('kick'))} "
          f"snares={len(analysis.transients_of_type('snare'))} "
          f"hihats={len(analysis.transients_of_type('hihat'))})")

    # Step 3: 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 4: Section structure
    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

    # Step 5: Build 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_melody_track(analysis, sections, offsets, key_root, key_minor, seed=args.seed or 42),
        build_vocal_track(selector, sections, offsets, key, bpm, analysis),
        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 6: Wire sends
    reverb_idx = len(tracks)
    delay_idx = len(tracks) + 1
    for track in all_tracks:
        if track.name not in ("Reverb", "Delay"):
            role = track.name.lower().replace("vocals", "vocal")
            sends = SEND_LEVELS.get(role, (0.0, 0.0))
            track.send_level = {reverb_idx: sends[0], delay_idx: sends[1]}

    # Step 7: Assemble
    meta = SongMeta(bpm=bpm, key=key, title="Reggaeton Drumloop Track")
    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)

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


if __name__ == "__main__":
    main()