#!/usr/bin/env python3
"""
add_stems.py — Add a new song scene with stems to LinkinPark.als

Usage:
    Edit the SONG_CONFIG at the bottom of this file, then run:
        python3 add_stems.py

The script:
 1. Reads LinkinPark Project/LinkinPark.xml (plain XML — run unpack.sh first if needed)
 2. Adds a new Scene with BPM
 3. For each stem track (Id 100–116), appends one ClipSlot:
      - filled (AudioClip with warp markers) if the stem file is mapped
      - empty otherwise
 4. Writes back to LinkinPark.xml (run build.sh afterward to create .als)

Track IDs:
  100=Kick  101=Snare  102=Hi-Hat  103=Cymbals  104=Toms  105=Other Kit
  106=Bass  107=Rhythm Guitar  108=Lead Guitar
  109=Vocals  110=Backing Vocals
  111=Keys  112=Piano  113=Strings  114=Wind  115=Other  116=Metronome

Group Track IDs (Ableton track groups):
  200=Drums      (contains 100–105, color=1/orange)
  201=Bass       (contains 106, color=11/cyan)
  202=Guitars    (contains 107–108, color=8/green)
  203=Vocals     (contains 109–110, color=5/blue)
  204=Keys & Synth (contains 111–116, color=17/yellow)
"""

import gzip, re, os, subprocess
from xml.etree import ElementTree as ET

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

XML_PATH = os.path.join(os.path.dirname(__file__),
                        "LinkinPark Project", "LinkinPark.xml")

STEM_TRACK_IDS = list(range(100, 117))  # 100..116


def get_duration(path):
    """Return duration in seconds via ffprobe."""
    result = subprocess.run(
        ["ffprobe", "-v", "quiet", "-show_entries", "format=duration",
         "-of", "csv=p=0", path],
        capture_output=True, text=True)
    return float(result.stdout.strip().split(",")[0])


def get_sample_rate(path):
    """Return sample rate via ffprobe."""
    result = subprocess.run(
        ["ffprobe", "-v", "quiet", "-show_entries", "stream=sample_rate",
         "-of", "csv=p=0", path],
        capture_output=True, text=True)
    lines = [l for l in result.stdout.strip().splitlines() if l.strip().isdigit()]
    return int(lines[0]) if lines else 48000


def get_file_size(path):
    return os.path.getsize(path)


# ---------------------------------------------------------------------------
# XML builders
# ---------------------------------------------------------------------------

def scene_xml(scene_id, name, bpm, color=0):
    return f'''\t\t\t<Scene Id="{scene_id}">
\t\t\t\t<FollowAction>
\t\t\t\t\t<FollowTime Value="4" />
\t\t\t\t\t<IsLinked Value="true" />
\t\t\t\t\t<LoopIterations Value="1" />
\t\t\t\t\t<FollowActionA Value="4" />
\t\t\t\t\t<FollowActionB Value="0" />
\t\t\t\t\t<FollowChanceA Value="100" />
\t\t\t\t\t<FollowChanceB Value="0" />
\t\t\t\t\t<JumpIndexA Value="0" />
\t\t\t\t\t<JumpIndexB Value="0" />
\t\t\t\t\t<FollowActionEnabled Value="false" />
\t\t\t\t</FollowAction>
\t\t\t\t<Name Value="{name}" />
\t\t\t\t<Annotation Value="" />
\t\t\t\t<Color Value="{color}" />
\t\t\t\t<Tempo>
\t\t\t\t\t<LomId Value="0" />
\t\t\t\t\t<Manual Value="{bpm}" />
\t\t\t\t\t<MidiControllerRange>
\t\t\t\t\t\t<Min Value="60" />
\t\t\t\t\t\t<Max Value="200" />
\t\t\t\t\t</MidiControllerRange>
\t\t\t\t\t<AutomationTarget Id="0" />
\t\t\t\t\t<ModulationTarget Id="0" />
\t\t\t\t</Tempo>
\t\t\t\t<IsTempoEnabled Value="true" />
\t\t\t\t<TimeSignature>
\t\t\t\t\t<TimeSignatures>
\t\t\t\t\t\t<RemoteableTimeSignature Id="0">
\t\t\t\t\t\t\t<Numerator Value="4" />
\t\t\t\t\t\t\t<Denominator Value="4" />
\t\t\t\t\t\t\t<Time Value="0" />
\t\t\t\t\t\t</RemoteableTimeSignature>
\t\t\t\t\t</TimeSignatures>
\t\t\t\t</TimeSignature>
\t\t\t\t<IsTimeSignatureEnabled Value="false" />
\t\t\t\t<LomId Value="0" />
\t\t\t\t<ClipSlotsListWrapper LomId="0" />
\t\t\t</Scene>
'''


