#!/usr/bin/env python
"""Compose a REAPER .rpp project from the sample library.

Single entrypoint: loads genre config, builds a SongDefinition from sections,
and writes a .rpp file.

Usage:
    python scripts/compose.py --genre reggaeton --bpm 95 --key Am
    python scripts/compose.py --genre reggaeton --bpm 95 --key Am --output output/my_track.rpp
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

# Ensure project root on 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.rhythm import get_notes, GENERATORS as RHYTHM_GENERATORS
from src.composer.melodic import bass_tresillo, lead_hook, chords_block, pad_sustain
from src.composer.converters import rhythm_to_midi, melodic_to_midi
from src.selector import SampleSelector
from src.reaper_builder import RPPBuilder
from src.reaper_builder.render import render_project


# ---------------------------------------------------------------------------
# VST3 plugin builder helpers (premium plugins)
# ---------------------------------------------------------------------------

# Premium VST3 plugins available:
#   Serum 2 (Xfer Records), Omnisphere (Spectrasonics)
#   FabFilter Pro-Q 3, Pro-C 2, Pro-R 2, Pro-L 2, Saturn 2, Timeless 3
#   The Glue (Cytomic)
#   Valhalla Delay

def serum2() -> PluginDef:
    """Serum 2 synth — used for bass, lead, harmony tracks."""
    return PluginDef(
        name="Serum2",
        path="Serum2.vst3",
        index=0,
    )


def omnisphere() -> PluginDef:
    """Omnisphere — used for pad tracks."""
    return PluginDef(
        name="Omnisphere",
        path="Omnisphere.vst3",
        index=0,
    )


ROLE_MELODIC_GENERATORS = {
    "bass":   bass_tresillo,
    "lead":   lead_hook,
    "harmony": chords_block,
    "pad":    pad_sustain,
}

ROLE_RHYTHM_GENERATORS = {
    "drums":  "kick_main_notes",
    "snare":  "snare_verse_notes",
    "hihat":  "hihat_16th_notes",
    "perc":   "perc_combo_notes",
}

# Roles that use audio items per hit instead of MIDI pattern
AUDIO_ROLES = {"drums", "snare", "hihat", "perc"}

# Role → sample key (used for SampleSelector)
ROLE_TO_SAMPLE_ROLE = {
    "drums":   "kick",
    "snare":   "snare",
    "hihat":   "hihat",
    "perc":    "perc",
    "bass":    "bass",
    "lead":    "lead",
    "harmony": "keys",
    "pad":     "pad",
}


# ---------------------------------------------------------------------------
# Effect chain builder
# ---------------------------------------------------------------------------

# Mapping of effect names to VST3 plugin entries
# Format: (registry_key, filename) tuples
# registry_key must match a key in VST3_REGISTRY for _build_plugin() lookup
_VST3_EFFECTS: dict[str, tuple[str, str]] = {
    "Pro-Q 3":         ("FabFilter Pro-Q 3", "FabFilter Pro-Q 3.vst3"),
    "Pro-C 2":         ("FabFilter Pro-C 2", "FabFilter Pro-C 2.vst3"),
    "Pro-R 2":         ("FabFilter Pro-R 2", "FabFilter Pro-R 2.vst3"),
    "Timeless 3":      ("FabFilter Timeless 3", "FabFilter Timeless 3.vst3"),
    "Saturn 2":        ("FabFilter Saturn 2", "FabFilter Saturn 2.vst3"),
    "Pro-L 2":         ("FabFilter Pro-L 2", "FabFilter Pro-L 2.vst3"),
    "The Glue":        ("The Glue", "The Glue.vst3"),
    "Valhalla Delay":  ("Valhalla Delay", "ValhallaDelay.vst3"),
}


def build_fx_chain(role: str, genre_config: dict, track_plugins: list[PluginDef]) -> list[PluginDef]:
    """Build a plugin chain for a role from genre config mix settings.

    Args:
        role: Track role (e.g. "drums", "bass", "lead")
        genre_config: Loaded genre JSON dict
        track_plugins: Already-added plugins (instruments) to skip

    Returns:
        List of PluginDef for the FX chain (effects only, no instruments).
    """
    mix = genre_config.get("mix", {})
    per_role = mix.get("per_role", {}).get(role, {})

    plugins: list[PluginDef] = []
    effects = per_role.get("effects", [])
    for idx, effect_name in enumerate(effects):
        key = effect_name
        # Normalize Fruity* aliases
        if key == "Fruity Parametric EQ 2":
            key = "Pro-Q 3"
        elif key == "Fruity Compressor":
            key = "Pro-C 2"
        elif key == "Fruity Delay 3":
            key = "Timeless 3"
        elif key == "Fruity Reverb 2":
            key = "Pro-R 2"

        vst3_info = _VST3_EFFECTS.get(key)
        if vst3_info:
            registry_key, filename = vst3_info
            plugins.append(PluginDef(
                name=registry_key,
                path=filename,
                index=idx,
            ))

    return plugins


# ---------------------------------------------------------------------------
# Return track builders
# ---------------------------------------------------------------------------

def create_return_tracks() -> list[TrackDef]:
    """Create reverb and delay return tracks.

    Returns:
        [Reverb return TrackDef (FabFilter Pro-R 2),
         Delay return TrackDef (FabFilter Timeless 3)]
    """
    reverb_track = TrackDef(
        name="Reverb",
        volume=0.7,
        pan=0.0,
        clips=[],
        plugins=[PluginDef(
            name="FabFilter Pro-R 2",
            path="FabFilter_Pro_R_2.vst3",
            index=0,
        )],
        send_reverb=0.0,
        send_delay=0.0,
    )
    delay_track = TrackDef(
        name="Delay",
        volume=0.7,
        pan=0.0,
        clips=[],
        plugins=[PluginDef(
            name="FabFilter Timeless 3",
            path="FabFilter_Timeless_3.vst3",
            index=0,
        )],
        send_reverb=0.0,
        send_delay=0.0,
    )
    return [reverb_track, delay_track]


# ---------------------------------------------------------------------------
# Section track builder
# ---------------------------------------------------------------------------

def build_section_tracks(
    genre_config: dict,
    selector: SampleSelector,
    key: str,
    bpm: float,
) -> tuple[list[TrackDef], list[SectionDef]]:
    """Build all tracks from genre config sections.

    Creates one set of tracks per role, with clips per section placed at
    cumulative bar offsets. Applies section energy via velocity_mult and vol_mult.

    Args:
        genre_config: Loaded genre JSON dict
        selector: SampleSelector for sample queries
        key: Musical key (e.g. "Am")
        bpm: BPM for sample selection

    Returns:
        (tracks, sections)
    """
    structure = genre_config.get("structure", {})
    sections_raw = structure.get("sections", [])
    roles = genre_config.get("roles", {})

    # Parse sections into SectionDef list
    sections: list[SectionDef] = []
    for s in sections_raw:
        sections.append(SectionDef(
            name=s.get("name", "unknown"),
            bars=s.get("bars", 4),
            energy=s.get("energy", 0.5),
        ))

    # Compute cumulative bar offsets for section positions
    section_offsets: list[float] = []
    offset = 0.0
    for sec in sections:
        section_offsets.append(offset)
        offset += sec.bars

    # Build one track per role
    tracks: list[TrackDef] = []

    # Track used sample IDs per role for diversity
    used_sample_ids: dict[str, list[str]] = {}

    for role, role_cfg in roles.items():
        sample_role = ROLE_TO_SAMPLE_ROLE.get(role, role)

        # Collect clips for each section
        section_clips: list[ClipDef] = []

        for sec_idx, (section, sec_offset) in enumerate(zip(sections, section_offsets)):
            # Derive velocity and volume multipliers from section energy
            vel_mult = section.energy
            vol_mult = section.energy

            # For audio roles, select a different sample per section
            sample_path = None
            if role in AUDIO_ROLES:
                exclude = used_sample_ids.get(role, [])
                diverse_results = selector.select_diverse(
                    role=sample_role, n=1, exclude=exclude, key=key, bpm=bpm
                )
                if diverse_results:
                    sample = diverse_results[0]
                    sample_path = sample.get("original_path")
                    sample_id = sample.get("file_hash", "")
                    if sample_id:
                        used_sample_ids.setdefault(role, []).append(sample_id)

            if role in ROLE_RHYTHM_GENERATORS:
                gen_name = ROLE_RHYTHM_GENERATORS[role]
                note_dict = get_notes(gen_name, section.bars, velocity_mult=vel_mult)

                # Audio roles: one clip per hit (one-shot samples placed at beat positions)
                if role in AUDIO_ROLES:
                    for bar_offset, bar_notes in note_dict.items():
                        for note_data in bar_notes:
                            note_pos = note_data.get("pos", 0.0)
                            audio_clip = ClipDef(
                                position=sec_offset * 4.0 + bar_offset * 4.0 + note_pos,
                                length=0.5,  # one-shot duration
                                name=f"{section.name.capitalize()} {role.capitalize()}",
                                audio_path=sample_path,
                            )
                            section_clips.append(audio_clip)
                else:
                    # MIDI roles: single clip with all notes
                    midi_notes = rhythm_to_midi(note_dict)
                    clip = ClipDef(
                        position=sec_offset * 4.0,  # bars → beats
                        length=section.bars * 4.0,
                        name=f"{section.name.capitalize()} {role.capitalize()}",
                        midi_notes=midi_notes,
                    )
                    section_clips.append(clip)
            elif role in ROLE_MELODIC_GENERATORS:
                gen_fn = ROLE_MELODIC_GENERATORS[role]
                note_list = gen_fn(key=key, bars=section.bars, velocity_mult=vel_mult)
                midi_notes = melodic_to_midi(note_list)
                # Melodic roles use MIDI instruments — no audio_path needed
                clip = ClipDef(
                    position=sec_offset * 4.0,
                    length=section.bars * 4.0,
                    name=f"{section.name.capitalize()} {role.capitalize()}",
                    midi_notes=midi_notes,
                )
                section_clips.append(clip)

        if not section_clips:
            continue

        # Build plugins: instrument (if melodic) + FX chain
        plugins: list[PluginDef] = []

        # Melodic tracks get instrument plugins (Serum 2 or Omnisphere)
        if role in ("bass", "lead", "harmony"):
            plugins.append(serum2())
        elif role == "pad":
            plugins.append(omnisphere())

        # FX chain from genre config (effects only, instruments already added above)
        fx_chain = build_fx_chain(role, genre_config, plugins)
        plugins.extend(fx_chain)

        # Send levels from per_role config
        per_role_cfg = genre_config.get("mix", {}).get("per_role", {}).get(role, {})
        send_reverb = 0.3 if per_role_cfg.get("reverb_on_lead") or per_role_cfg.get("reverb_on_snare") else 0.0
        send_delay = 0.0

        track = TrackDef(
            name=role.capitalize(),
            volume=0.85 * vol_mult,
            pan=0.0,
            color=0,
            clips=section_clips,
            plugins=plugins,
            send_reverb=send_reverb,
            send_delay=send_delay,
        )
        tracks.append(track)

    return tracks, sections


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

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Compose a REAPER .rpp project from the genre config."
    )
    parser.add_argument(
        "--genre",
        default="reggaeton",
        help="Genre (default: reggaeton)",
    )
    parser.add_argument(
        "--bpm",
        type=float,
        default=96.0,
        help="BPM (default: 96)",
    )
    parser.add_argument(
        "--key",
        default="Am",
        help="Musical key (default: Am)",
    )
    parser.add_argument(
        "--output",
        default="output/track.rpp",
        help="Output .rpp path (default: output/track.rpp)",
    )
    parser.add_argument(
        "--render",
        action="store_true",
        help="Render the project to WAV after generating the .rpp file.",
    )
    parser.add_argument(
        "--render-output",
        default=None,
        help="Output WAV path for rendering.",
    )
    args = parser.parse_args()

    # Validate BPM
    if args.bpm <= 0:
        raise ValueError(f"bpm must be > 0, got {args.bpm}")

    # Ensure output directory exists
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    # Load genre config
    genre_path = _ROOT / "knowledge" / "genres" / f"{args.genre.lower()}_2009.json"
    if not genre_path.exists():
        print(f"ERROR: genre config not found at {genre_path}", file=sys.stderr)
        sys.exit(1)

    with open(genre_path, "r", encoding="utf-8") as f:
        genre_config = json.load(f)

    # Load sample index
    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))

    # Build tracks and sections from genre config
    tracks, sections = build_section_tracks(genre_config, selector, args.key, args.bpm)

    # Create return tracks
    return_tracks = create_return_tracks()

    # Assemble SongDefinition
    meta = SongMeta(
        bpm=args.bpm,
        key=args.key,
        title=f"{genre_config.get('display_name', args.genre.capitalize())}",
        time_sig_num=genre_config.get("time_signature", [4, 4])[0],
        time_sig_den=genre_config.get("time_signature", [4, 4])[1],
        ppq=genre_config.get("ppq", 96),
    )

    song = SongDefinition(
        meta=meta,
        tracks=tracks + return_tracks,
        sections=sections,
    )

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

    # Write .rpp
    builder = RPPBuilder(song)
    builder.write(str(output_path))

    # Render if requested
    if args.render:
        render_output_path = args.render_output
        if render_output_path is None:
            render_output_path = str(output_path).replace('.rpp', '.wav')
        render_project(str(output_path), render_output_path)

    print(str(output_path.resolve()))


if __name__ == "__main__":
    main()