#include "pxl8_replay.h"
#include "pxl8_log.h"
#include "pxl8_sys.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

typedef struct pxl8_replay_chunk {
    u8 type;
    u32 size;
    u8* data;
    struct pxl8_replay_chunk* next;
} pxl8_replay_chunk;

typedef struct pxl8_keyframe_entry {
    pxl8_keyframe keyframe;
    pxl8_replay_chunk* input_deltas;
    struct pxl8_keyframe_entry* next;
    struct pxl8_keyframe_entry* prev;
} pxl8_keyframe_entry;

struct pxl8_replay {
    FILE* file;
    pxl8_replay_header header;
    bool recording;
    bool playing;
    u32 current_frame;

    pxl8_keyframe_entry* keyframes;
    pxl8_keyframe_entry* current_keyframe;
    u32 keyframe_count;
    u32 max_keyframes;

    pxl8_replay_chunk* pending_inputs;
    pxl8_replay_chunk* pending_inputs_tail;

    pxl8_replay_chunk* audio_events;
    pxl8_replay_chunk* audio_events_tail;
};

static void pxl8_replay_chunk_free(pxl8_replay_chunk* chunk) {
    while (chunk) {
        pxl8_replay_chunk* next = chunk->next;
        free(chunk->data);
        free(chunk);
        chunk = next;
    }
}

static void pxl8_replay_keyframe_entry_free(pxl8_keyframe_entry* entry) {
    while (entry) {
        pxl8_keyframe_entry* next = entry->next;
        pxl8_replay_chunk_free(entry->input_deltas);
        free(entry);
        entry = next;
    }
}

pxl8_replay* pxl8_replay_create(const char* path, u32 keyframe_interval) {
    FILE* f = fopen(path, "wb");
    if (!f) {
        pxl8_error("Failed to create replay file: %s", path);
        return NULL;
    }

    pxl8_replay* r = calloc(1, sizeof(pxl8_replay));
    if (!r) {
        fclose(f);
        return NULL;
    }

    r->file = f;
    r->recording = true;
    r->playing = false;
    r->header.magic = PXL8_REPLAY_MAGIC;
    r->header.version = PXL8_REPLAY_VERSION;
    r->header.keyframe_interval = keyframe_interval;

    fwrite(&r->header, sizeof(pxl8_replay_header), 1, f);
    fflush(f);

    pxl8_info("Created replay file: %s (keyframe interval: %u)", path, keyframe_interval);
    return r;
}

pxl8_replay* pxl8_replay_create_buffer(u32 keyframe_interval, u32 max_keyframes) {
    pxl8_replay* r = calloc(1, sizeof(pxl8_replay));
    if (!r) return NULL;

    r->recording = true;
    r->playing = false;
    r->max_keyframes = max_keyframes;
    r->header.magic = PXL8_REPLAY_MAGIC;
    r->header.version = PXL8_REPLAY_VERSION;
    r->header.keyframe_interval = keyframe_interval;

    pxl8_debug("Created replay buffer (keyframe interval: %u, max: %u)", keyframe_interval, max_keyframes);
    return r;
}

