#include "midi/midi.h"
#include <windows.h>
#include <mmeapi.h>
#include <string.h>

#define MIDI_MAX_DEVICES 64
#define MIDI_RELEASE_FLASH_DURATION 0.15f

struct MidiEngine {
    MidiDeviceInfo devices[MIDI_MAX_DEVICES];
    S32        device_count;

    HMIDIIN        input_handles[MIDI_MAX_DEVICES];

    // Set atomically from callback thread
    volatile LONG  pending_note_on_vel[MIDI_MAX_DEVICES];  // last note-on velocity (0 = consumed)
    volatile LONG  pending_note_num[MIDI_MAX_DEVICES];     // last note-on note number + 1 (0 = consumed)
    volatile LONG  pending_note_off[MIDI_MAX_DEVICES];     // note-off received flag
    volatile LONG  held_note_count[MIDI_MAX_DEVICES];      // number of notes currently held

    // Per-note state (across all devices), set atomically from callback
    volatile LONG  note_states[128];      // held count
    volatile LONG  note_velocities[128];  // last note-on velocity

    // Main thread only
    S32        display_velocity[MIDI_MAX_DEVICES];
    S32        display_note[MIDI_MAX_DEVICES];
    F32          release_timers[MIDI_MAX_DEVICES];
};

////////////////////////////////
// MIDI input callback — called from Win32 MIDI driver thread

static MidiEngine *g_midi_engine = nullptr;

static void CALLBACK midi_in_callback(HMIDIIN hMidiIn, UINT wMsg, DWORD_PTR dwInstance,
                                       DWORD_PTR dwParam1, DWORD_PTR dwParam2) {
    (void)hMidiIn;
    (void)dwParam2;

    if (wMsg != MIM_DATA) return;
    if (!g_midi_engine) return;

    S32 idx = (S32)dwInstance;
    if (idx < 0 || idx >= MIDI_MAX_DEVICES) return;

    BYTE status   = (BYTE)(dwParam1 & 0xFF);
    BYTE note     = (BYTE)((dwParam1 >> 8) & 0xFF);
    BYTE velocity = (BYTE)((dwParam1 >> 16) & 0xFF);
    BYTE kind     = status & 0xF0;

    // Note-on with velocity > 0
    if (kind == 0x90 && velocity > 0) {
        InterlockedExchange(&g_midi_engine->pending_note_on_vel[idx], (LONG)velocity);
        InterlockedExchange(&g_midi_engine->pending_note_num[idx], (LONG)(note + 1)); // +1 so 0 means "no pending"
        InterlockedIncrement(&g_midi_engine->held_note_count[idx]);
        if (note < 128) {
            InterlockedIncrement(&g_midi_engine->note_states[note]);
            InterlockedExchange(&g_midi_engine->note_velocities[note], (LONG)velocity);
        }
    }
    // Note-off (0x80) or note-on with velocity 0 (running status note-off)
    else if (kind == 0x80 || (kind == 0x90 && velocity == 0)) {
        InterlockedExchange(&g_midi_engine->pending_note_off[idx], 1);
        LONG count = InterlockedDecrement(&g_midi_engine->held_note_count[idx]);
        if (count < 0) InterlockedExchange(&g_midi_engine->held_note_count[idx], 0);
        if (note < 128) {
            LONG c = InterlockedDecrement(&g_midi_engine->note_states[note]);
            if (c < 0) InterlockedExchange(&g_midi_engine->note_states[note], 0);
        }
    }
}

MidiEngine *midi_create() {
    MidiEngine *engine = new MidiEngine();
    memset(engine, 0, sizeof(*engine));
    g_midi_engine = engine;
    midi_refresh_devices(engine);
    return engine;
}

void midi_destroy(MidiEngine *engine) {
    midi_close_all_inputs(engine);
    if (g_midi_engine == engine) g_midi_engine = nullptr;
    delete engine;
}

void midi_open_all_inputs(MidiEngine *engine) {
    for (S32 i = 0; i < engine->device_count; i++) {
        MidiDeviceInfo *dev = &engine->devices[i];
        if (!dev->is_input) continue;
        if (engine->input_handles[i]) continue; // already open

        HMIDIIN handle = nullptr;
        MMRESULT res = midiInOpen(&handle, (UINT)dev->id,
                                   (DWORD_PTR)midi_in_callback,
                                   (DWORD_PTR)i, CALLBACK_FUNCTION);
        if (res == MMSYSERR_NOERROR) {
            engine->input_handles[i] = handle;
            midiInStart(handle);
        }
    }
}

