#include "MifiFile.h"

#include "FileHelper.h"
#include "MidiMessageSequence.h"
#include "ByteOrder.h"
#include "Math.h"

namespace MidiFileHelpers
{
    static void writeVariableLengthInt (TArray<uint8>& out, uint32 v)
    {
        auto buffer = v & 0x7f;

        while ((v >>= 7) != 0)
        {
            buffer <<= 8;
            buffer |= ((v & 0x7f) | 0x80);
        }

        for (;;)
        {
            out.Add((char)buffer);

            if (buffer & 0x80)
                buffer >>= 8;
            else
                break;
        }
    }

    template <typename Integral>
    struct ReadTrait;

    template <>
    struct ReadTrait<uint32> { static constexpr auto read = FByteOrder::bigEndianInt; };

    template <>
    struct ReadTrait<uint16> { static constexpr auto read = FByteOrder::bigEndianShort; };

    template <typename Integral>
    TOptional<Integral> tryRead (const uint8*& data, size_t& remaining)
    {
        using Trait = ReadTrait<Integral>;
        constexpr auto size = sizeof (Integral);

        if (remaining < size)
            return {};

        const TOptional<Integral> result { Trait::read (data) };

        data += size;
        remaining -= size;

        return result;
    }

    struct HeaderDetails
    {
        size_t bytesRead = 0;
        short timeFormat = 0;
        short fileType = 0;
        short numberOfTracks = 0;
    };

    static TOptional<HeaderDetails> parseMidiHeader (const uint8* const initialData,
                                                    const size_t maxSize)
    {
        auto* data = initialData;
        auto remaining = maxSize;

        auto ch = tryRead<uint32> (data, remaining);

        if (! ch.IsSet())
            return {};

        if (*ch != FByteOrder::bigEndianInt ("MThd"))
        {
            auto ok = false;

            if (*ch == FByteOrder::bigEndianInt ("RIFF"))
            {
                for (int i = 0; i < 8; ++i)
                {
                    ch = tryRead<uint32> (data, remaining);

                    if (! ch.IsSet())
                        return {};

                    if (*ch == FByteOrder::bigEndianInt ("MThd"))
                    {
                        ok = true;
                        break;
                    }
                }
            }

            if (! ok)
                return {};
        }

        const auto bytesRemaining = tryRead<uint32> (data, remaining);

        if (! bytesRemaining.IsSet() || *bytesRemaining > remaining)
            return {};

        const auto optFileType = tryRead<uint16> (data, remaining);

        if (! optFileType.IsSet() || 2 < *optFileType)
            return {};

        const auto optNumTracks = tryRead<uint16> (data, remaining);

        if (! optNumTracks.IsSet() || (*optFileType == 0 && *optNumTracks != 1))
            return {};

        const auto optTimeFormat = tryRead<uint16> (data, remaining);

        if (! optTimeFormat.IsSet())
            return {};

        HeaderDetails result;

        result.fileType = (short) *optFileType;
        result.timeFormat = (short) *optTimeFormat;
        result.numberOfTracks = (short) *optNumTracks;
        result.bytesRead = maxSize - remaining;

        return { result };
    }

    static double convertTicksToSeconds (double time,
                                         const FMidiMessageSequence& tempoEvents,
                                         int timeFormat)
    {
        if (timeFormat < 0)
            return time / (-(timeFormat >> 8) * (timeFormat & 0xff));

        double lastTime = 0, correctedTime = 0;
        auto tickLen = 1.0 / (timeFormat & 0x7fff);
        auto secsPerTick = 0.5 * tickLen;
        auto numEvents = tempoEvents.getNumEvents();

        for (int i = 0; i < numEvents; ++i)
        {
            auto& m = tempoEvents.getEventPointer(i)->message;
            auto eventTime = m.getTimeStamp();

            if (eventTime >= time)
                break;

            correctedTime += (eventTime - lastTime) * secsPerTick;
            lastTime = eventTime;

            if (m.isTempoMetaEvent())
                secsPerTick = tickLen * m.getTempoSecondsPerQuarterNote();

            while (i + 1 < numEvents)
            {
                auto& m2 = tempoEvents.getEventPointer(i + 1)->message;

                if (m2.getTimeStamp() != eventTime)
                    break;

                if (m2.isTempoMetaEvent())
                    secsPerTick = tickLen * m2.getTempoSecondsPerQuarterNote();

                ++i;
            }
        }

        return correctedTime + (time - lastTime) * secsPerTick;
    }

