#include <aos/aos.h>
#include <aos/paging.h>
#include <aos/ump_binding.h>
#include <aos/ump_chan.h>
#include <drivers/block_driver.h>
#include <aos/cache.h>

#include <maps/imx8x_map.h>
#include <drivers/sdhc.h>

// #define BLOCK_DRIVER_PERFORMANCE
#ifdef BLOCK_DRIVER_PERFORMANCE
#include <aos/performance.h>
struct performance_context pcontext;
#endif

struct block_driver_state {
    struct ump_send_chan *send_chan;
    struct ump_recv_chan *recv_chan;
    struct block_driver_request *current_request;
    struct block_driver_state *next;
};

static struct sdhc_s *sdhc;

static void *dma_buffer;
static genpaddr_t dma_buffer_phys;

static bool locked = false;
static struct block_driver_state *lock_queue_head = NULL;
static struct block_driver_state *lock_queue_tail = NULL;

// list of open file handles because we have to synchronize on open dirs/files vs deletion
struct handle {
    uint64_t directory_entry_id;
    size_t count;
    struct handle *next;
    struct handle *prev;
};
struct handle *handle_list_head = NULL;
struct handle *handle_list_tail = NULL;

#define BLOCK_DRIVER_CACHE
#ifdef BLOCK_DRIVER_CACHE
uint32_t current_block_in_buffer = UINT32_MAX;
#endif

static errval_t read_buffer(int block_number) {
    assert(block_number != UINT32_MAX);

    errval_t err;

    // debug_printf("[block_driver_server] read block %d\n", block_number);

    #ifdef BLOCK_DRIVER_CACHE
    // very primitive cache for the last read block
    if(current_block_in_buffer == block_number) {
        return SYS_ERR_OK;
    }
    #endif

    err = sdhc_read_block(sdhc, block_number, (lpaddr_t)dma_buffer_phys);
    if(err_is_fail(err)) {
        #ifdef BLOCK_DRIVER_CACHE
        // maybe there was a partial modification so invalidate the cache
        current_block_in_buffer = UINT32_MAX;
        #endif
        return err;
    }
    #ifdef BLOCK_DRIVER_CACHE
    current_block_in_buffer = block_number;
    #endif

    __asm volatile (
        "dmb sy\n"
    );
    cpu_dcache_wbinv_range((genvaddr_t)dma_buffer, SDHC_BLOCK_SIZE);

    return SYS_ERR_OK;
}

static errval_t write_buffer(int block_number) {
    assert(block_number != UINT32_MAX);

    errval_t err;

    // debug_printf("[block_driver_server] write block %d\n", block_number);

    __asm volatile (
        "dmb sy\n"
    );
    cpu_dcache_wbinv_range((genvaddr_t)dma_buffer, SDHC_BLOCK_SIZE);

    err = sdhc_write_block(sdhc, block_number, (lpaddr_t)dma_buffer_phys);
    if(err_is_fail(err)) return err;

    return SYS_ERR_OK;
}

static struct handle *search_element(uint64_t directory_entry_id) {
    struct handle *cur_handle = handle_list_head;
    while(cur_handle != NULL && cur_handle->directory_entry_id != directory_entry_id) {
        cur_handle = cur_handle->next;
    }
    return cur_handle;
}

static void insert_element(uint64_t directory_entry_id) {
    struct handle *element_handle = calloc(1, sizeof(struct handle));
    if (element_handle == NULL) USER_PANIC("Ran out of memory");

    element_handle->count = 1;
    element_handle->directory_entry_id = directory_entry_id;

    element_handle->prev = handle_list_tail;
    element_handle->next = NULL;
    if(handle_list_head == NULL) {
        handle_list_head = element_handle;
    } else {
        handle_list_tail->next = element_handle;
    }
    handle_list_tail = element_handle;
}

static void remove_element(struct handle *element_handle) {
    if (element_handle->prev == NULL) {
        handle_list_head = element_handle->next;
    } else {
        element_handle->prev->next = element_handle->next;
    }
    if (element_handle->next == NULL) {
        handle_list_tail = element_handle->prev;
    } else {
        element_handle->next->prev = element_handle->prev;
    }
}