void midi_close_all_inputs(MidiEngine *engine) {
    for (S32 i = 0; i < MIDI_MAX_DEVICES; i++) {
        if (engine->input_handles[i]) {
            midiInStop(engine->input_handles[i]);
            midiInClose(engine->input_handles[i]);
            engine->input_handles[i] = nullptr;
        }
        engine->pending_note_on_vel[i] = 0;
        engine->pending_note_num[i] = 0;
        engine->pending_note_off[i] = 0;
        engine->held_note_count[i] = 0;
        engine->display_velocity[i] = 0;
        engine->display_note[i] = 0;
        engine->release_timers[i] = 0.0f;
    }
    for (S32 i = 0; i < 128; i++) {
        engine->note_states[i] = 0;
        engine->note_velocities[i] = 0;
    }
}

void midi_update(MidiEngine *engine, F32 dt) {
    for (S32 i = 0; i < engine->device_count; i++) {
        if (!engine->devices[i].is_input) continue;

        // Consume pending note-on velocity and note number
        LONG vel = InterlockedExchange(&engine->pending_note_on_vel[i], 0);
        LONG note_p1 = InterlockedExchange(&engine->pending_note_num[i], 0);
        if (vel > 0) {
            engine->display_velocity[i] = (S32)vel;
        }
        if (note_p1 > 0) {
            engine->display_note[i] = (S32)(note_p1 - 1);
        }

        // Consume pending note-off
        LONG off = InterlockedExchange(&engine->pending_note_off[i], 0);

        // Read held note count
        LONG held = engine->held_note_count[i];

        if (held > 0) {
            engine->devices[i].active = true;
            engine->devices[i].releasing = false;
            engine->release_timers[i] = 0.0f;
        } else if (off || (engine->devices[i].active && held <= 0)) {
            // All notes just released — start release flash
            engine->devices[i].active = false;
            engine->devices[i].releasing = true;
            engine->release_timers[i] = MIDI_RELEASE_FLASH_DURATION;
        }

        // Decay release flash timer
        if (engine->release_timers[i] > 0.0f) {
            engine->release_timers[i] -= dt;
            if (engine->release_timers[i] <= 0.0f) {
                engine->release_timers[i] = 0.0f;
                engine->devices[i].releasing = false;
                engine->display_velocity[i] = 0;
                engine->display_note[i] = 0;
            }
        }

        engine->devices[i].velocity = engine->display_velocity[i];
        engine->devices[i].note = engine->display_note[i];
    }
}

B32 midi_is_input_active(MidiEngine *engine, S32 device_index) {
    if (device_index < 0 || device_index >= engine->device_count)
        return false;
    return engine->devices[device_index].active;
}

B32 midi_is_note_held(MidiEngine *engine, S32 note) {
    if (note < 0 || note > 127) return false;
    return engine->note_states[note] > 0;
}

S32 midi_get_note_velocity(MidiEngine *engine, S32 note) {
    if (note < 0 || note > 127) return 0;
    if (engine->note_states[note] <= 0) return 0;
    return (S32)engine->note_velocities[note];
}

void midi_refresh_devices(MidiEngine *engine) {
    midi_close_all_inputs(engine);
    engine->device_count = 0;

    UINT num_in = midiInGetNumDevs();
    for (UINT i = 0; i < num_in && engine->device_count < MIDI_MAX_DEVICES; i++) {
        MIDIINCAPSA caps = {};
        if (midiInGetDevCapsA(i, &caps, sizeof(caps)) == MMSYSERR_NOERROR) {
            MidiDeviceInfo *dev = &engine->devices[engine->device_count++];
            strncpy_s(dev->name, sizeof(dev->name), caps.szPname, _TRUNCATE);
            dev->id = (S32)i;
            dev->is_input = true;
            dev->active = false;
        }
    }

    UINT num_out = midiOutGetNumDevs();
    for (UINT i = 0; i < num_out && engine->device_count < MIDI_MAX_DEVICES; i++) {
        MIDIOUTCAPSA caps = {};
        if (midiOutGetDevCapsA(i, &caps, sizeof(caps)) == MMSYSERR_NOERROR) {
            MidiDeviceInfo *dev = &engine->devices[engine->device_count++];
            strncpy_s(dev->name, sizeof(dev->name), caps.szPname, _TRUNCATE);
            dev->id = (S32)i;
            dev->is_input = false;
            dev->active = false;
        }
    }

    midi_open_all_inputs(engine);
}

S32 midi_get_device_count(MidiEngine *engine) {
    return engine->device_count;
}

MidiDeviceInfo *midi_get_device(MidiEngine *engine, S32 index) {
    if (index < 0 || index >= engine->device_count)
        return nullptr;
    return &engine->devices[index];
}