pxl8_replay* pxl8_replay_open(const char* path) {
    FILE* f = fopen(path, "rb");
    if (!f) {
        pxl8_error("Failed to open replay file: %s", path);
        return NULL;
    }

    pxl8_replay_header header;
    if (fread(&header, sizeof(header), 1, f) != 1) {
        pxl8_error("Failed to read replay header");
        fclose(f);
        return NULL;
    }

    if (header.magic != PXL8_REPLAY_MAGIC) {
        pxl8_error("Invalid replay magic: 0x%08X (expected 0x%08X)", header.magic, PXL8_REPLAY_MAGIC);
        fclose(f);
        return NULL;
    }

    if (header.version > PXL8_REPLAY_VERSION) {
        pxl8_error("Unsupported replay version: %u (max supported: %u)", header.version, PXL8_REPLAY_VERSION);
        fclose(f);
        return NULL;
    }

    pxl8_replay* r = calloc(1, sizeof(pxl8_replay));
    if (!r) {
        fclose(f);
        return NULL;
    }

    r->file = f;
    r->header = header;
    r->recording = false;
    r->playing = true;

    while (!feof(f)) {
        u8 chunk_type;
        if (fread(&chunk_type, 1, 1, f) != 1) break;

        if (chunk_type == PXL8_REPLAY_CHUNK_END) break;

        u8 size_bytes[3];
        if (fread(size_bytes, 3, 1, f) != 1) break;
        u32 size = size_bytes[0] | (size_bytes[1] << 8) | (size_bytes[2] << 16);

        u8* data = malloc(size);
        if (!data || fread(data, size, 1, f) != 1) {
            free(data);
            break;
        }

        if (chunk_type == PXL8_REPLAY_CHUNK_KEYFRAME) {
            pxl8_keyframe_entry* entry = calloc(1, sizeof(pxl8_keyframe_entry));
            if (entry && size >= sizeof(pxl8_keyframe)) {
                memcpy(&entry->keyframe, data, sizeof(pxl8_keyframe));
                entry->prev = r->current_keyframe;
                if (r->current_keyframe) {
                    r->current_keyframe->next = entry;
                }
                r->current_keyframe = entry;
                if (!r->keyframes) {
                    r->keyframes = entry;
                }
                r->keyframe_count++;
            }
            free(data);
        } else if (chunk_type == PXL8_REPLAY_CHUNK_INPUT) {
            if (r->current_keyframe) {
                pxl8_replay_chunk* chunk = calloc(1, sizeof(pxl8_replay_chunk));
                if (chunk) {
                    chunk->type = chunk_type;
                    chunk->size = size;
                    chunk->data = data;
                    data = NULL;

                    if (!r->current_keyframe->input_deltas) {
                        r->current_keyframe->input_deltas = chunk;
                    } else {
                        pxl8_replay_chunk* tail = r->current_keyframe->input_deltas;
                        while (tail->next) tail = tail->next;
                        tail->next = chunk;
                    }
                }
            }
            free(data);
        } else if (chunk_type == PXL8_REPLAY_CHUNK_AUDIO_EVENT) {
            pxl8_replay_chunk* chunk = calloc(1, sizeof(pxl8_replay_chunk));
            if (chunk) {
                chunk->type = chunk_type;
                chunk->size = size;
                chunk->data = data;
                data = NULL;

                if (!r->audio_events) {
                    r->audio_events = chunk;
                    r->audio_events_tail = chunk;
                } else {
                    r->audio_events_tail->next = chunk;
                    r->audio_events_tail = chunk;
                }
            }
            free(data);
        } else {
            free(data);
        }
    }

    r->current_keyframe = r->keyframes;
    pxl8_info("Opened replay: %u frames, %u keyframes", r->header.total_frames, r->keyframe_count);
    return r;
}

void pxl8_replay_destroy(pxl8_replay* r) {
    if (!r) return;

    if (r->file) {
        if (r->recording) {
            u8 end_chunk = PXL8_REPLAY_CHUNK_END;
            fwrite(&end_chunk, 1, 1, r->file);

            fseek(r->file, 0, SEEK_SET);
            fwrite(&r->header, sizeof(pxl8_replay_header), 1, r->file);
        }
        fclose(r->file);
    }

    pxl8_replay_keyframe_entry_free(r->keyframes);
    pxl8_replay_chunk_free(r->pending_inputs);
    pxl8_replay_chunk_free(r->audio_events);

    free(r);
}

bool pxl8_replay_is_recording(pxl8_replay* r) {
    return r && r->recording;
}

bool pxl8_replay_is_playing(pxl8_replay* r) {
    return r && r->playing;
}

u32 pxl8_replay_get_frame(pxl8_replay* r) {
    return r ? r->current_frame : 0;
}

u32 pxl8_replay_get_total_frames(pxl8_replay* r) {
    return r ? r->header.total_frames : 0;
}

static void write_chunk(FILE* f, u8 type, const void* data, u32 size) {
    fwrite(&type, 1, 1, f);
    u8 size_bytes[3] = {
        (u8)(size & 0xFF),
        (u8)((size >> 8) & 0xFF),
        (u8)((size >> 16) & 0xFF)
    };
    fwrite(size_bytes, 3, 1, f);
    fwrite(data, size, 1, f);
}