__attribute__((__unused__))
static void print_handles(void) {
    struct handle *element_handle = handle_list_head;
    debug_printf("Handle List\n");
    while (element_handle != NULL) {
        debug_printf("Handle %lu: %lu\n", element_handle->directory_entry_id, element_handle->count);
        element_handle = element_handle->next;
    }
}

static errval_t register_handle(uint64_t directory_entry_id) {
    // debug_printf("Register %lu\n", directory_entry_id);
    struct handle *element_handle = search_element(directory_entry_id);
    if (element_handle == NULL) {
        insert_element(directory_entry_id);
    } else if(element_handle->count == SIZE_MAX) {
        return FAT_ERR_TOO_MANY_HANDLES;
    }else {
        ++element_handle->count;
    }
    // print_handles();

    return SYS_ERR_OK;
}

static errval_t unregister_handle(uint64_t directory_entry_id) {
    // debug_printf("Unegister %lu\n", directory_entry_id);
    struct handle *element_handle = search_element(directory_entry_id);
    if (element_handle == NULL) {
        return ERR_INVALID_ARGS;
    } else if (element_handle->count == 1) {
        remove_element(element_handle);
    } else {
        --element_handle->count;
    }
    // print_handles();
    return SYS_ERR_OK;
}

static size_t count_handles(uint64_t directory_entry_id) {
    struct handle *element_handle = search_element(directory_entry_id);
    if (element_handle == NULL) {
        return 0;
    } else {
        return element_handle->count;
    }
}

static errval_t read_object(uint32_t block_number, size_t offset, size_t size, void *dst) {
    errval_t err;

    // printf("[block_driver_server] reading %lu bytes from %lu@%u to %p\n", size, offset, block_number, dst);

    // sector bounds check
    if (offset + size > SDHC_BLOCK_SIZE) {
        return BLOCK_ERR_OVERFLOW_BLOCK;
    }

    #ifdef BLOCK_DRIVER_PERFORMANCE
    perf_init(&pcontext, "block_driver_read_object");
    perf_add_now(&pcontext, "start");
    #endif

    err = read_buffer(block_number);
    if (err_is_fail(err)) return err;

    #ifdef BLOCK_DRIVER_PERFORMANCE
    perf_add_now(&pcontext, "memcpy");
    #endif

    memcpy(dst, dma_buffer + offset, size);

    // debug_printf("Read start: '%.5s'\n", dst);

    #ifdef BLOCK_DRIVER_PERFORMANCE
    perf_add_now(&pcontext, "done");
    perf_print(&pcontext);
    #endif

    return SYS_ERR_OK;
}

static errval_t write_object(uint32_t block_number, size_t offset, size_t size, void *src) {
    errval_t err;

    // printf("[block_driver_server] writing %lu bytes from %p to %lu@%u\n", size, src, offset, block_number);

    // sector bounds check
    if (offset + size > SDHC_BLOCK_SIZE) {
        return BLOCK_ERR_OVERFLOW_BLOCK;
    }

    // we never modify the bpb and boot sector
    if (block_number == 0) {
        return BLOCK_ERR_WRITE_BOOTSECTOR;
    }

    #ifdef BLOCK_DRIVER_PERFORMANCE
    perf_init(&pcontext, "block_driver_write_object");
    perf_add_now(&pcontext, "start");
    #endif

    // if we are not writing the full block then we need to copy the current block
    if (size < SDHC_BLOCK_SIZE) {
        err = read_buffer(block_number);
        if (err_is_fail(err)) return err;
    }

    #ifdef BLOCK_DRIVER_PERFORMANCE
    perf_add_now(&pcontext, "read");
    #endif

    // debug_printf("Write start: '%.5s'\n", src);
    memcpy(dma_buffer + offset, src, size);
    
    #ifdef BLOCK_DRIVER_CACHE
    // update the cache metadata
    current_block_in_buffer = block_number;
    #endif

    #ifdef BLOCK_DRIVER_PERFORMANCE
    perf_add_now(&pcontext, "memcpy");
    #endif

    err = write_buffer(block_number);
    if (err_is_fail(err)) return err;

    #ifdef BLOCK_DRIVER_PERFORMANCE
    perf_add_now(&pcontext, "done");
    perf_print(&pcontext);
    #endif

    return SYS_ERR_OK;
}

