#include "midi/midi.h"
#include <CoreMIDI/CoreMIDI.h>
#include <CoreFoundation/CoreFoundation.h>
#include <string.h>
#include <stdatomic.h>

#define MIDI_MAX_DEVICES 64
#define MIDI_RELEASE_FLASH_DURATION 0.15f

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

    MIDIClientRef     client;
    MIDIPortRef       input_port;

    // Map: source endpoint index -> our device array index
    S32           source_to_device[MIDI_MAX_DEVICES];

    // Set atomically from callback thread
    _Atomic S32   pending_note_on_vel[MIDI_MAX_DEVICES];
    _Atomic S32   pending_note_num[MIDI_MAX_DEVICES];
    _Atomic S32   pending_note_off[MIDI_MAX_DEVICES];
    _Atomic S32   held_note_count[MIDI_MAX_DEVICES];

    // Per-note state (across all devices), set atomically from callback
    _Atomic S32   note_states[128];      // held count
    _Atomic S32   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 read callback

static void midi_read_callback(const MIDIPacketList *pktlist, void *readProcRefCon, void *srcConnRefCon) {
    (void)readProcRefCon;

    MidiEngine *engine = (MidiEngine *)readProcRefCon;
    S32 device_idx = (S32)(intptr_t)srcConnRefCon;
    if (!engine || device_idx < 0 || device_idx >= MIDI_MAX_DEVICES) return;

    const MIDIPacket *packet = &pktlist->packet[0];
    for (UInt32 i = 0; i < pktlist->numPackets; i++) {
        // Parse MIDI bytes
        for (UInt16 j = 0; j < packet->length; ) {
            U8 status = packet->data[j];

            // Skip non-status bytes (running status not handled for simplicity)
            if (status < 0x80) { j++; continue; }

            U8 kind = status & 0xF0;

            if (kind == 0x90 && j + 2 < packet->length) {
                U8 note = packet->data[j + 1];
                U8 velocity = packet->data[j + 2];
                j += 3;

                if (velocity > 0) {
                    atomic_store(&engine->pending_note_on_vel[device_idx], (S32)velocity);
                    atomic_store(&engine->pending_note_num[device_idx], (S32)(note + 1));
                    atomic_fetch_add(&engine->held_note_count[device_idx], 1);
                    if (note < 128) {
                        atomic_fetch_add(&engine->note_states[note], 1);
                        atomic_store(&engine->note_velocities[note], (S32)velocity);
                    }
                } else {
                    // Note-on with velocity 0 = note-off
                    atomic_store(&engine->pending_note_off[device_idx], 1);
                    S32 count = atomic_fetch_sub(&engine->held_note_count[device_idx], 1);
                    if (count <= 1) atomic_store(&engine->held_note_count[device_idx], 0);
                    if (note < 128) {
                        S32 c = atomic_fetch_sub(&engine->note_states[note], 1);
                        if (c <= 1) atomic_store(&engine->note_states[note], 0);
                    }
                }
            } else if (kind == 0x80 && j + 2 < packet->length) {
                U8 note_off = packet->data[j + 1];
                j += 3;
                atomic_store(&engine->pending_note_off[device_idx], 1);
                S32 count = atomic_fetch_sub(&engine->held_note_count[device_idx], 1);
                if (count <= 1) atomic_store(&engine->held_note_count[device_idx], 0);
                if (note_off < 128) {
                    S32 c = atomic_fetch_sub(&engine->note_states[note_off], 1);
                    if (c <= 1) atomic_store(&engine->note_states[note_off], 0);
                }
            } else if (kind == 0xC0 || kind == 0xD0) {
                j += 2; // Program Change, Channel Pressure (2 bytes)
            } else if (kind == 0xF0) {
                // System messages — skip to end or next status byte
                j++;
                while (j < packet->length && packet->data[j] < 0x80) j++;
            } else {
                j += 3; // Default: 3-byte message
            }
        }

        packet = MIDIPacketNext(packet);
    }
}

////////////////////////////////
// Device enumeration

