#include "engine.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>

// JACK process callback
static int process_callback(jack_nframes_t nframes, void *arg) {
    Engine *engine = (Engine *)arg;
    
    // Get ports
    jack_default_audio_sample_t *audio_in = (jack_default_audio_sample_t *)
        jack_port_get_buffer(engine->audio_in_port, nframes);
    jack_default_audio_sample_t *audio_out = (jack_default_audio_sample_t *)
        jack_port_get_buffer(engine->audio_out_port, nframes);
    void *midi_in_buf = jack_port_get_buffer(engine->midi_in_port, nframes);
    void *midi_out_buf = jack_port_get_buffer(engine->midi_out_port, nframes);
    
    // Clear output MIDI buffer
    jack_midi_clear_buffer(midi_out_buf);
    
    // Process MIDI input
    jack_midi_event_t midi_event;
    jack_nframes_t event_index = 0;
    
    while (jack_midi_event_get(&midi_event, midi_in_buf, event_index) == 0) {
        event_index++;
        
        uint8_t *data = midi_event.buffer;
        uint8_t status = data[0] & 0xF0;
        uint8_t channel = data[0] & 0x0F;
        uint8_t note = data[1];
        uint8_t velocity = data[2];
        
        // Only process note on messages on the control channel
        if (status == 0x90 && channel == engine->control_channel && velocity > 0) {
            int clip_index = note % MAX_CLIPS;
            engine_trigger_clip(engine, clip_index);
            
            // Send note with velocity representing state
            uint8_t out_velocity = clip_state_to_velocity(engine->clips[clip_index].state);
            uint8_t out_msg[3] = {0x90 | channel, note, out_velocity};
            
            if (jack_midi_event_write(midi_out_buf, midi_event.time, out_msg, 3) != 0) {
                fprintf(stderr, "Failed to write MIDI event\n");
            }
        } else {
            // Pass through all other MIDI messages
            if (jack_midi_event_write(midi_out_buf, midi_event.time, 
                                      midi_event.buffer, midi_event.size) != 0) {
                fprintf(stderr, "Failed to write MIDI event\n");
            }
        }
    }
    
    // Process audio
    memset(audio_out, 0, sizeof(jack_default_audio_sample_t) * nframes);
    
    for (jack_nframes_t i = 0; i < nframes; i++) {
        // Record input to recording clips
        for (int c = 0; c < MAX_CLIPS; c++) {
            Clip *clip = &engine->clips[c];
            
            if (clip->state == CLIP_RECORDING) {
                if (clip->write_position < MAX_BUFFER_SIZE) {
                    clip->buffer[clip->write_position++] = audio_in[i];
                } else {
                    // Buffer full, stop recording
                    clip->state = CLIP_LOOPING;
                    clip->buffer_size = clip->write_position;
                    clip->read_position = 0;
                }
            }
            
            // Play looping clips
            if (clip->state == CLIP_LOOPING && clip->buffer_size > 0) {
                audio_out[i] += clip->buffer[clip->read_position];
                clip->read_position = (clip->read_position + 1) % clip->buffer_size;
            }
        }
    }
    
    return 0;
}

// JACK shutdown callback
static void shutdown_callback(jack_status_t code, const char *reason, void *arg) {
    Engine *engine = (Engine *)arg;
    engine->running = false;
    fprintf(stderr, "JACK shutdown: %s\n", reason);
}

