import mido
from mido import MidiFile, MidiTrack, Message, MetaMessage
import sys
import os
from collections import defaultdict

class Note:
    def __init__(self, note, start, end, channel, velocity_on, velocity_off):
        self.note = note
        self.start = start
        self.end = end
        self.channel = channel
        self.velocity_on = velocity_on
        self.velocity_off = velocity_off
        self.voice = None

def get_notes_and_events(track):
    """
    Extract notes with start and end times and collect non-note events from a MIDI track.
    Returns a list of Note objects and a list of non-note events.
    Each non-note event is a tuple (absolute_time, message).
    """
    notes = []
    ongoing_notes = {}
    non_note_events = []
    absolute_time = 0
    for msg in track:
        absolute_time += msg.time
        if msg.type == 'note_on' and msg.velocity > 0:
            key = (msg.note, msg.channel)
            if key in ongoing_notes:
                print(f"Warning: Note {msg.note} on channel {msg.channel} started without previous note_off.")
            ongoing_notes[key] = (absolute_time, msg.velocity)
        elif msg.type == 'note_off' or (msg.type == 'note_on' and msg.velocity == 0):
            key = (msg.note, msg.channel)
            if key in ongoing_notes:
                start_time, velocity_on = ongoing_notes.pop(key)
                notes.append(Note(msg.note, start_time, absolute_time, msg.channel, velocity_on, msg.velocity))
            else:
                print(f"Warning: Note {msg.note} on channel {msg.channel} ended without a start.")
        else:
            # Collect non-note events
            non_note_events.append((absolute_time, msg.copy(time=0)))  # Use time=0 temporarily
    return notes, non_note_events

def assign_voices(notes):
    """
    Assign voices to notes based on overlapping.
    Lower pitch notes get lower voice numbers.
    Returns notes with assigned voice and total number of voices.
    """
    # Sort notes by start time, then by pitch (ascending)
    notes.sort(key=lambda x: (x.start, x.note))
    
    voices = []  # List of end times for each voice
    for note in notes:
        # Find available voices where the current voice's last note ends before the new note starts
        available = [i for i, end in enumerate(voices) if end <= note.start]
        if available:
            # Assign to the lowest available voice
            voice = min(available)
            voices[voice] = note.end
        else:
            # No available voice, create a new one
            voice = len(voices)
            voices.append(note.end)
        note.voice = voice
    return notes, len(voices)

def create_track_name(original_name, voice):
    if voice == 0:
        return original_name
    else:
        return f"{original_name}_{voice}"

def merge_events(note_events, non_note_events, voice):
    """
    Merge note events and non-note events for a specific voice.
    Returns a list of messages sorted by absolute time.
    """
    events = []

    # Add non-note events
    for abs_time, msg in non_note_events:
        events.append((abs_time, msg))

    # Add note_on and note_off events for this voice
    for note in note_events:
        if note.voice != voice:
            continue
        # Note on
        events.append((note.start, Message('note_on', note=note.note, velocity=note.velocity_on, channel=note.channel, time=0)))
        # Note off
        events.append((note.end, Message('note_off', note=note.note, velocity=note.velocity_off, channel=note.channel, time=0)))

    # Sort all events by absolute time
    events.sort(key=lambda x: x[0])

    # Convert absolute times to delta times
    merged_msgs = []
    prev_time = 0
    for abs_time, msg in events:
        delta = abs_time - prev_time
        msg.time = delta
        merged_msgs.append(msg)
        prev_time = abs_time

    return merged_msgs

def process_track(track, original_track_index):
    """
    Process a single MIDI track and split overlapping notes into new tracks.
    Preserves non-note messages by copying them to each suffixed track.
    Returns a list of new tracks.
    """
    # Extract track name
    track_name = f"Track{original_track_index}"
    for msg in track:
        if msg.type == 'track_name':
            track_name = msg.name
            break

    # Extract notes and non-note events
    notes, non_note_events = get_notes_and_events(track)
    if not notes:
        return [track]  # No notes to process

    # Assign voices
    assigned_notes, num_voices = assign_voices(notes)

    # Collect notes per voice
    voice_to_notes = defaultdict(list)
    for note in assigned_notes:
        voice_to_notes[note.voice].append(note)

    # Create new tracks
    new_tracks = []
    for voice in range(num_voices):
        new_track = MidiTrack()
        # Add track name
        new_track.append(MetaMessage('track_name', name=create_track_name(track_name, voice), time=0))
        # Merge and sort events
        merged_msgs = merge_events(voice_to_notes[voice], non_note_events, voice)
        new_track.extend(merged_msgs)
        new_tracks.append(new_track)
    return new_tracks

def generate_output_filename(input_file):
    """
    Generate output filename by inserting '_monofy' before the file extension.
    For example: 'song.mid' -> 'song_monofy.mid'
    """
    base, ext = os.path.splitext(input_file)
    return f"{base}_monofy{ext}"

def main(input_file):
    output_file = generate_output_filename(input_file)

    mid = MidiFile(input_file)
    new_mid = MidiFile()
    new_mid.ticks_per_beat = mid.ticks_per_beat

    for i, track in enumerate(mid.tracks):
        new_tracks = process_track(track, i)
        new_mid.tracks.extend(new_tracks)

    new_mid.save(output_file)
    print(f"Processed MIDI saved to '{output_file}'")

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python split_midi_tracks.py input.mid")
    else:
        input_file = sys.argv[1]
        if not os.path.isfile(input_file):
            print(f"Input file '{input_file}' does not exist.")
            sys.exit(1)
        main(input_file)