static void handle_send_completed(void *arg, struct ump_send_queue_entry *entry) {
    free((void *)entry->header);
    free((void *)entry->payload);
    free(entry);
}

static void send_response(struct block_driver_state *state, errval_t err, size_t payload_size, void *payload) {
    assert(state != NULL);

    struct block_driver_result *result = malloc(sizeof(struct block_driver_result));
    if (result == NULL) USER_PANIC("Ran out of memory");
    struct ump_send_queue_entry *entry = malloc(sizeof(struct ump_send_queue_entry));
    if (entry == NULL) USER_PANIC("Ran out of memory");

    result->err = err;
    ump_send(
        state->send_chan,
        entry,
        sizeof(struct block_driver_result),
        result,
        payload_size,
        payload,
        handle_send_completed,
        NULL
    );
}
static void send_error_response(struct block_driver_state *state, errval_t err) {
    assert(state != NULL);

    send_response(state, err, 0, NULL);
}

static void handle_payload(void *arg, size_t payload_size, void *payload);
static void handle_request(void *arg, size_t header_size, void *header, size_t payload_size) {
    assert(arg != NULL);

    errval_t err;

    struct block_driver_state *state = arg;
    struct block_driver_request *request = header;

    state->current_request = request;

    if (request->bytes > SDHC_BLOCK_SIZE) {
        send_error_response(state, ERR_INVALID_ARGS);
    } else if (request->action == BLOCK_DRIVER_ACTION_READ) {
        // perform read directly
        void *response_buffer = malloc(request->bytes);
        if (response_buffer == NULL) USER_PANIC("Ran out of memory");

        err = read_object(request->block_number, request->offset, request->bytes, response_buffer);
        if (err_is_fail(err)) {
            send_error_response(state, err);
        } else {
            send_response(state, SYS_ERR_OK, request->bytes, response_buffer);
        }
    } else if (request->action == BLOCK_DRIVER_ACTION_WRITE) {
        if (request->bytes != payload_size) {
            send_error_response(state, ERR_INVALID_ARGS);
            ump_recv_payload(state->recv_chan, NULL, handle_payload, arg);
            return;
        }
        // we first need to receive the payload before we can handle the request
        void *payload_buffer = malloc(request->bytes);
        if (payload_buffer == NULL) USER_PANIC("Ran out of memory");
        ump_recv_payload(state->recv_chan, payload_buffer, handle_payload, arg);
        return;
    } else if (request->action == BLOCK_DRIVER_ACTION_LOCK) {
        if(locked) {
            // enqueue the request
            if (lock_queue_tail != NULL) {
                lock_queue_tail->next = state;
            } else {
                lock_queue_head = state;
            }
            lock_queue_tail = state;
        } else {
            locked = true;
            send_error_response(state, SYS_ERR_OK);
        }
    } else if (request->action == BLOCK_DRIVER_ACTION_UNLOCK) {
        if (!locked) USER_PANIC("Received unlock command but was not locked");
        send_error_response(state, SYS_ERR_OK);

        // send response to first thread on the lock queue if there is one
        if (lock_queue_head != NULL) {
            send_error_response(lock_queue_head, SYS_ERR_OK);

            // dequeue the head
            if (lock_queue_head->next == NULL) {
                lock_queue_tail = NULL;
            }
            lock_queue_head = lock_queue_head->next;
        } else {
            locked = false;
        }
    } else if (request->action == BLOCK_DRIVER_ACTION_REGISTER_HANDLE) {
        err = register_handle(request->directory_entry_id);
        send_error_response(state, err);
    } else if (request->action == BLOCK_DRIVER_ACTION_UNREGISTER_HANDLE) {
        err = unregister_handle(request->directory_entry_id);
        send_error_response(state, err);
    } else if (request->action == BLOCK_DRIVER_ACTION_COUNT_HANDLES) {
        size_t *count = malloc(sizeof(size_t));
        if (count == NULL) {
            send_error_response(state, LIB_ERR_MALLOC_FAIL);
        } else {
            *count = count_handles(request->directory_entry_id);
            send_response(state, SYS_ERR_OK, sizeof(size_t), count);
        }
    } else {
        send_error_response(state, ERR_INVALID_ARGS);
    }

    // skip payload
    ump_recv_payload(state->recv_chan, NULL, handle_payload, arg);
}