def empty_clipslot_xml(slot_id):
    return f'''\t\t\t\t\t\t\t<ClipSlot Id="{slot_id}">
\t\t\t\t\t\t\t\t<LomId Value="0" />
\t\t\t\t\t\t\t\t<ClipSlot>
\t\t\t\t\t\t\t\t\t<Value />
\t\t\t\t\t\t\t\t</ClipSlot>
\t\t\t\t\t\t\t\t<HasStop Value="true" />
\t\t\t\t\t\t\t\t<NeedRefreeze Value="true" />
\t\t\t\t\t\t\t</ClipSlot>
'''


def filled_clipslot_xml(slot_id, clip_id, name, bpm, rel_path, abs_path,
                         duration_sec, file_size, sample_rate=48000, color=16):
    # Escape & in paths so XML stays valid (e.g. "Guitar & Piano.flac")
    rel_path = rel_path.replace('&', '&amp;')
    abs_path = abs_path.replace('&', '&amp;')
    beats = duration_sec * bpm / 60.0
    beat_dur = 60.0 / bpm
    default_duration = int(round(duration_sec * sample_rate))

    return f'''\t\t\t\t\t\t\t<ClipSlot Id="{slot_id}">
\t\t\t\t\t\t\t\t<LomId Value="0" />
\t\t\t\t\t\t\t\t<ClipSlot>
\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t<Value>
\t\t\t\t\t\t\t<AudioClip Id="{clip_id}" Time="0">
\t\t\t\t\t\t\t\t<LomId Value="0" />
\t\t\t\t\t\t\t\t<LomIdView Value="0" />
\t\t\t\t\t\t\t\t<CurrentStart Value="0" />
\t\t\t\t\t\t\t\t<CurrentEnd Value="{beats:.4f}" />
\t\t\t\t\t\t\t\t<Loop>
\t\t\t\t\t\t\t\t\t<LoopStart Value="0" />
\t\t\t\t\t\t\t\t\t<LoopEnd Value="{beats:.4f}" />
\t\t\t\t\t\t\t\t\t<StartRelative Value="0" />
\t\t\t\t\t\t\t\t\t<LoopOn Value="false" />
\t\t\t\t\t\t\t\t\t<OutMarker Value="{beats:.4f}" />
\t\t\t\t\t\t\t\t\t<HiddenLoopStart Value="0" />
\t\t\t\t\t\t\t\t\t<HiddenLoopEnd Value="{beats:.4f}" />
\t\t\t\t\t\t\t\t</Loop>
\t\t\t\t\t\t\t\t<Name Value="{name}" />
\t\t\t\t\t\t\t\t<Annotation Value="" />
\t\t\t\t\t\t\t\t<Color Value="{color}" />
\t\t\t\t\t\t\t\t<LaunchMode Value="0" />
\t\t\t\t\t\t\t\t<LaunchQuantisation Value="0" />
\t\t\t\t\t\t\t\t<TimeSignature>
\t\t\t\t\t\t\t\t\t<TimeSignatures>
\t\t\t\t\t\t\t\t\t\t<RemoteableTimeSignature Id="0">
\t\t\t\t\t\t\t\t\t\t\t<Numerator Value="4" />
\t\t\t\t\t\t\t\t\t\t\t<Denominator Value="4" />
\t\t\t\t\t\t\t\t\t\t\t<Time Value="0" />
\t\t\t\t\t\t\t\t\t\t</RemoteableTimeSignature>
\t\t\t\t\t\t\t\t\t</TimeSignatures>
\t\t\t\t\t\t\t\t</TimeSignature>
\t\t\t\t\t\t\t\t<Envelopes>
\t\t\t\t\t\t\t\t\t<Envelopes />
\t\t\t\t\t\t\t\t</Envelopes>
\t\t\t\t\t\t\t\t<ScrollerTimePreserver>
\t\t\t\t\t\t\t\t\t<LeftTime Value="0" />
\t\t\t\t\t\t\t\t\t<RightTime Value="0" />
\t\t\t\t\t\t\t\t</ScrollerTimePreserver>
\t\t\t\t\t\t\t\t<TimeSelection>
\t\t\t\t\t\t\t\t\t<AnchorTime Value="0" />
\t\t\t\t\t\t\t\t\t<OtherTime Value="0" />
\t\t\t\t\t\t\t\t</TimeSelection>
\t\t\t\t\t\t\t\t<Legato Value="false" />
\t\t\t\t\t\t\t\t<Ram Value="false" />
\t\t\t\t\t\t\t\t<GrooveSettings>
\t\t\t\t\t\t\t\t\t<GrooveId Value="-1" />
\t\t\t\t\t\t\t\t</GrooveSettings>
\t\t\t\t\t\t\t\t<Disabled Value="false" />
\t\t\t\t\t\t\t\t<VelocityAmount Value="0" />
\t\t\t\t\t\t\t\t<FollowAction>
\t\t\t\t\t\t\t\t\t<FollowTime Value="4" />
\t\t\t\t\t\t\t\t\t<IsLinked Value="true" />
\t\t\t\t\t\t\t\t\t<LoopIterations Value="1" />
\t\t\t\t\t\t\t\t\t<FollowActionA Value="4" />
\t\t\t\t\t\t\t\t\t<FollowActionB Value="0" />
\t\t\t\t\t\t\t\t\t<FollowChanceA Value="100" />
\t\t\t\t\t\t\t\t\t<FollowChanceB Value="0" />
\t\t\t\t\t\t\t\t\t<JumpIndexA Value="0" />
\t\t\t\t\t\t\t\t\t<JumpIndexB Value="0" />
\t\t\t\t\t\t\t\t\t<FollowActionEnabled Value="false" />
\t\t\t\t\t\t\t\t</FollowAction>
\t\t\t\t\t\t\t\t<Grid>
\t\t\t\t\t\t\t\t\t<FixedNumerator Value="1" />
\t\t\t\t\t\t\t\t\t<FixedDenominator Value="16" />
\t\t\t\t\t\t\t\t\t<GridIntervalPixel Value="20" />
\t\t\t\t\t\t\t\t\t<Ntoles Value="2" />
\t\t\t\t\t\t\t\t\t<SnapToGrid Value="true" />
\t\t\t\t\t\t\t\t\t<Fixed Value="false" />
\t\t\t\t\t\t\t\t</Grid>
\t\t\t\t\t\t\t\t<FreezeStart Value="0" />
\t\t\t\t\t\t\t\t<FreezeEnd Value="0" />
\t\t\t\t\t\t\t\t<IsWarped Value="true" />
\t\t\t\t\t\t\t\t<TakeId Value="1" />
\t\t\t\t\t\t\t\t<IsInKey Value="false" />
\t\t\t\t\t\t\t\t<ScaleInformation>
\t\t\t\t\t\t\t\t\t<Root Value="0" />
\t\t\t\t\t\t\t\t\t<Name Value="0" />
\t\t\t\t\t\t\t\t</ScaleInformation>
\t\t\t\t\t\t\t\t<SampleRef>
\t\t\t\t\t\t\t\t\t<FileRef>
\t\t\t\t\t\t\t\t\t\t<RelativePathType Value="3" />
\t\t\t\t\t\t\t\t\t\t<RelativePath Value="{rel_path}" />
\t\t\t\t\t\t\t\t\t\t<Path Value="{abs_path}" />
\t\t\t\t\t\t\t\t\t\t<Type Value="1" />
\t\t\t\t\t\t\t\t\t\t<LivePackName Value="" />
\t\t\t\t\t\t\t\t\t\t<LivePackId Value="" />
\t\t\t\t\t\t\t\t\t\t<OriginalFileSize Value="{file_size}" />
\t\t\t\t\t\t\t\t\t\t<OriginalCrc Value="0" />
\t\t\t\t\t\t\t\t\t\t<SourceHint Value="" />
\t\t\t\t\t\t\t\t\t</FileRef>
\t\t\t\t\t\t\t\t\t<LastModDate Value="0" />
\t\t\t\t\t\t\t\t\t<SourceContext />
\t\t\t\t\t\t\t\t\t<SampleUsageHint Value="0" />
\t\t\t\t\t\t\t\t\t<DefaultDuration Value="{default_duration}" />
\t\t\t\t\t\t\t\t\t<DefaultSampleRate Value="{sample_rate}" />
\t\t\t\t\t\t\t\t\t<SamplesToAutoWarp Value="0" />
\t\t\t\t\t\t\t\t</SampleRef>
\t\t\t\t\t\t\t\t<Onsets>
\t\t\t\t\t\t\t\t\t<UserOnsets />
\t\t\t\t\t\t\t\t\t<HasUserOnsets Value="false" />
\t\t\t\t\t\t\t\t</Onsets>
\t\t\t\t\t\t\t\t<WarpMode Value="0" />
\t\t\t\t\t\t\t\t<GranularityTones Value="30" />
\t\t\t\t\t\t\t\t<GranularityTexture Value="65" />
\t\t\t\t\t\t\t\t<FluctuationTexture Value="25" />
\t\t\t\t\t\t\t\t<TransientResolution Value="6" />
\t\t\t\t\t\t\t\t<TransientLoopMode Value="2" />
\t\t\t\t\t\t\t\t<TransientEnvelope Value="100" />
\t\t\t\t\t\t\t\t<ComplexProFormants Value="100" />
\t\t\t\t\t\t\t\t<ComplexProEnvelope Value="128" />
\t\t\t\t\t\t\t\t<Sync Value="true" />
\t\t\t\t\t\t\t\t<HiQ Value="false" />
\t\t\t\t\t\t\t\t<Fade Value="true" />
\t\t\t\t\t\t\t\t<Fades>
\t\t\t\t\t\t\t\t\t<FadeInLength Value="0" />
\t\t\t\t\t\t\t\t\t<FadeOutLength Value="0" />
\t\t\t\t\t\t\t\t\t<ClipFadesAreInitialized Value="true" />
\t\t\t\t\t\t\t\t\t<CrossfadeInState Value="0" />
\t\t\t\t\t\t\t\t\t<FadeInCurveSkew Value="0" />
\t\t\t\t\t\t\t\t\t<FadeInCurveSlope Value="0" />
\t\t\t\t\t\t\t\t\t<FadeOutCurveSkew Value="0" />
\t\t\t\t\t\t\t\t\t<FadeOutCurveSlope Value="0" />
\t\t\t\t\t\t\t\t\t<IsDefaultFadeIn Value="true" />
\t\t\t\t\t\t\t\t\t<IsDefaultFadeOut Value="true" />
\t\t\t\t\t\t\t\t</Fades>
\t\t\t\t\t\t\t\t<PitchCoarse Value="0" />
\t\t\t\t\t\t\t\t<PitchFine Value="0" />
\t\t\t\t\t\t\t\t<SampleVolume Value="1" />
\t\t\t\t\t\t\t\t<WarpMarkers>
\t\t\t\t\t\t\t\t\t\t<WarpMarker Id="0" SecTime="0" BeatTime="0" />
\t\t\t\t\t\t\t\t\t\t<WarpMarker Id="1" SecTime="{beat_dur:.10f}" BeatTime="1" />
\t\t\t\t\t\t\t\t\t</WarpMarkers>
\t\t\t\t\t\t\t\t<SavedWarpMarkersForStretched />
\t\t\t\t\t\t\t\t<MarkersGenerated Value="false" />
\t\t\t\t\t\t\t\t<IsSongTempoLeader Value="false" />
\t\t\t\t\t\t\t</AudioClip>
\t\t\t\t\t\t\t\t\t</Value>
\t\t\t\t\t\t\t\t</ClipSlot>
\t\t\t\t\t\t\t\t<HasStop Value="true" />
\t\t\t\t\t\t\t\t<NeedRefreeze Value="true" />
\t\t\t\t\t\t\t</ClipSlot>
'''


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

