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

// Forward declarations
static void process_queued_triggers(Engine *engine, jack_nframes_t current_frame);
static jack_nframes_t get_next_quantize_frame(Engine *engine, jack_nframes_t current_frame);

// JACK process callback
static int process_callback(jack_nframes_t nframes, void *arg) {
    Engine *engine = (Engine *)arg;
    
    // Process commands from frontend threads
    engine_process_commands(engine);
    
    // Get per-channel audio buffers
    jack_default_audio_sample_t *audio_in[MAX_CHANNELS];
    jack_default_audio_sample_t *audio_out[MAX_CHANNELS];
    
    for (int ch = 0; ch < MAX_CHANNELS; ch++) {
        audio_in[ch] = (jack_default_audio_sample_t *)
            jack_port_get_buffer(engine->audio_in_ports[ch], nframes);
        audio_out[ch] = (jack_default_audio_sample_t *)
            jack_port_get_buffer(engine->audio_out_ports[ch], nframes);
    }
    
    // Get MIDI buffers
    void *midi_in_buf = jack_port_get_buffer(engine->midi_in_port, nframes);
    void *midi_scene_buf = jack_port_get_buffer(engine->midi_scene_in_port, nframes);
    void *midi_clock_buf = jack_port_get_buffer(engine->midi_clock_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 clock input
    jack_midi_event_t midi_event;
    jack_nframes_t event_index = 0;
    
    while (jack_midi_event_get(&midi_event, midi_clock_buf, event_index) == 0) {
        event_index++;
        
        uint8_t *data = midi_event.buffer;
        uint8_t status = data[0];
        
        if (status == 0xF8) {  // MIDI Clock
            engine->transport.clock_count++;
            engine->transport.sample_position = 
                (engine->transport.clock_count * engine->sample_rate * 4) / 
                (MIDI_CLOCKS_PER_BEAT * BEATS_PER_BAR);
            
            // Update atomic mirrors for frontend reads
            atomic_store(&engine->transport_clock_count, engine->transport.clock_count);
            atomic_store(&engine->transport_sample_position, engine->transport.sample_position);
            
            if (engine->transport.clock_count % MIDI_CLOCKS_PER_BEAT == 0) {
                engine->transport.beat_position = 
                    (engine->transport.beat_position + 1) % BEATS_PER_BAR;
                atomic_store(&engine->transport_beat_position, engine->transport.beat_position);
                if (engine->transport.beat_position == 0) {
                    engine->transport.bar_position++;
                    atomic_store(&engine->transport_bar_position, engine->transport.bar_position);
                }
            }
        } else if (status == 0xFA) {  // MIDI Start
            engine->transport.rolling = true;
            engine->transport.clock_count = 0;
            engine->transport.beat_position = 0;
            engine->transport.bar_position = 0;
            engine->transport.sample_position = 0;
            atomic_store(&engine->transport_rolling, 1);
            atomic_store(&engine->transport_clock_count, 0);
            atomic_store(&engine->transport_beat_position, 0);
            atomic_store(&engine->transport_bar_position, 0);
            atomic_store(&engine->transport_sample_position, 0);
        } else if (status == 0xFC) {  // MIDI Stop
            engine->transport.rolling = false;
            atomic_store(&engine->transport_rolling, 0);
        } else if (status == 0xFB) {  // MIDI Continue
            engine->transport.rolling = true;
            atomic_store(&engine->transport_rolling, 1);
        }
        
        // Pass through clock 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 control channel MIDI input (clip triggers)
    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;
            
            // Read quantize mode atomically (frontend may update it)
            QuantizeMode current_quantize = (QuantizeMode)atomic_load(&engine->quantize_mode_atomic);
            
            if (current_quantize != QUANTIZE_OFF && engine->transport.rolling) {
                // Queue for quantization
                jack_nframes_t trigger_time = midi_event.time;
                queue_trigger(engine, clip_index, false, trigger_time);
            } else {
                // Trigger immediately
                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 scene launch MIDI input
    event_index = 0;
    while (jack_midi_event_get(&midi_event, midi_scene_buf, event_index) == 0) {
        event_index++;
        
        uint8_t *data = midi_event.buffer;
        uint8_t status = data[0] & 0xF0;
        uint8_t note = data[1];
        uint8_t velocity = data[2];
        
        // Process note on messages (any channel) for scene launch
        if (status == 0x90 && velocity > 0) {
            int scene_index = note % MAX_SCENES;
            
            // Read quantize mode atomically (frontend may update it)
            QuantizeMode current_quantize = (QuantizeMode)atomic_load(&engine->quantize_mode_atomic);
            
            if (current_quantize != QUANTIZE_OFF && engine->transport.rolling) {
                // Queue for quantization
                jack_nframes_t trigger_time = midi_event.time;
                queue_trigger(engine, scene_index, true, trigger_time);
            } else {
                // Trigger immediately
                engine_trigger_scene(engine, scene_index);
            }
        }
    }
    
    // Process queued triggers at quantization boundaries
    process_queued_triggers(engine, nframes);
    
    // Process audio per-channel
    for (int ch = 0; ch < MAX_CHANNELS; ch++) {
        memset(audio_out[ch], 0, sizeof(jack_default_audio_sample_t) * nframes);
        
        for (jack_nframes_t i = 0; i < nframes; i++) {
            // Record input to recording clips in this channel
            for (int s = 0; s < MAX_SCENES; s++) {
                int clip_idx = CLIP_INDEX(s, ch);
                Clip *clip = &engine->clips[clip_idx];
                
                if (clip->state == CLIP_RECORDING) {
                    if (clip->write_position < MAX_BUFFER_SIZE) {
                        clip->buffer[clip->write_position++] = audio_in[ch][i];
                    } else {
                        // Buffer full, stop recording
                        clip->state = CLIP_LOOPING;
                        clip->buffer_size = clip->write_position;
                        clip->read_position = 0;
                    }
                }
                
                // Play looping clips to this channel's output
                if (clip->state == CLIP_LOOPING && clip->buffer_size > 0) {
                    audio_out[ch][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(void *arg) {
    Engine *engine = (Engine *)arg;
    engine->running = false;
    fprintf(stderr, "JACK shutdown\n");
}

// Get the next quantization boundary frame
static jack_nframes_t get_next_quantize_frame(Engine *engine, jack_nframes_t current_frame) {
    if (!engine->transport.rolling || engine->quantize_mode == QUANTIZE_OFF) {
        return current_frame;
    }
    
    // Calculate frames per beat
    jack_nframes_t frames_per_beat = engine->sample_rate * 60 / 120;  // Assume 120 BPM from clock
    if (engine->transport.clock_count > 0) {
        // Derive from actual clock
        frames_per_beat = (engine->transport.sample_position * MIDI_CLOCKS_PER_BEAT) / 
                          (engine->transport.clock_count / MIDI_CLOCKS_PER_BEAT + 1);
    }
    
    jack_nframes_t frames_per_bar = frames_per_beat * BEATS_PER_BAR;
    
    // Current position in frames
    jack_nframes_t current_pos = engine->transport.sample_position + current_frame;
    
    if (engine->quantize_mode == QUANTIZE_BEAT) {
        // Next beat boundary
        jack_nframes_t beat_frames = frames_per_beat;
        jack_nframes_t next_beat = ((current_pos / beat_frames) + 1) * beat_frames;
        return next_beat - engine->transport.sample_position;
    } else {  // QUANTIZE_BAR
        // Next bar boundary
        jack_nframes_t bar_frames = frames_per_bar;
        jack_nframes_t next_bar = ((current_pos / bar_frames) + 1) * bar_frames;
        return next_bar - engine->transport.sample_position;
    }
}

// Queue a trigger for quantization
void queue_trigger(Engine *engine, int clip_index, bool is_scene, jack_nframes_t time) {
    if (!engine) return;
    
    QueuedTrigger *qt = (QueuedTrigger *)malloc(sizeof(QueuedTrigger));
    if (!qt) return;
    
    qt->clip_index = clip_index;
    qt->is_scene = is_scene;
    qt->trigger_time = time;
    qt->next = NULL;
    
    // Add to end of queue
    if (!engine->queued_triggers) {
        engine->queued_triggers = qt;
    } else {
        QueuedTrigger *last = engine->queued_triggers;
        while (last->next) last = last->next;
        last->next = qt;
    }
}

// Process queued triggers at quantization boundaries
static void process_queued_triggers(Engine *engine, jack_nframes_t nframes) {
    if (!engine->queued_triggers || !engine->transport.rolling) return;
    
    jack_nframes_t quantize_frame = get_next_quantize_frame(engine, 0);
    
    // Check if we've reached the quantization boundary
    if (quantize_frame <= nframes) {
        QueuedTrigger *qt = engine->queued_triggers;
        engine->queued_triggers = NULL;
        
        while (qt) {
            if (qt->is_scene) {
                engine_trigger_scene(engine, qt->clip_index);
            } else {
                engine_trigger_clip(engine, qt->clip_index);
            }
            
            QueuedTrigger *next = qt->next;
            free(qt);
            qt = next;
        }
    }
}

// Initialize command queue
static void command_queue_init(CommandQueue *q) {
    atomic_store(&q->write_index, 0);
    atomic_store(&q->read_index, 0);
}

// Submit command from frontend thread (non-blocking)
int engine_submit_command(Engine *engine, CommandType type, int index, jack_nframes_t value) {
    if (!engine) return -1;
    
    CommandQueue *q = &engine->command_queue;
    unsigned int write = atomic_load(&q->write_index);
    unsigned int read = atomic_load(&q->read_index);
    
    // Check if queue is full
    if ((write - read) >= MAX_QUEUED_COMMANDS) {
        fprintf(stderr, "Command queue full, dropping command\n");
        return -1;
    }
    
    unsigned int slot = write % MAX_QUEUED_COMMANDS;
    q->buffer[slot].type = type;
    q->buffer[slot].index = index;
    q->buffer[slot].value = value;
    
    // Memory barrier ensures buffer write completes before write_index update
    atomic_store(&q->write_index, write + 1);
    
    return 0;
}

// Process pending commands (called from audio thread only)
void engine_process_commands(Engine *engine) {
    if (!engine) return;
    
    CommandQueue *q = &engine->command_queue;
    unsigned int write = atomic_load(&q->write_index);
    unsigned int read = atomic_load(&q->read_index);
    
    while (read < write) {
        unsigned int slot = read % MAX_QUEUED_COMMANDS;
        Command cmd = q->buffer[slot];
        
        // Process the command directly (we're in the audio thread)
        switch (cmd.type) {
            case CMD_TRIGGER_CLIP: {
                if (cmd.index < 0 || cmd.index >= MAX_CLIPS) break;
                Clip *clip = &engine->clips[cmd.index];
                
                switch (clip->state) {
                    case CLIP_EMPTY:
                        clip->state = CLIP_RECORDING;
                        clip->write_position = 0;
                        clip->buffer_size = 0;
                        clip->read_position = 0;
                        printf("Clip %d (scene %d, channel %d): Recording started\n", 
                               cmd.index, cmd.index / MAX_CHANNELS, cmd.index % MAX_CHANNELS);
                        break;
                    case CLIP_RECORDING:
                        clip->state = CLIP_LOOPING;
                        clip->buffer_size = clip->write_position;
                        clip->read_position = 0;
                        printf("Clip %d (scene %d, channel %d): Recording stopped, looping %zu samples\n", 
                               cmd.index, cmd.index / MAX_CHANNELS, cmd.index % MAX_CHANNELS,
                               clip->buffer_size);
                        break;
                    case CLIP_LOOPING:
                        clip->state = CLIP_STOPPED;
                        clip->read_position = 0;
                        printf("Clip %d (scene %d, channel %d): Looping stopped\n", 
                               cmd.index, cmd.index / MAX_CHANNELS, cmd.index % MAX_CHANNELS);
                        break;
                    case CLIP_STOPPED:
                        clip->state = CLIP_LOOPING;
                        clip->read_position = 0;
                        printf("Clip %d (scene %d, channel %d): Looping resumed\n", 
                               cmd.index, cmd.index / MAX_CHANNELS, cmd.index % MAX_CHANNELS);
                        break;
                }
                break;
            }
            
            case CMD_TRIGGER_SCENE: {
                if (cmd.index < 0 || cmd.index >= MAX_SCENES) break;
                printf("Scene %d: Triggering all clips\n", cmd.index);
                for (int ch = 0; ch < MAX_CHANNELS; ch++) {
                    int clip_idx = CLIP_INDEX(cmd.index, ch);
                    Clip *clip = &engine->clips[clip_idx];
                    
                    switch (clip->state) {
                        case CLIP_EMPTY:
                            clip->state = CLIP_RECORDING;
                            clip->write_position = 0;
                            clip->buffer_size = 0;
                            clip->read_position = 0;
                            break;
                        case CLIP_RECORDING:
                            clip->state = CLIP_LOOPING;
                            clip->buffer_size = clip->write_position;
                            clip->read_position = 0;
                            break;
                        case CLIP_LOOPING:
                            clip->state = CLIP_STOPPED;
                            clip->read_position = 0;
                            break;
                        case CLIP_STOPPED:
                            clip->state = CLIP_LOOPING;
                            clip->read_position = 0;
                            break;
                    }
                }
                break;
            }
            
            case CMD_RESET_CLIP: {
                if (cmd.index < 0 || cmd.index >= MAX_CLIPS) break;
                Clip *clip = &engine->clips[cmd.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));
                break;
            }
            
            case CMD_SET_QUANTIZE_MODE:
                engine->quantize_mode = (QuantizeMode)cmd.index;
                break;
                
            case CMD_SET_QUANTIZE_THRESHOLD:
                engine->quantize_threshold = cmd.value;
                break;
                
            case CMD_RESET_TRANSPORT:
                engine->transport.rolling = false;
                engine->transport.clock_count = 0;
                engine->transport.beat_position = 0;
                engine->transport.bar_position = 0;
                engine->transport.sample_position = 0;
                atomic_store(&engine->transport_rolling, 0);
                atomic_store(&engine->transport_clock_count, 0);
                atomic_store(&engine->transport_beat_position, 0);
                atomic_store(&engine->transport_bar_position, 0);
                atomic_store(&engine->transport_sample_position, 0);
                break;
        }
        
        read++;
    }
    
    // Update read index after processing all commands
    atomic_store(&q->read_index, read);
}

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;
    engine->quantize_mode = QUANTIZE_OFF;
    engine->quantize_threshold = 0;
    engine->queued_triggers = NULL;
    
    // Initialize command queue
    command_queue_init(&engine->command_queue);
    
    // Initialize atomic state mirrors
    atomic_store(&engine->transport_rolling, 0);
    atomic_store(&engine->transport_clock_count, 0);
    atomic_store(&engine->transport_beat_position, 0);
    atomic_store(&engine->transport_bar_position, 0);
    atomic_store(&engine->transport_sample_position, 0);
    atomic_store(&engine->quantize_mode_atomic, (int)QUANTIZE_OFF);
    atomic_store(&engine->quantize_threshold_atomic, 0);
    
    // Initialize transport
    engine->transport.rolling = false;
    engine->transport.clock_count = 0;
    engine->transport.beat_position = 0;
    engine->transport.bar_position = 0;
    engine->transport.sample_position = 0;
    
    // 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 per-channel audio ports
    char port_name[32];
    for (int ch = 0; ch < MAX_CHANNELS; ch++) {
        snprintf(port_name, sizeof(port_name), "audio_in_%d", ch);
        engine->audio_in_ports[ch] = jack_port_register(engine->client, port_name,
                                                         JACK_DEFAULT_AUDIO_TYPE,
                                                         JackPortIsInput, 0);
        
        snprintf(port_name, sizeof(port_name), "audio_out_%d", ch);
        engine->audio_out_ports[ch] = jack_port_register(engine->client, port_name,
                                                          JACK_DEFAULT_AUDIO_TYPE,
                                                          JackPortIsOutput, 0);
        
        if (!engine->audio_in_ports[ch] || !engine->audio_out_ports[ch]) {
            fprintf(stderr, "Failed to register audio port %d\n", ch);
            engine_cleanup(engine);
            return -1;
        }
    }
    
    // Register MIDI ports
    engine->midi_in_port = jack_port_register(engine->client, "midi_control_in",
                                               JACK_DEFAULT_MIDI_TYPE,
                                               JackPortIsInput, 0);
    engine->midi_scene_in_port = jack_port_register(engine->client, "midi_scene_in",
                                                     JACK_DEFAULT_MIDI_TYPE,
                                                     JackPortIsInput, 0);
    engine->midi_clock_in_port = jack_port_register(engine->client, "midi_clock_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->midi_in_port || !engine->midi_scene_in_port || 
        !engine->midi_clock_in_port || !engine->midi_out_port) {
        fprintf(stderr, "Failed to register MIDI 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;
    
    // Free any queued triggers
    QueuedTrigger *qt = engine->queued_triggers;
    while (qt) {
        QueuedTrigger *next = qt->next;
        free(qt);
        qt = next;
    }
    engine->queued_triggers = NULL;
    
    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;
    
    // Queue command for audio thread processing
    engine_submit_command(engine, CMD_TRIGGER_CLIP, clip_index, 0);
}

void engine_trigger_scene(Engine *engine, int scene_index) {
    if (!engine || scene_index < 0 || scene_index >= MAX_SCENES) return;
    
    engine_submit_command(engine, CMD_TRIGGER_SCENE, scene_index, 0);
}

void engine_reset_clip(Engine *engine, int clip_index) {
    if (!engine || clip_index < 0 || clip_index >= MAX_CLIPS) return;
    
    engine_submit_command(engine, CMD_RESET_CLIP, clip_index, 0);
}

void engine_set_quantize_mode(Engine *engine, QuantizeMode mode) {
    if (!engine) return;
    
    // Atomically update the mode so audio thread sees it immediately
    atomic_store(&engine->quantize_mode_atomic, (int)mode);
    
    // Also queue for any additional processing
    engine_submit_command(engine, CMD_SET_QUANTIZE_MODE, (int)mode, 0);
    
    printf("Quantize mode set to: %s\n", quantize_mode_to_string(mode));
}

void engine_set_quantize_threshold(Engine *engine, jack_nframes_t samples) {
    if (!engine) return;
    
    atomic_store(&engine->quantize_threshold_atomic, samples);
    engine_submit_command(engine, CMD_SET_QUANTIZE_THRESHOLD, 0, samples);
}

void engine_reset_transport(Engine *engine) {
    if (!engine) return;
    
    engine_submit_command(engine, CMD_RESET_TRANSPORT, 0, 0);
    printf("Transport reset\n");
}

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;
    }
}

const char* quantize_mode_to_string(QuantizeMode mode) {
    switch (mode) {
        case QUANTIZE_OFF: return "Off";
        case QUANTIZE_BEAT: return "Beat";
        case QUANTIZE_BAR: return "Bar";
        default: return "Unknown";
    }
}