    template <typename MethodType>
    static void findAllMatchingEvents (const TArray<FMidiMessageSequence*>& tracks,
                                       FMidiMessageSequence& results,
                                       MethodType method)
    {
        for (auto* track : tracks)
        {
            auto numEvents = track->getNumEvents();

            for (int j = 0; j < numEvents; ++j)
            {
                auto& m = track->getEventPointer(j)->message;

                if ((m.*method)())
                    results.addEvent (m);
            }
        }
    }

    static FMidiMessageSequence readTrack (const uint8* data, int size)
    {
        double time = 0;
        uint8 lastStatusByte = 0;

        FMidiMessageSequence result;

        while (size > 0)
        {
            const auto delay = FMidiMessage::readVariableLengthValue (data, (int) size);

            if (! delay.isValid())
                break;

            data += delay.bytesUsed;
            size -= delay.bytesUsed;
            time += delay.value;

            if (size <= 0)
                break;

            int messSize = 0;
            const FMidiMessage mm (data, size, messSize, lastStatusByte, time);

            if (messSize <= 0)
                break;

            size -= messSize;
            data += messSize;

            result.addEvent (mm);

            auto firstByte = *(mm.getRawData());

            if ((firstByte & 0xf0) != 0xf0)
                lastStatusByte = firstByte;
        }

        return result;
    }
}

//==============================================================================
FMidiFile::FMidiFile() : timeFormat ((short) (unsigned short) 0xe728) {}

FMidiFile::FMidiFile (const FMidiFile& other) : timeFormat (other.timeFormat)
{
    tracks.Append(other.tracks);
}

FMidiFile& FMidiFile::operator= (const FMidiFile& other)
{
    tracks.Empty();
    tracks.Append(other.tracks);
    timeFormat = other.timeFormat;
    return *this;
}

FMidiFile::FMidiFile (FMidiFile&& other)
    : tracks (std::move (other.tracks)),
      timeFormat (other.timeFormat)
{
}

FMidiFile& FMidiFile::operator= (FMidiFile&& other)
{
    tracks = std::move (other.tracks);
    timeFormat = other.timeFormat;
    return *this;
}

void FMidiFile::clear()
{
    tracks.Empty();
}

//==============================================================================
int FMidiFile::getNumTracks() const noexcept
{
    return tracks.Num();
}

const FMidiMessageSequence* FMidiFile::getTrack (int index) const noexcept
{
    return tracks[index];
}

void FMidiFile::addTrack (const FMidiMessageSequence& trackSequence)
{
    tracks.Add(new FMidiMessageSequence (trackSequence));
}

//==============================================================================
short FMidiFile::getTimeFormat() const noexcept
{
    return timeFormat;
}

void FMidiFile::setTicksPerQuarterNote (int ticks) noexcept
{
    timeFormat = (short) ticks;
}

void FMidiFile::setSmpteTimeFormat (int framesPerSecond, int subframeResolution) noexcept
{
    timeFormat = (short) (((-framesPerSecond) << 8) | subframeResolution);
}

//==============================================================================
void FMidiFile::findAllTempoEvents (FMidiMessageSequence& results) const
{
    MidiFileHelpers::findAllMatchingEvents (tracks, results, &FMidiMessage::isTempoMetaEvent);
}

void FMidiFile::findAllTimeSigEvents (FMidiMessageSequence& results) const
{
    MidiFileHelpers::findAllMatchingEvents (tracks, results, &FMidiMessage::isTimeSignatureMetaEvent);
}

void FMidiFile::findAllKeySigEvents (FMidiMessageSequence& results) const
{
    MidiFileHelpers::findAllMatchingEvents (tracks, results, &FMidiMessage::isKeySignatureMetaEvent);
}

double FMidiFile::getLastTimestamp() const
{
    double t = 0.0;

    for (auto* ms : tracks)
        t = FMath::Max(t, ms->getEndTime());

    return t;
}