def add_song(song_config):
    """
    song_config = {
        "name": "Somewhere I Belong",
        "bpm": 162,
        "color": 16,                      # scene/clip color (-1 = default)
        "stems_dir": "Samples/Somewhere I Belong",  # relative to project root
        # track_id -> (filename_stem_without_ext, clip_name)
        "stems": {
            100: ("01 - Kick Drum", "kick"),
            101: ("02 - Snare Drum", "snare"),
            ...
        }
    }
    """
    project_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                "LinkinPark Project")
    xml_path = os.path.join(project_dir, "LinkinPark.xml")

    print(f"Reading {xml_path}...")
    with open(xml_path, "r", encoding="utf-8") as f:
        content = f.read()

    # --- Current NextPointeeId ---
    m = re.search(r'NextPointeeId Value="(\d+)"', content)
    next_id = int(m.group(1))
    print(f"  NextPointeeId: {next_id}")

    # --- Current scene count ---
    scene_ids = re.findall(r'<Scene Id="(\d+)"', content)
    new_scene_id = len(scene_ids)
    print(f"  Current scenes: {len(scene_ids)} → new scene Id={new_scene_id}")

    # --- Insert new scene into <Scenes> ---
    new_scene = scene_xml(new_scene_id, song_config["name"],
                          song_config["bpm"], song_config.get("color", 0))
    content = content.replace("</Scenes>", new_scene + "\t\t\t</Scenes>")

    # --- For each stem track: append one ClipSlot ---
    stems_dir_rel = song_config["stems_dir"]
    stems_dir_abs = os.path.join(project_dir, stems_dir_rel)
    bpm = song_config["bpm"]
    stems_map = song_config["stems"]  # track_id -> (filename_no_ext, clip_name)

    # Include default track 18 (always empty)
    all_audio_track_ids = [18] + STEM_TRACK_IDS

    for track_id in all_audio_track_ids:
        pattern = f'<AudioTrack Id="{track_id}"'
        t_start = content.find(pattern)
        if t_start == -1:
            print(f"  WARNING: track {track_id} not found, skipping")
            continue

        track_end = content.find("</AudioTrack>", t_start)
        track_xml = content[t_start:track_end]

        # Find ALL </ClipSlotList> positions within this track
        cs_ends = []
        search_from = t_start
        while True:
            pos = content.find("</ClipSlotList>", search_from, track_end)
            if pos == -1:
                break
            cs_ends.append(pos)
            search_from = pos + 1

        if not cs_ends:
            print(f"  WARNING: no </ClipSlotList> found for track {track_id}")
            continue

        # Slot id = number of existing slots in the first (MainSequencer) list
        existing_slots = len(re.findall(r'<ClipSlot Id="\d+">', track_xml))
        slot_id = existing_slots

        if track_id in stems_map:
            filename_no_ext, clip_name = stems_map[track_id]
            flac_path = os.path.join(stems_dir_abs, filename_no_ext + ".flac")
            if not os.path.exists(flac_path):
                print(f"  WARNING: {flac_path} not found, inserting empty slot")
                new_slot = empty_clipslot_xml(slot_id)
            else:
                dur = get_duration(flac_path)
                sr = get_sample_rate(flac_path)
                size = get_file_size(flac_path)
                rel = stems_dir_rel + "/" + filename_no_ext + ".flac"
                abs_p = flac_path
                clip_id = next_id
                next_id += 1
                print(f"  Track {track_id}: clip_id={clip_id} '{clip_name}' "
                      f"dur={dur:.2f}s beats={dur*bpm/60:.2f}")
                new_slot = filled_clipslot_xml(
                    slot_id, clip_id, clip_name, bpm,
                    rel, abs_p, dur, size, sr,
                    song_config.get("color", 16))
        else:
            new_slot = empty_clipslot_xml(slot_id)

        # Insert filled/empty slot into first ClipSlotList (MainSequencer),
        # then insert empty slot into remaining ClipSlotLists (FreezeSequencer).
        # Process in order, tracking the running offset shift from each insertion.
        shift = 0
        for i, cs_end in enumerate(cs_ends):
            slot = new_slot if i == 0 else empty_clipslot_xml(slot_id)
            pos = cs_end + shift
            content = content[:pos] + slot + content[pos:]
            shift += len(slot)

    # --- Add GroupTrackSlot + FreezeSequencer ClipSlot to each GroupTrack ---
    group_track_ids = re.findall(r'<GroupTrack Id="(\d+)"', content)
    for gid in group_track_ids:
        t_start = content.find(f'<GroupTrack Id="{gid}"')
        t_end   = content.find('</GroupTrack>', t_start)
        gt_body = content[t_start:t_end]

        if '<Slots>' not in gt_body:
            print(f"  WARNING: GroupTrack {gid} has no <Slots> section — skipping")
            continue

        # 1. Add GroupTrackSlot to <Slots>
        existing_slots = len(re.findall(r'<GroupTrackSlot Id=', gt_body))
        gts_xml = (f'\t\t\t\t<GroupTrackSlot Id="{existing_slots}">\n'
                   f'\t\t\t\t\t<LomId Value="0" />\n'
                   f'\t\t\t\t</GroupTrackSlot>\n')
        slots_end = content.find('</Slots>', t_start)
        content = content[:slots_end] + gts_xml + '\t\t\t' + content[slots_end:]
        print(f"  GroupTrack {gid}: added GroupTrackSlot Id={existing_slots}")

        # 2. Add empty ClipSlot to FreezeSequencer's ClipSlotList
        # Re-find t_start/t_end since content shifted
        t_start = content.find(f'<GroupTrack Id="{gid}"')
        t_end   = content.find('</GroupTrack>', t_start)
        existing_clip_slots = len(re.findall(r'<ClipSlot Id="\d+">', content[t_start:t_end]))
        cs_end = content.find('</ClipSlotList>', t_start)
        if cs_end != -1 and cs_end < t_end:
            content = content[:cs_end] + empty_clipslot_xml(existing_clip_slots) + content[cs_end:]

    # --- Update NextPointeeId ---
    content = content.replace(
        f'NextPointeeId Value="{int(m.group(1))}"',
        f'NextPointeeId Value="{next_id}"')

    # --- Validate & write back ---
    ET.fromstring(content.encode("utf-8"))  # raises if invalid XML
    print(f"Writing {xml_path}...")
    with open(xml_path, "w", encoding="utf-8") as f:
        f.write(content)
    print(f"Done! NextPointeeId now {next_id}")
    print("Run ./build.sh to create LinkinPark.als for Ableton.")

    # Verify
    with open(xml_path, "r", encoding="utf-8") as f:
        check = f.read()
    clips = re.findall(r"<AudioClip Id=", check)
    scenes = re.findall(r'<Name Value="[^"]+" />', check)
    print(f"Verification: {len(clips)} AudioClips in project")


# ===========================================================================
# SONG CONFIGURATION — edit this to add a new song
# ===========================================================================

if __name__ == "__main__":
    SONG_CONFIG = {
        "name": "Where'd You Go",
        "bpm": 80,
        "color": 16,
        "stems_dir": "Samples/The Rising Tied/Where'd You Go",
        "stems": {
            105: ("drums", "drums"),
            106: ("bass", "bass"),
            107: ("guitar", "guitar"),
            109: ("vocals", "vocals"),
            112: ("piano", "piano"),
            115: ("other", "other"),
        }
    }

    add_song(SONG_CONFIG)