int engine_init(Engine *engine, const char *client_name) {
    if (!engine || !client_name) return -1;
    
    memset(engine, 0, sizeof(Engine));
    engine->control_channel = 0;
    engine->running = false;
    
    // Initialize clips
    for (int i = 0; i < MAX_CLIPS; i++) {
        engine->clips[i].state = CLIP_EMPTY;
        engine->clips[i].buffer = (float *)calloc(MAX_BUFFER_SIZE, sizeof(float));
        if (!engine->clips[i].buffer) {
            // Cleanup on allocation failure
            for (int j = 0; j < i; j++) {
                free(engine->clips[j].buffer);
            }
            return -1;
        }
        engine->clips[i].buffer_size = 0;
        engine->clips[i].write_position = 0;
        engine->clips[i].read_position = 0;
    }
    
    // Open JACK client
    jack_status_t status;
    engine->client = jack_client_open(client_name, JackNullOption, &status, NULL);
    if (!engine->client) {
        fprintf(stderr, "Failed to open JACK client, status = 0x%2.0x\n", status);
        return -1;
    }
    
    // Register ports
    engine->audio_in_port = jack_port_register(engine->client, "audio_in",
                                                JACK_DEFAULT_AUDIO_TYPE,
                                                JackPortIsInput, 0);
    engine->audio_out_port = jack_port_register(engine->client, "audio_out",
                                                 JACK_DEFAULT_AUDIO_TYPE,
                                                 JackPortIsOutput, 0);
    engine->midi_in_port = jack_port_register(engine->client, "midi_in",
                                               JACK_DEFAULT_MIDI_TYPE,
                                               JackPortIsInput, 0);
    engine->midi_out_port = jack_port_register(engine->client, "midi_out",
                                                JACK_DEFAULT_MIDI_TYPE,
                                                JackPortIsOutput, 0);
    
    if (!engine->audio_in_port || !engine->audio_out_port ||
        !engine->midi_in_port || !engine->midi_out_port) {
        fprintf(stderr, "Failed to register ports\n");
        engine_cleanup(engine);
        return -1;
    }
    
    // Set callbacks
    jack_set_process_callback(engine->client, process_callback, engine);
    jack_on_shutdown(engine->client, shutdown_callback, engine);
    
    // Get sample rate
    engine->sample_rate = jack_get_sample_rate(engine->client);
    
    return 0;
}

void engine_cleanup(Engine *engine) {
    if (!engine) return;
    
    if (engine->client) {
        jack_client_close(engine->client);
        engine->client = NULL;
    }
    
    for (int i = 0; i < MAX_CLIPS; i++) {
        free(engine->clips[i].buffer);
        engine->clips[i].buffer = NULL;
    }
}

int engine_start(Engine *engine) {
    if (!engine || !engine->client) return -1;
    
    if (jack_activate(engine->client) != 0) {
        fprintf(stderr, "Failed to activate JACK client\n");
        return -1;
    }
    
    engine->running = true;
    return 0;
}

void engine_stop(Engine *engine) {
    if (!engine || !engine->client) return;
    
    engine->running = false;
    jack_deactivate(engine->client);
}

void engine_trigger_clip(Engine *engine, int clip_index) {
    if (!engine || clip_index < 0 || clip_index >= MAX_CLIPS) return;
    
    Clip *clip = &engine->clips[clip_index];
    
    switch (clip->state) {
        case CLIP_EMPTY:
            // Start recording
            clip->state = CLIP_RECORDING;
            clip->write_position = 0;
            clip->buffer_size = 0;
            clip->read_position = 0;
            printf("Clip %d: Recording started\n", clip_index);
            break;
            
        case CLIP_RECORDING:
            // Stop recording, start looping
            clip->state = CLIP_LOOPING;
            clip->buffer_size = clip->write_position;
            clip->read_position = 0;
            printf("Clip %d: Recording stopped, looping %zu samples\n", 
                   clip_index, clip->buffer_size);
            break;
            
        case CLIP_LOOPING:
            // Stop looping
            clip->state = CLIP_STOPPED;
            clip->read_position = 0;
            printf("Clip %d: Looping stopped\n", clip_index);
            break;
            
        case CLIP_STOPPED:
            // Start looping again
            clip->state = CLIP_LOOPING;
            clip->read_position = 0;
            printf("Clip %d: Looping resumed\n", clip_index);
            break;
    }
}

void engine_reset_clip(Engine *engine, int clip_index) {
    if (!engine || clip_index < 0 || clip_index >= MAX_CLIPS) return;
    
    Clip *clip = &engine->clips[clip_index];
    clip->state = CLIP_EMPTY;
    clip->buffer_size = 0;
    clip->write_position = 0;
    clip->read_position = 0;
    memset(clip->buffer, 0, MAX_BUFFER_SIZE * sizeof(float));
}

const char* clip_state_to_string(ClipState state) {
    switch (state) {
        case CLIP_EMPTY: return "Empty";
        case CLIP_RECORDING: return "Recording";
        case CLIP_LOOPING: return "Looping";
        case CLIP_STOPPED: return "Stopped";
        default: return "Unknown";
    }
}

uint8_t clip_state_to_velocity(ClipState state) {
    switch (state) {
        case CLIP_EMPTY: return 0;
        case CLIP_RECORDING: return 64;
        case CLIP_LOOPING: return 127;
        case CLIP_STOPPED: return 32;
        default: return 0;
    }
}