static void add_chunk_to_buffer(pxl8_replay* r, u8 type, const void* data, u32 size) {
    pxl8_replay_chunk* chunk = calloc(1, sizeof(pxl8_replay_chunk));
    if (!chunk) return;

    chunk->type = type;
    chunk->size = size;
    chunk->data = malloc(size);
    if (!chunk->data) {
        free(chunk);
        return;
    }
    memcpy(chunk->data, data, size);

    if (!r->pending_inputs) {
        r->pending_inputs = chunk;
        r->pending_inputs_tail = chunk;
    } else {
        r->pending_inputs_tail->next = chunk;
        r->pending_inputs_tail = chunk;
    }
}

static void pack_keys(const bool* keys, u8* packed) {
    memset(packed, 0, 32);
    for (int i = 0; i < 256; i++) {
        if (keys[i]) {
            packed[i / 8] |= (1 << (i % 8));
        }
    }
}

static void unpack_keys(const u8* packed, bool* keys) {
    for (int i = 0; i < 256; i++) {
        keys[i] = (packed[i / 8] >> (i % 8)) & 1;
    }
}

static u8 pack_mouse_buttons(const bool* buttons) {
    return (buttons[0] ? 1 : 0) | (buttons[1] ? 2 : 0) | (buttons[2] ? 4 : 0);
}

static void unpack_mouse_buttons(u8 packed, bool* buttons) {
    buttons[0] = (packed & 1) != 0;
    buttons[1] = (packed & 2) != 0;
    buttons[2] = (packed & 4) != 0;
}

void pxl8_replay_write_keyframe(pxl8_replay* r, u32 frame, f32 time, pxl8_rng* rng, pxl8_input_state* input) {
    if (!r || !r->recording) return;

    pxl8_keyframe kf = {0};
    kf.frame_number = frame;
    kf.time = time;
    kf.rng_state = rng ? rng->state : 0;

    if (input) {
        pack_keys(input->keys_down, kf.keys_down);
        kf.mouse_buttons = pack_mouse_buttons(input->mouse_buttons_down);
        kf.mouse_x = (i16)input->mouse_x;
        kf.mouse_y = (i16)input->mouse_y;
    }

    if (r->file) {
        write_chunk(r->file, PXL8_REPLAY_CHUNK_KEYFRAME, &kf, sizeof(kf));
        fflush(r->file);
    } else {
        pxl8_keyframe_entry* entry = calloc(1, sizeof(pxl8_keyframe_entry));
        if (!entry) return;

        entry->keyframe = kf;
        entry->input_deltas = r->pending_inputs;
        r->pending_inputs = NULL;
        r->pending_inputs_tail = NULL;

        if (r->keyframe_count >= r->max_keyframes && r->keyframes) {
            pxl8_keyframe_entry* oldest = r->keyframes;
            r->keyframes = oldest->next;
            if (r->keyframes) {
                r->keyframes->prev = NULL;
            }
            pxl8_replay_chunk_free(oldest->input_deltas);
            free(oldest);
            r->keyframe_count--;
        }

        entry->prev = r->current_keyframe;
        if (r->current_keyframe) {
            r->current_keyframe->next = entry;
        }
        r->current_keyframe = entry;
        if (!r->keyframes) {
            r->keyframes = entry;
        }
        r->keyframe_count++;
    }

    r->current_frame = frame;
    r->header.total_frames = frame;
}

