import mido
from mido import MidiFile, MidiTrack, Message, MetaMessage
import sys
import os
import argparse

def parse_args():
    parser = argparse.ArgumentParser(description='Bake tempo changes into MIDI file with constant tempo.')
    parser.add_argument('input_file', type=str, help='Path to the input MIDI file.')
    return parser.parse_args()

class MIDIMessage:
    def __init__(self, message, track, abs_tick, abs_time):
        self.message = message
        self.track = track
        self.abs_tick = abs_tick
        self.abs_time = abs_time
        self.new_tick = None  # To be calculated later

def compute_absolute_times(mid):
    """
    Compute absolute time for each message across all tracks, considering tempo changes.
    Returns a list of MIDIMessage instances with absolute times.
    """
    ticks_per_beat = mid.ticks_per_beat
    DEFAULT_TEMPO = 500000  # Default tempo (microseconds per beat)

    # Collect all messages with their absolute ticks and track index
    all_msgs = []
    for track_index, track in enumerate(mid.tracks):
        abs_tick = 0
        for msg in track:
            abs_tick += msg.time
            all_msgs.append(MIDIMessage(msg, track_index, abs_tick, 0.0))

    # Sort all messages by absolute tick
    all_msgs.sort(key=lambda m: m.abs_tick)

    # Now, compute absolute times
    current_tempo = DEFAULT_TEMPO
    abs_time = 0.0
    prev_tick = 0

    for msg in all_msgs:
        delta_ticks = msg.abs_tick - prev_tick
        delta_time = mido.tick2second(delta_ticks, ticks_per_beat, current_tempo)
        abs_time += delta_time
        msg.abs_time = abs_time

        # If the message is a tempo change, update the current tempo
        if msg.message.type == 'set_tempo':
            current_tempo = msg.message.tempo
        prev_tick = msg.abs_tick

    return all_msgs

def bake_tempo(all_msgs, ticks_per_beat, constant_tempo):
    """
    Assign new ticks to each message based on absolute time and constant tempo.
    """
    for msg in all_msgs:
        # Calculate new tick based on absolute time and constant tempo
        # new_tick = abs_time / seconds_per_tick
        # seconds_per_tick = constant_tempo / ticks_per_beat / 1e6
        seconds_per_tick = constant_tempo / ticks_per_beat / 1e6
        msg.new_tick = int(round(msg.abs_time / seconds_per_tick))

    return all_msgs

def assign_ticks_to_tracks(all_msgs, mid, ticks_per_beat):
    """
    Assign new delta ticks to each track based on the baked tempo.
    Returns a list of tracks with updated messages.
    """
    # Prepare a list for each track
    new_tracks = [[] for _ in mid.tracks]

    # Sort messages back into their respective tracks
    for msg in all_msgs:
        if msg.message.type == 'set_tempo':
            # Skip tempo messages
            continue
        new_tracks[msg.track].append(msg)

    # Now, for each track, sort messages by new_tick and assign delta ticks
    for track_index, track_msgs in enumerate(new_tracks):
        # Sort messages by new_tick
        track_msgs.sort(key=lambda m: m.new_tick)

        # Assign delta ticks
        prev_tick = 0
        new_track = []
        for msg in track_msgs:
            delta_tick = msg.new_tick - prev_tick
            prev_tick = msg.new_tick

            # Create a copy of the message to avoid modifying the original
            new_msg = msg.message.copy(time=delta_tick)
            new_track.append(new_msg)

        # Ensure the track ends with an end_of_track message
        if not new_track or new_track[-1].type != 'end_of_track':
            new_track.append(MetaMessage('end_of_track', time=0))

        new_tracks[track_index] = new_track

    return new_tracks

def get_initial_tempo(all_msgs, default_tempo=500000):
    """
    Retrieve the initial tempo from the list of MIDI messages.
    If no set_tempo message is found, return the default tempo.
    """
    for msg in all_msgs:
        if msg.message.type == 'set_tempo':
            return msg.message.tempo
    return default_tempo

def main():
    args = parse_args()
    input_path = args.input_file

    if not os.path.isfile(input_path):
        print(f"Error: File '{input_path}' does not exist.")
        sys.exit(1)

    try:
        mid = MidiFile(input_path)
    except Exception as e:
        print(f"Error reading MIDI file: {e}")
        sys.exit(1)

    ticks_per_beat = mid.ticks_per_beat

    # Compute absolute times for all messages
    all_msgs = compute_absolute_times(mid)

    # Get the initial tempo (first set_tempo message or default)
    initial_tempo = get_initial_tempo(all_msgs)

    # Bake the tempo by assigning new ticks based on constant tempo
    all_msgs = bake_tempo(all_msgs, ticks_per_beat, initial_tempo)

    # Assign new delta ticks to each track
    new_tracks = assign_ticks_to_tracks(all_msgs, mid, ticks_per_beat)

    # Create a new MIDI file
    new_mid = MidiFile(ticks_per_beat=ticks_per_beat)

    # Create a tempo track with the initial tempo
    tempo_track = MidiTrack()
    tempo_track.append(MetaMessage('set_tempo', tempo=initial_tempo, time=0))
    new_mid.tracks.append(tempo_track)

    # Add the updated performance tracks
    for track in new_tracks:
        new_mid.tracks.append(track)

    # Define output file name
    base, ext = os.path.splitext(input_path)
    output_path = f"{base}_tempo{ext}"

    try:
        new_mid.save(output_path)
        print(f"Successfully saved baked MIDI file as '{output_path}'.")
    except Exception as e:
        print(f"Error saving MIDI file: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()