#!/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 random
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.composer.patterns import generate_structure
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_pattern_bank_notes",
    "snare":  "snare_pattern_bank_notes",
    "hihat":  "hihat_pattern_bank_notes",
    "perc":   "perc_combo_notes",
    "clap":   "clap_24_notes",
}

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

# Role → color index mapping (REAPER color palette 0-67)
ROLE_COLORS: dict[str, int] = {
    "drums":    3,   # red
    "clap":     4,   # orange-red
    "bass":     5,   # orange
    "harmony":  9,   # green
    "lead":    11,   # cyan
    "pad":     13,   # blue
    "perc":     7,   # yellow
    "vocal":   15,  # pink/magenta
    "drumloop": 3,   # red (same as drums)
}

# 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",
    "clap":    "snare",   # clap uses snare samples (no dedicated clap role)
    "vocal":   "vocal",
    "drumloop": "drumloop",
}


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

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


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",
            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",
            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,
    sections_data: list[dict] | None = None,
    humanize: float = 0.3,
    groove_strength: float = 0.3,
    bank_weights: list[tuple[str, float]] | None = None,
) -> 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
        sections_data: List of section dicts with 'name', 'bars', 'energy' keys.
                      If None, falls back to reading 'sections' from genre_config.
        humanize: Humanization amount for melodic generators (0.0-1.0)
        groove_strength: Groove amount for rhythm generators (0.0-1.0)
        bank_weights: List of (bank_name, weight) tuples for weighted random bank selection

    Returns:
        (tracks, sections)
    """
    roles = genre_config.get("roles", {})

    # Fall back to fixed sections from genre config for backward compatibility
    if sections_data is None:
        sections_data = genre_config.get("structure", {}).get("sections", [])

    # Default bank weights for drums — weighted random selection
    if bank_weights is None:
        bank_weights = [
            ("dembow_classico", 3),
            ("dense", 3),
            ("perreo", 2),
            ("trapico", 1),
        ]

    # Add extended roles (clap, vocal, drumloop) if not already in roles
    # These are handled as audio roles with special sample selection
    for _role in ("clap", "vocal", "drumloop"):
        if _role not in roles:
            roles[_role] = {}

    # Parse sections into SectionDef list
    sections: list[SectionDef] = []
    for s in sections_data:
        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]
                # Weighted random bank selection for variation
                bank_names = [b[0] for b in bank_weights]
                bank_weight_values = [b[1] for b in bank_weights]
                bank = random.choices(bank_names, weights=bank_weight_values, k=1)[0]

                note_dict = get_notes(
                    gen_name, section.bars,
                    velocity_mult=vel_mult,
                    bank=bank,
                    groove_strength=groove_strength,
                )

                # 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,
                    section_type=section.name,
                    humanize=humanize,
                )
                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)
            else:
                # vocal, drumloop: audio clips spanning full sections
                if role in ("vocal", "drumloop") and sample_path:
                    # Select character based on section type for vocal
                    if role == "vocal":
                        if section.name in ("verse", "bridge"):
                            character = "melodic"
                        elif section.name in ("chorus", "drop"):
                            character = "powerful"
                        else:
                            character = "neutral"
                        vocal_samples = selector.select_diverse(
                            role="vocal", n=1, exclude=used_sample_ids.get("vocal", []),
                            key=key, bpm=bpm, character=character
                        )
                        if vocal_samples:
                            sample = vocal_samples[0]
                            sample_path = sample.get("original_path")
                            sample_id = sample.get("file_hash", "")
                            if sample_id:
                                used_sample_ids.setdefault("vocal", []).append(sample_id)

                    audio_clip = ClipDef(
                        position=sec_offset * 4.0,
                        length=section.bars * 4.0,
                        name=f"{section.name.capitalize()} {role.capitalize()}",
                        audio_path=sample_path,
                        loop=(role == "drumloop"),
                    )
                    section_clips.append(audio_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

        # Apply role color (REAPER color palette 0-67)
        role_color = ROLE_COLORS.get(role, 0)

        track = TrackDef(
            name=role.capitalize(),
            volume=0.85 * vol_mult,
            pan=0.0,
            color=role_color,
            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.",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Random seed for reproducible output (default: unseeded for max variation).",
    )
    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)

    # Randomize chord progression selection from 5 options in reggaeton_2009.json
    progressions = genre_config.get("chord_progressions", [])
    if progressions:
        # Weighted random selection by popularity
        prog_names = [p.get("name", "") for p in progressions]
        pop_values = [p.get("popularity", 0.5) for p in progressions]
        selected_prog = random.choices(prog_names, weights=pop_values, k=1)[0]
        progression_name = selected_prog
    else:
        progression_name = "i-VII-VI-VII"

    # 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))

    # Generate section structure from template with randomization
    # Note: generate_structure reseeds random internally if seed is provided
    sections_data = generate_structure(genre_config, args.bpm, args.key, seed=args.seed)

    # Build tracks and sections
    tracks, sections = build_section_tracks(
        genre_config, selector, args.key, args.bpm, sections_data,
        humanize=0.3, groove_strength=0.3,
    )

    # 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,
        progression_name=progression_name,
    )

    # Wire sends: find return track indices and set send_level on non-return tracks
    ret_idx_map: dict[str, int] = {}
    for idx, track in enumerate(song.tracks):
        if track.name == "Reverb":
            ret_idx_map["reverb"] = idx
        elif track.name == "Delay":
            ret_idx_map["delay"] = idx

    reverb_idx = ret_idx_map.get("reverb", 0)
    delay_idx = ret_idx_map.get("delay", 1)
    for track in song.tracks:
        if track.name not in ("Reverb", "Delay", "master"):
            track.send_level = {
                reverb_idx: 0.15,  # send to reverb
                delay_idx: 0.10,   # send to delay
            }

    # Wire master chain
    song.master_plugins = ["Pro-Q_3", "Pro-C_2", "Pro-L_2"]

    # 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, seed=args.seed)
    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()