static void handle_payload(void *arg, size_t payload_size, void *payload) {
    assert(arg != NULL);

    errval_t err;

    struct block_driver_state *state = arg;
    struct block_driver_request *request = state->current_request;

    // we are guaranteed to have a valid structure request here
    if (request->action == BLOCK_DRIVER_ACTION_WRITE) {
        // we got the payload, write it to disk
        err = write_object(request->block_number, request->offset, request->bytes, payload);
        send_response(state, err, 0, NULL);
        // free the payload buffer, we don't use it anymore
        free(payload);
    }
    
    // listen for the next request on this channel
    ump_recv_header(state->recv_chan, handle_request, arg);
}

static errval_t connect_callback(void *arg, struct capref cap) {
    errval_t err;

    printf("[block_driver_server] incoming connection\n");

    struct ump_send_chan *send_chan;
    struct ump_recv_chan *recv_chan;

    // we can run the server on the default waitset since we are dispatching on
    // it for listening to connections anyways
    err = ump_chan_init(UMP_ROLE_SERVER, &send_chan, &recv_chan, sizeof(struct block_driver_request), cap, get_default_waitset());
    if (err_is_fail(err)) return err;

    struct block_driver_state *state = malloc(sizeof(struct block_driver_state));
    if (state == NULL) return LIB_ERR_MALLOC_FAIL;
    state->recv_chan = recv_chan;
    state->send_chan = send_chan;
    ump_recv_header(recv_chan, handle_request, state);

    printf("[block_driver_server] connection ready to receive requests\n");

    return SYS_ERR_OK;
}

int main(int argc, char *argv[]) {

    errval_t err;

    // get the device capability from the arg cnode
    struct capref cap_arg0 = {
        .cnode = cnode_arg,
        .slot  = 0
    };
    // map device registers uncachable
    void *device_register_vaddr;
    err = paging_map_frame_attr(
        get_current_paging_state(),
        &device_register_vaddr,
        IMX8X_SDHC_SIZE,
        cap_arg0,
        VREGION_FLAGS_READ_WRITE_NOCACHE
    );
    if (err_is_fail(err)) USER_PANIC_ERR(err, "Failed to map device registers");

    // initialize the sdhc driver
    printf("[block_driver_server] initializing sdhc driver\n");
    err = sdhc_init(&sdhc, device_register_vaddr);
    if (err_is_fail(err)) USER_PANIC_ERR(err, "Failed to initialize sdhc driver");

    // create the dma buffer
    struct capref dma_buffer_frame;
    size_t size;
    err = frame_alloc(&dma_buffer_frame, SDHC_BLOCK_SIZE, &size);
    if(err_is_fail(err)) return err_push(err, LIB_ERR_FRAME_ALLOC);

    // we need the physical address for the DMA read operation
    struct frame_identity dma_buffer_frame_id;
    err = frame_identify(dma_buffer_frame, &dma_buffer_frame_id);
    if (err_is_fail(err)) return err_push(err, LIB_ERR_FRAME_IDENTIFY);
    dma_buffer_phys = dma_buffer_frame_id.base;

    // map the scratchpad frame for us
    err = paging_map_frame_attr(get_current_paging_state(), &dma_buffer, size, dma_buffer_frame, VREGION_FLAGS_READ_WRITE);
    if (err_is_fail(err)) return err_push(err, LIB_ERR_PMAP_MAP);

    printf("[block_driver_server] registering\n");
    struct ump_binding_server server;
    err = ump_binding_register(&server, UMP_SERVER_BLOCK_DRIVER, connect_callback, NULL);
    if (err_is_fail(err)) USER_PANIC_ERR(err, "Failed to register UMP server");

    // wait for connections to the server
    printf("[block_driver_server] listening\n");
    struct waitset *default_ws = get_default_waitset();
    while (true) {
        err = event_dispatch(default_ws);
        if (err_is_fail(err)) {
            DEBUG_ERR(err, "in event_dispatch");
            abort();
        }
    }

    return EXIT_SUCCESS;
}