#include "Track.h"

#include <choc/audio/choc_MIDI.h>
#include <choc/audio/choc_MIDIFile.h>
#include <fstream>
#include <iostream>
#include <map>
#include <ranges>
#include <span>

using namespace std;
//--------------------------------------------------------------
/* Load a MIDI file from disk */
void Track::loadFromFile(const std::filesystem::path &filePath)
{
  cout << "Loading MIDI file: " << filePath.string() << endl;

  // Open the MIDI file in binary mode and place the file pointer at the end to determine its size
  std::ifstream file(filePath, std::ifstream::binary | std::ifstream::ate);
  if (!file.is_open())
    throw runtime_error("Failed to open MIDI file: " + filePath.string());

  const auto fileSize = file.tellg();
  file.seekg(0, std::ios::beg);

  // Read the entire file into a buffer
  std::vector<uint8_t> buffer(fileSize);
  if (!file.read(reinterpret_cast<char *>(buffer.data()), fileSize))
    throw runtime_error("Failed to read MIDI file: " + filePath.string());

  // Load the MIDI data from memory
  return loadFromMemory(buffer);
}
//--------------------------------------------------------------
/* Load a MIDI file from memory */
void Track::loadFromMemory(std::span<uint8_t> midiBytes)
{
  cout << "Loading MIDI data from memory, size: " << midiBytes.size() << " bytes" << endl;

  scoped_lock lock(m_mtx);     // Lock the mutex to protect access to the data structures

  // Clear any existing data before loading new MIDI data
  m_noteEvents.clear();
  m_tempoEvents.clear();
  m_notesByStart.clear();
  m_notesByEnd.clear();

  // Load midi sequence from memory using choc::midi::File
  choc::midi::File midiFile;
  midiFile.load(midiBytes.data(), midiBytes.size());

  // --- Step 1: Extract note events and tempo changes from the parsed MIDI data ---
  // Default tempo in microseconds per quarter note (500,000 us = 120 BPM)
  constexpr double defaultTempoPerQuarterNote = 500000.0;

  int trackID = 0;
  for (const auto &track : midiFile.tracks)
  {
    for (const auto &event : track.events)
    {
      const auto &msg = event.message;

      // Security check to ensure the message has enough bytes for processing
      if (msg.length() == 0) [[unlikely]]
        continue;

      // Check for Meta events (e.g., tempo changes, time signature changes, etc.)
      if (msg.isMetaEvent())
      {
        // Process tempo change events (Magic number 0x51 indicates a tempo change event)
        if (msg.length() >= 6 && msg.getMetaEventType() == 0x51)
        {
          const auto msPerQuarterNote = static_cast<uint32_t>(msg.data()[3] << 16) |
                                        static_cast<uint32_t>(msg.data()[4] << 8) |
                                        static_cast<uint32_t>(msg.data()[5]);
          MidiTempoEvents tempoEvent{ .tick                       = event.tickPosition,
                                      .microsecondsPerQuarterNote = static_cast<double>(msPerQuarterNote),
                                      .timeMs                     = Timestamp(0) };     // Time in milliseconds will be calculated later
          m_tempoEvents.push_back(tempoEvent);
        }
      }

      // Check for Note On events
      // Due to a bug on the choc::midi::Message class, we need to check the
      // length of the message before checking for Note On and Note Off events
      else if (msg.length() >= 3 && msg.isNoteOn())
      {
        MidiNoteEvent noteEvent{ .tick     = event.tickPosition,
                                 .pitch    = msg.getNoteNumber(),
                                 .velocity = msg.getVelocity(),
                                 .track    = trackID,
                                 .channel  = msg.getChannel0to15(),
                                 .noteOn   = true };
        m_noteEvents.push_back(noteEvent);
      }

      // Check for Note Off events
      // Due to a bug on the choc::midi::Message class, we need to check the
      // length of the message before checking for Note On and Note Off events
      else if (msg.length() >= 3 && msg.isNoteOff())
      {
        MidiNoteEvent noteEvent{ .tick     = event.tickPosition,
                                 .pitch    = msg.getNoteNumber(),
                                 .velocity = msg.getVelocity(),
                                 .track    = trackID,
                                 .channel  = msg.getChannel0to15(),
                                 .noteOn   = false };
        m_noteEvents.push_back(noteEvent);
      }
    }

    ++trackID;
  }

  // --- Step 2: Convert ticks to timestamps in milliseconds ---
  // Extract the time format from the MIDI file to determine ticks per beat
  double ticksPerBeat = 480.0;
  if (midiFile.timeFormat > 0)
  {
    ticksPerBeat = static_cast<double>(midiFile.timeFormat);
  }
  else if (midiFile.timeFormat < 0)
  {
    // Manage SMPTE time format (negative value) to calculate ticks per beat
    const int framesPerSecond = -static_cast<int>(midiFile.timeFormat >> 8);
    const int ticksPerFrame   = static_cast<int>(midiFile.timeFormat & 0xFF);
    ticksPerBeat              = static_cast<double>(framesPerSecond * ticksPerFrame);
  }
  // If ticksPerBeat is zero or negative, set it to a default value of 480.0
  if (ticksPerBeat <= 0.0) [[unlikely]]
  {
    ticksPerBeat = 480.0;
  }

  // Calculate the time in milliseconds for each tempo change event
  if (m_tempoEvents.empty())
  {
    // If no tempo events were found, use the default tempo for the entire track
    m_tempoEvents.push_back({ .tick                       = 0,
                              .microsecondsPerQuarterNote = defaultTempoPerQuarterNote,
                              .timeMs                     = Timestamp(0) });
  }
  else
  {
    // Sort the tempo events by tick to ensure they are in chronological order
    std::ranges::sort(m_tempoEvents, std::less<>{}, &MidiTempoEvents::tick);

    // Security check: if the first tempo event is not at tick 0, insert a default tempo event at tick 0
    if (m_tempoEvents.front().tick > 0)
    {
      constexpr MidiTempoEvents startTempoEvent{ .tick                       = 0,
                                                 .microsecondsPerQuarterNote = defaultTempoPerQuarterNote,
                                                 .timeMs                     = Timestamp(0) };     // Time in milliseconds will be calculated later
      m_tempoEvents.insert(m_tempoEvents.begin(), startTempoEvent);
    }

    auto    currentTimestamp = Timestamp(0);
    int64_t currentTick      = 0;
    auto    currentTempo     = m_tempoEvents[0].microsecondsPerQuarterNote;
    m_tempoEvents[0].timeMs  = Timestamp(0);     // The first tempo event starts at time 0
    for (size_t i = 1; i < m_tempoEvents.size(); i++)
    {
      const double tickDelta = static_cast<double>(m_tempoEvents[i].tick - currentTick);

      // Calculate the time delta in milliseconds based on the current tempo
      // Time in milliseconds = (tickDelta * microsecondsPerQuarterNote) / (ticksPerBeat * 1000)
      const auto timeDeltaMs = static_cast<int64_t>((tickDelta * currentTempo) / (ticksPerBeat * 1000.0));

      // Update the current timestamp and store it in the tempo event
      currentTimestamp += Timestamp(timeDeltaMs);
      currentTick  = m_tempoEvents[i].tick;
      currentTempo = m_tempoEvents[i].microsecondsPerQuarterNote;

      // Store the calculated time in milliseconds for the tempo event
      m_tempoEvents[i].timeMs = currentTimestamp;
    }
  }

  // --- Step 3: Convert note events from ticks to timestamps in milliseconds ---
  // Helper function to convert a tick value to a timestamp in milliseconds based on the tempo events
  const auto tickToMs = [this, ticksPerBeat](const int64_t tick) -> Timestamp
  {
    // Find the last tempo event that occurs before or at the given tick
    size_t tempoIndex = 0;
    for (size_t i = 0; i < m_tempoEvents.size(); i++)
    {
      if (m_tempoEvents[i].tick <= tick)
        tempoIndex = i;
      else
        break;
    }

    const auto  &tempoEvent = m_tempoEvents[tempoIndex];
    const double tickDelta  = static_cast<double>(tick - tempoEvent.tick);

    // Convert tick delta to milliseconds
    const auto timeDeltaMs = static_cast<int64_t>((tickDelta * tempoEvent.microsecondsPerQuarterNote) / (ticksPerBeat * 1000.0));

    return tempoEvent.timeMs + Timestamp(timeDeltaMs);
  };

  // --- Step 4: Create NoteEvent objects with start and end timestamps ---
  // Sort the note events by tick to ensure they are in chronological order
  std::ranges::sort(m_noteEvents, std::less<>{}, &MidiNoteEvent::tick);

  // Create a map to keep track of active notes (notes that have been started but not yet ended)
  // unordered_map<int, MidiNoteEvent> activeNotesMap;
  map<std::pair<int, int>, MidiNoteEvent> activeNotesMap;

  // Estimate the number of notes to reserve space in the vectors for performance
  // Half of the note events are expected to be Note On events, so we reserve half the size
  m_notesByStart.reserve(m_noteEvents.size() / 2);

  for (auto &noteEvent : m_noteEvents)
  {
    if (noteEvent.noteOn)
    {
      activeNotesMap[{ noteEvent.pitch, noteEvent.channel }] = noteEvent;     // Store the Note On event in the active notes map
    }
    else
    {
      // Note Off event: extract the corresponding Note On event from the active notes map
      auto node = activeNotesMap.extract({ noteEvent.pitch, noteEvent.channel });
      if (!node.empty())
      {
        const auto &noteOnEvent = node.mapped();

        // Create a NoteEvent with start and end timestamps
        NoteEvent note;
        note.startTimestamp = tickToMs(noteOnEvent.tick);
        note.endTimestamp   = tickToMs(noteEvent.tick);
        note.trackType      = TrackType::Unknown;     // Can be determined based on channel or other criteria
        note.pitch          = noteOnEvent.pitch;
        note.velocity       = noteOnEvent.velocity;
        note.noteOn         = true;

        // Manage channel-specific logic to determine track type (LeftHand or RightHand) based on the track number
        if (noteOnEvent.track == 0)
          note.trackType = TrackType::RightHand;
        else if (noteOnEvent.track == 1)
          note.trackType = TrackType::LeftHand;

        // Add the NoteEvent to the vector and remove it from the active notes map
        m_notesByStart.push_back(note);
      }
    }
  }

  // Handle any remaining active notes (notes without a corresponding note off)
  // These will be extended to a reasonable default duration (e.g., 100ms)
  for (const auto &noteOnEvent : activeNotesMap | views::values)
  {
    // Create a NoteEvent with a default end time (e.g., 100ms after the start time)
    NoteEvent note;
    note.startTimestamp = tickToMs(noteOnEvent.tick);
    note.endTimestamp   = tickToMs(noteOnEvent.tick) + 100ms;
    note.trackType      = TrackType::Unknown;
    note.pitch          = noteOnEvent.pitch;
    note.velocity       = noteOnEvent.velocity;
    note.noteOn         = true;

    // Add the NoteEvent to the vector
    m_notesByStart.push_back(note);
  }

  // --- Step 5: Sort notes for efficient retrieval ---
  // Sort by start time
  ranges::sort(m_notesByStart, std::less<>{}, &NoteEvent::startTimestamp);

  // Create a sorted-by-end-time vector
  m_notesByEnd = m_notesByStart;
  ranges::sort(m_notesByEnd, std::less<>{}, &NoteEvent::endTimestamp);

  cout << "Loaded " << m_notesByStart.size() << " notes and " << m_tempoEvents.size() << " tempo changes." << endl;
}
//--------------------------------------------------------------
/* Clear the track data */
void Track::clear()
{
  scoped_lock lock(m_mtx);     // Lock the mutex to protect access to the data structures

  m_noteEvents.clear();
  m_tempoEvents.clear();
  m_notesByStart.clear();
  m_notesByEnd.clear();
}
//--------------------------------------------------------------
/* Check if a MIDI file is loaded */
bool Track::isLoaded() const
{
  scoped_lock lock(m_mtx);     // Lock the mutex to protect access to the data structures

  return !m_noteEvents.empty();
}
//--------------------------------------------------------------
/* Get the duration of the track in milliseconds */
ITrack::Timestamp Track::getDuration() const
{
  scoped_lock lock(m_mtx);     // Lock the mutex to protect access to the data structures

  if (m_notesByEnd.empty())
    return Timestamp(0);

  const auto &lastNote = m_notesByEnd.back();
  return lastNote.endTimestamp;     // Return the end timestamp of the last note
}
//--------------------------------------------------------------
/* Debug function to print the track data */
void Track::debug()
{
  const auto timeToStr = [](const Timestamp &t) -> std::string
  {
    auto ms = t.count();
    auto s  = ms / 1000;
    auto m  = s / 60;
    auto h  = m / 60;
    ms      = ms % 1000;
    s       = s % 60;
    m       = m % 60;
    return std::format("{:02}:{:02}:{:02}.{:03}", h, m, s, ms);
  };

  for (const auto &noteEvent : m_notesByStart)
  {
    cout << format("Type: {}, pitch={}, velocity={}, start={}, end={}\n",
                   static_cast<int>(noteEvent.trackType),
                   noteEvent.pitch,
                   noteEvent.velocity,
                   timeToStr(noteEvent.startTimestamp),
                   timeToStr(noteEvent.endTimestamp));
  }
  cout << endl;
}
//--------------------------------------------------------------
/* Get notes in a given time window */
ITrack::TrackWindow Track::getTrackWindow(const Timestamp startTime, const Timestamp endTime) const
{
  std::unique_lock lock(m_mtx);

  // Check if notes are loaded
  if (m_noteEvents.empty())
    return { .activeNotes   = {},
             .upcomingNotes = {},
             .lock          = {} };

  // --- Preparation of the upcoming notes ---
  // Find the first note that starts after or at the startTime
  const auto upStart = std::ranges::lower_bound(m_notesByStart,
                                                startTime,
                                                std::less<>{},
                                                &NoteEvent::startTimestamp);

  // Find the first note that starts after or at the endTime
  const auto upEnd = std::ranges::lower_bound(upStart,
                                              m_notesByStart.end(),
                                              endTime,
                                              std::less<>{},
                                              &NoteEvent::startTimestamp);

  // Create a span for the upcoming notes
  const auto count = static_cast<size_t>(std::distance(upStart, upEnd));
  if (count == 0)
    return { .activeNotes   = {},
             .upcomingNotes = {},
             .lock          = {} };
  const std::span upcomingNotes(m_notesByStart.data() + std::distance(m_notesByStart.begin(), upStart), count);

  // --- Preparation of the active notes ---
  // Find the first note that ends after or at the startTime
  const auto activeStart = std::ranges::lower_bound(m_notesByEnd,
                                                    startTime,
                                                    std::less<>{},
                                                    &NoteEvent::endTimestamp);

  // Filter the active notes to include only those that have started before or at the endTime
  std::vector<NoteEvent> activeNotes;
  activeNotes.reserve(32);     // Reserve space for 32 notes, which is often sufficient for a frame

  for (auto it = activeStart; it != m_notesByEnd.end(); ++it)
  {
    // Check if the note is active at the startTime
    // (i.e., it has started before or at startTime and has not
    // ended yet)
    if (it->startTimestamp <= startTime && it->endTimestamp > startTime)
      activeNotes.push_back(*it);

    // Stop if the note starts after the endTime, as we only
    // want active notes in the window
    if (it->startTimestamp > endTime)
      break;
  }

  // --- Create and return the track window with the active and upcoming notes ---
  return { .activeNotes   = std::move(activeNotes),
           .upcomingNotes = upcomingNotes,
           .lock          = std::move(lock) };
}
//--------------------------------------------------------------