void pxl8_replay_write_input(pxl8_replay* r, u32 frame, pxl8_input_state* prev, pxl8_input_state* curr) {
    if (!r || !r->recording || !prev || !curr) return;

    u8 buffer[256];
    u32 offset = 0;

    buffer[offset++] = (u8)(frame & 0xFF);
    buffer[offset++] = (u8)((frame >> 8) & 0xFF);
    buffer[offset++] = (u8)((frame >> 16) & 0xFF);
    buffer[offset++] = (u8)((frame >> 24) & 0xFF);

    u8 key_event_count = 0;
    u8 key_events[64];
    u32 key_offset = 0;

    for (int i = 0; i < 256 && key_event_count < 32; i++) {
        if (prev->keys_down[i] != curr->keys_down[i]) {
            key_events[key_offset++] = (u8)i;
            key_events[key_offset++] = curr->keys_down[i] ? 1 : 0;
            key_event_count++;
        }
    }

    buffer[offset++] = key_event_count;
    memcpy(buffer + offset, key_events, key_offset);
    offset += key_offset;

    u8 prev_mouse_btns = pack_mouse_buttons(prev->mouse_buttons_down);
    u8 curr_mouse_btns = pack_mouse_buttons(curr->mouse_buttons_down);

    u8 mouse_flags = 0;
    if (curr->mouse_x != prev->mouse_x || curr->mouse_y != prev->mouse_y) {
        mouse_flags |= 0x01;
    }
    if (curr_mouse_btns != prev_mouse_btns) {
        mouse_flags |= 0x02;
    }
    if (curr->mouse_wheel_x != 0 || curr->mouse_wheel_y != 0) {
        mouse_flags |= 0x04;
    }

    buffer[offset++] = mouse_flags;

    if (mouse_flags & 0x01) {
        i16 dx = (i16)(curr->mouse_x - prev->mouse_x);
        i16 dy = (i16)(curr->mouse_y - prev->mouse_y);
        buffer[offset++] = (u8)(dx & 0xFF);
        buffer[offset++] = (u8)((dx >> 8) & 0xFF);
        buffer[offset++] = (u8)(dy & 0xFF);
        buffer[offset++] = (u8)((dy >> 8) & 0xFF);
    }

    if (mouse_flags & 0x02) {
        buffer[offset++] = curr_mouse_btns;
    }

    if (mouse_flags & 0x04) {
        buffer[offset++] = (i8)curr->mouse_wheel_x;
        buffer[offset++] = (i8)curr->mouse_wheel_y;
    }

    if (r->file) {
        write_chunk(r->file, PXL8_REPLAY_CHUNK_INPUT, buffer, offset);
    } else {
        add_chunk_to_buffer(r, PXL8_REPLAY_CHUNK_INPUT, buffer, offset);
    }

    r->current_frame = frame;
    r->header.total_frames = frame;
}

void pxl8_replay_write_audio_event(pxl8_replay* r, u32 frame, u8 event_type, u8 context_id, u8 note, f32 volume) {
    if (!r || !r->recording) return;

    pxl8_audio_event evt = {
        .frame_number = frame,
        .event_type = event_type,
        .context_id = context_id,
        .note = note,
        .volume = volume
    };

    if (r->file) {
        write_chunk(r->file, PXL8_REPLAY_CHUNK_AUDIO_EVENT, &evt, sizeof(evt));
    } else {
        pxl8_replay_chunk* chunk = calloc(1, sizeof(pxl8_replay_chunk));
        if (!chunk) return;

        chunk->type = PXL8_REPLAY_CHUNK_AUDIO_EVENT;
        chunk->size = sizeof(evt);
        chunk->data = malloc(sizeof(evt));
        if (!chunk->data) {
            free(chunk);
            return;
        }
        memcpy(chunk->data, &evt, sizeof(evt));

        if (!r->audio_events) {
            r->audio_events = chunk;
            r->audio_events_tail = chunk;
        } else {
            r->audio_events_tail->next = chunk;
            r->audio_events_tail = chunk;
        }
    }
}

void pxl8_replay_close(pxl8_replay* r) {
    pxl8_replay_destroy(r);
}

bool pxl8_replay_seek_frame(pxl8_replay* r, u32 frame) {
    if (!r || !r->playing) return false;

    pxl8_keyframe_entry* target = NULL;
    for (pxl8_keyframe_entry* e = r->keyframes; e; e = e->next) {
        if (e->keyframe.frame_number <= frame) {
            target = e;
        } else {
            break;
        }
    }

    if (!target) return false;

    r->current_keyframe = target;
    r->current_frame = target->keyframe.frame_number;
    return true;
}