static void enumerate_midi_devices(MidiEngine *engine) {
    engine->device_count = 0;

    // Input sources
    ItemCount num_sources = MIDIGetNumberOfSources();
    for (ItemCount i = 0; i < num_sources && engine->device_count < MIDI_MAX_DEVICES; i++) {
        MIDIEndpointRef endpoint = MIDIGetSource((ItemCount)i);

        CFStringRef name_ref = nullptr;
        MIDIObjectGetStringProperty(endpoint, kMIDIPropertyName, &name_ref);

        MidiDeviceInfo *dev = &engine->devices[engine->device_count];
        memset(dev, 0, sizeof(*dev));

        if (name_ref) {
            CFStringGetCString(name_ref, dev->name, sizeof(dev->name), kCFStringEncodingUTF8);
            CFRelease(name_ref);
        } else {
            snprintf(dev->name, sizeof(dev->name), "MIDI Source %d", (S32)i);
        }

        dev->id = engine->device_count;
        dev->is_input = true;
        dev->active = false;

        engine->source_to_device[i] = engine->device_count;
        engine->device_count++;
    }

    // Output destinations
    ItemCount num_dests = MIDIGetNumberOfDestinations();
    for (ItemCount i = 0; i < num_dests && engine->device_count < MIDI_MAX_DEVICES; i++) {
        MIDIEndpointRef endpoint = MIDIGetDestination((ItemCount)i);

        CFStringRef name_ref = nullptr;
        MIDIObjectGetStringProperty(endpoint, kMIDIPropertyName, &name_ref);

        MidiDeviceInfo *dev = &engine->devices[engine->device_count];
        memset(dev, 0, sizeof(*dev));

        if (name_ref) {
            CFStringGetCString(name_ref, dev->name, sizeof(dev->name), kCFStringEncodingUTF8);
            CFRelease(name_ref);
        } else {
            snprintf(dev->name, sizeof(dev->name), "MIDI Dest %d", (S32)i);
        }

        dev->id = engine->device_count;
        dev->is_input = false;
        dev->active = false;

        engine->device_count++;
    }
}

////////////////////////////////
// Public API

MidiEngine *midi_create() {
    MidiEngine *engine = new MidiEngine();
    memset(engine, 0, sizeof(*engine));

    MIDIClientCreate(CFSTR("autosample"), nullptr, nullptr, &engine->client);
    MIDIInputPortCreate(engine->client, CFSTR("Input"), midi_read_callback, engine, &engine->input_port);

    midi_refresh_devices(engine);
    return engine;
}

void midi_destroy(MidiEngine *engine) {
    midi_close_all_inputs(engine);
    if (engine->input_port) MIDIPortDispose(engine->input_port);
    if (engine->client) MIDIClientDispose(engine->client);
    delete engine;
}

void midi_open_all_inputs(MidiEngine *engine) {
    ItemCount num_sources = MIDIGetNumberOfSources();
    for (ItemCount i = 0; i < num_sources && (S32)i < MIDI_MAX_DEVICES; i++) {
        MIDIEndpointRef endpoint = MIDIGetSource(i);
        S32 dev_idx = engine->source_to_device[i];
        MIDIPortConnectSource(engine->input_port, endpoint, (void *)(intptr_t)dev_idx);
    }
}

void midi_close_all_inputs(MidiEngine *engine) {
    ItemCount num_sources = MIDIGetNumberOfSources();
    for (ItemCount i = 0; i < num_sources; i++) {
        MIDIEndpointRef endpoint = MIDIGetSource(i);
        MIDIPortDisconnectSource(engine->input_port, endpoint);
    }

    for (S32 i = 0; i < MIDI_MAX_DEVICES; i++) {
        atomic_store(&engine->pending_note_on_vel[i], 0);
        atomic_store(&engine->pending_note_num[i], 0);
        atomic_store(&engine->pending_note_off[i], 0);
        atomic_store(&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++) {
        atomic_store(&engine->note_states[i], 0);
        atomic_store(&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;

        S32 vel = atomic_exchange(&engine->pending_note_on_vel[i], 0);
        S32 note_p1 = atomic_exchange(&engine->pending_note_num[i], 0);
        if (vel > 0) engine->display_velocity[i] = vel;
        if (note_p1 > 0) engine->display_note[i] = note_p1 - 1;

        S32 off = atomic_exchange(&engine->pending_note_off[i], 0);
        S32 held = atomic_load(&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)) {
            engine->devices[i].active = false;
            engine->devices[i].releasing = true;
            engine->release_timers[i] = MIDI_RELEASE_FLASH_DURATION;
        }

        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 atomic_load(&engine->note_states[note]) > 0;
}

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

void midi_refresh_devices(MidiEngine *engine) {
    midi_close_all_inputs(engine);
    enumerate_midi_devices(engine);
    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];
}