//==============================================================================
bool FMidiFile::readFrom(FString FilePathName,
                         bool createMatchingNoteOffs,
                         int* fileType)
{
    clear();
    TArray<uint8> data;

    if (!FFileHelper::LoadFileToArray(data, *FilePathName))
        return false;

    size_t size = data.Num();
    auto d = static_cast<const uint8*> (data.GetData());

    const auto optHeader = MidiFileHelpers::parseMidiHeader (d, size);

    if (! optHeader.IsSet())
        return false;

    const auto header = *optHeader;
    timeFormat = header.timeFormat;

    d += header.bytesRead;
    size -= (size_t) header.bytesRead;

    for (int track = 0; track < header.numberOfTracks; ++track)
    {
        const auto optChunkType = MidiFileHelpers::tryRead<uint32> (d, size);

        if (!optChunkType.IsSet())
            return false;

        const auto optChunkSize = MidiFileHelpers::tryRead<uint32> (d, size);

        if (! optChunkSize.IsSet())
            return false;

        const auto chunkSize = *optChunkSize;

        if (size < chunkSize)
            return false;

        if (*optChunkType == FByteOrder::bigEndianInt ("MTrk"))
            readNextTrack (d, (int) chunkSize, createMatchingNoteOffs);

        size -= chunkSize;
        d += chunkSize;
    }

    const auto successful = (size == 0);

    if (successful && fileType != nullptr)
        *fileType = header.fileType;

    return successful;
}

void FMidiFile::readNextTrack (const uint8* data, int size, bool createMatchingNoteOffs)
{
    auto sequence = MidiFileHelpers::readTrack (data, size);

    // sort so that we put all the note-offs before note-ons that have the same time
    Algo::StableSort(sequence.list, [] (const FMidiMessageSequence::MidiEventHolder* a,
                          const FMidiMessageSequence::MidiEventHolder* b)
    {
        auto t1 = a->message.getTimeStamp();
        auto t2 = b->message.getTimeStamp();

        if (t1 < t2)  return true;
        if (t2 < t1)  return false;

        return a->message.isNoteOff() && b->message.isNoteOn();
    });

    if (createMatchingNoteOffs)
        sequence.updateMatchedPairs();

    addTrack (sequence);
}

//==============================================================================
void FMidiFile::convertTimestampTicksToSeconds()
{
    FMidiMessageSequence tempoEvents;
    findAllTempoEvents (tempoEvents);
    findAllTimeSigEvents (tempoEvents);

    if (timeFormat != 0)
    {
        for (auto* ms : tracks)
        {
            for (int j = ms->getNumEvents(); --j >= 0;)
            {
                auto& m = ms->getEventPointer(j)->message;
                m.setTimeStamp (MidiFileHelpers::convertTicksToSeconds (m.getTimeStamp(), tempoEvents, timeFormat));
            }
        }
    }
}

//==============================================================================
bool FMidiFile::writeTo (FArchive& out, int midiFileType) const
{
    check(midiFileType >= 0 && midiFileType <= 2);

    int MThd = FByteOrder::bigEndianInt("MThd");
    int m = 6;
    short format = (short)midiFileType;
    short numTracks = (short)tracks.Num();
    short tempTimeFormat = timeFormat;
    
    out << MThd;
    out << m;
    out << format;
    out << numTracks;
    out << tempTimeFormat;

    for (auto* ms : tracks)
        if (!writeTrack(out, *ms))
            return false;

    out.Flush();
    return true;
}

bool FMidiFile::writeTrack (FArchive& mainOut, const FMidiMessageSequence& ms) const
{
    TArray<uint8> out;

    int lastTick = 0;
    uint8 lastStatusByte = 0;
    bool endOfTrackEventWritten = false;

    for (int i = 0; i < ms.getNumEvents(); ++i)
    {
        auto& mm = ms.getEventPointer(i)->message;

        if (mm.isEndOfTrackMetaEvent())
            endOfTrackEventWritten = true;

        auto tick = RoundToInt(mm.getTimeStamp());
        auto delta = FMath::Max(0, tick - lastTick);
        MidiFileHelpers::writeVariableLengthInt(out, (uint32) delta);
        lastTick = tick;

        auto* data = mm.getRawData();
        auto dataSize = mm.getRawDataSize();
        auto statusByte = data[0];

        if (statusByte == lastStatusByte
             && (statusByte & 0xf0) != 0xf0
             && dataSize > 1
             && i > 0)
        {
            ++data;
            --dataSize;
        }
        else if (statusByte == 0xf0)  // Write sysex message with length bytes.
        {
            out.Add((char)statusByte);

            ++data;
            --dataSize;

            MidiFileHelpers::writeVariableLengthInt (out, (uint32) dataSize);
        }

        out.Append(data, (size_t)dataSize);
        lastStatusByte = statusByte;
    }

    if (! endOfTrackEventWritten)
    {
        out.Add(0); // (tick delta)
        auto m = FMidiMessage::endOfTrack();
        out.Append(m.getRawData(), (size_t)m.getRawDataSize());
    }
    int MTrk = FByteOrder::bigEndianInt("MTrk");
    
    mainOut << MTrk;
    mainOut.Serialize(out.GetData(), out.Num());

    return true;
}