bool pxl8_replay_read_frame(pxl8_replay* r, pxl8_input_state* input_out) {
    if (!r || !r->playing || !r->current_keyframe) return false;

    u32 target_frame = r->current_frame + 1;

    if (target_frame > r->header.total_frames) return false;

    pxl8_replay_chunk* delta = r->current_keyframe->input_deltas;
    while (delta) {
        if (delta->size >= 5) {
            u32 delta_frame = delta->data[0] | (delta->data[1] << 8) | (delta->data[2] << 16) | (delta->data[3] << 24);
            if (delta_frame == target_frame) {
                u32 offset = 4;
                u8 key_event_count = delta->data[offset++];

                for (u8 i = 0; i < key_event_count && offset + 1 < delta->size; i++) {
                    u8 scancode = delta->data[offset++];
                    u8 pressed = delta->data[offset++];
                    input_out->keys_down[scancode] = pressed != 0;
                }

                if (offset < delta->size) {
                    u8 mouse_flags = delta->data[offset++];

                    if (mouse_flags & 0x01 && offset + 3 < delta->size) {
                        i16 dx = (i16)(delta->data[offset] | (delta->data[offset + 1] << 8));
                        i16 dy = (i16)(delta->data[offset + 2] | (delta->data[offset + 3] << 8));
                        offset += 4;
                        input_out->mouse_x += dx;
                        input_out->mouse_y += dy;
                    }

                    if (mouse_flags & 0x02 && offset < delta->size) {
                        u8 mouse_btns = delta->data[offset++];
                        unpack_mouse_buttons(mouse_btns, input_out->mouse_buttons_down);
                    }

                    if (mouse_flags & 0x04 && offset + 1 < delta->size) {
                        input_out->mouse_wheel_x = (i8)delta->data[offset++];
                        input_out->mouse_wheel_y = (i8)delta->data[offset++];
                    }
                }

                r->current_frame = target_frame;
                return true;
            }
        }
        delta = delta->next;
    }

    if (r->current_keyframe->next && r->current_keyframe->next->keyframe.frame_number <= target_frame) {
        r->current_keyframe = r->current_keyframe->next;
    }

    r->current_frame = target_frame;
    return true;
}

bool pxl8_replay_get_keyframe(pxl8_replay* r, u32 frame, pxl8_keyframe* out) {
    if (!r || !out) return false;

    pxl8_keyframe_entry* target = NULL;
    for (pxl8_keyframe_entry* e = r->keyframes; e; e = e->next) {
        if (e->keyframe.frame_number <= frame) {
            target = e;
        } else {
            break;
        }
    }

    if (!target) return false;

    *out = target->keyframe;
    return true;
}

void pxl8_replay_apply_keyframe(pxl8_replay* r, pxl8_keyframe* kf, pxl8_rng* rng, pxl8_input_state* input) {
    if (!kf) return;

    if (rng) {
        rng->state = kf->rng_state;
    }

    if (input) {
        unpack_keys(kf->keys_down, input->keys_down);
        unpack_mouse_buttons(kf->mouse_buttons, input->mouse_buttons_down);
        input->mouse_x = kf->mouse_x;
        input->mouse_y = kf->mouse_y;
    }

    if (r) {
        r->current_frame = kf->frame_number;
    }
}

bool pxl8_replay_export(pxl8_replay* r, const char* path) {
    if (!r || !r->keyframes) return false;

    FILE* f = fopen(path, "wb");
    if (!f) {
        pxl8_error("Failed to create export file: %s", path);
        return false;
    }

    pxl8_replay_header header = r->header;
    fwrite(&header, sizeof(header), 1, f);

    for (pxl8_keyframe_entry* e = r->keyframes; e; e = e->next) {
        write_chunk(f, PXL8_REPLAY_CHUNK_KEYFRAME, &e->keyframe, sizeof(pxl8_keyframe));

        for (pxl8_replay_chunk* c = e->input_deltas; c; c = c->next) {
            write_chunk(f, c->type, c->data, c->size);
        }
    }

    for (pxl8_replay_chunk* c = r->audio_events; c; c = c->next) {
        write_chunk(f, c->type, c->data, c->size);
    }

    u8 end_chunk = PXL8_REPLAY_CHUNK_END;
    fwrite(&end_chunk, 1, 1, f);

    fclose(f);
    pxl8_info("Exported replay to: %s (%u frames)", path, header.total_frames);
    return true;
}