import mido
from mido import MidiFile, MidiTrack, MetaMessage


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


def compute_absolute_times(mid):
    ticks_per_beat = mid.ticks_per_beat
    DEFAULT_TEMPO = 500000

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

    all_msgs.sort(key=lambda m: m.abs_tick)

    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 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):
    for msg in all_msgs:
        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):
    new_tracks = [[] for _ in mid.tracks]

    for msg in all_msgs:
        if msg.message.type == 'set_tempo':
            continue
        new_tracks[msg.track].append(msg)

    for track_index, track_msgs in enumerate(new_tracks):
        track_msgs.sort(key=lambda m: m.new_tick)

        prev_tick = 0
        new_track = []
        for msg in track_msgs:
            delta_tick = msg.new_tick - prev_tick
            prev_tick = msg.new_tick
            new_msg = msg.message.copy(time=delta_tick)
            new_track.append(new_msg)

        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):
    for msg in all_msgs:
        if msg.message.type == 'set_tempo':
            return msg.message.tempo
    return default_tempo


def process(mid: MidiFile) -> MidiFile:
    ticks_per_beat = mid.ticks_per_beat

    all_msgs = compute_absolute_times(mid)
    initial_tempo = get_initial_tempo(all_msgs)
    all_msgs = bake_tempo(all_msgs, ticks_per_beat, initial_tempo)
    new_tracks = assign_ticks_to_tracks(all_msgs, mid, ticks_per_beat)

    new_mid = MidiFile(ticks_per_beat=ticks_per_beat)

    tempo_track = MidiTrack()
    tempo_track.append(MetaMessage('set_tempo', tempo=initial_tempo, time=0))
    new_mid.tracks.append(tempo_track)

    for track in new_tracks:
        new_mid.tracks.append(track)

    return new_mid