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

#include "heatshrink_encoder.h"
#include "heatshrink_decoder.h"

#include <rtthread.h>

#define DBG_TAG    "HEATSHRINK.demo"
#if HEATSHRINK_DEBUGGING_LOGS
#define DBG_LVL    DBG_LOG
#else
#define DBG_LVL    DBG_INFO
#endif

#include <rtdbg.h>

#define DEF_WINDOW_SZ2                  11
#define DEF_LOOKAHEAD_SZ2               4
#define DEF_DECODER_INPUT_BUFFER_SIZE   256
#define DEF_BUFFER_SIZE                 (2 * 1024)

/*
 * We have to open binary files with the O_BINARY flag on Windows. Most other
 * platforms don't differentiate between binary and non-binary files.
 */
#ifndef O_BINARY
#define O_BINARY 0
#endif

static const int version_major = HEATSHRINK_VERSION_MAJOR;
static const int version_minor = HEATSHRINK_VERSION_MINOR;
static const int version_patch = HEATSHRINK_VERSION_PATCH;
static const char author[] = HEATSHRINK_AUTHOR;
static const char url[] = HEATSHRINK_URL;

static void usage(void)
{
    fprintf(stderr, "heatshrink version %u.%u.%u by %s\n",
        version_major, version_minor, version_patch, author);
    fprintf(stderr, "Home page: %s\n\n", url);
    fprintf(stderr,
        "Usage:\n"
        "  heatshrink [-h] [-e|-d] [IN_FILE] [OUT_FILE]\n"
        "\n"
        "heatshrink compresses or decompresses byte streams using LZSS, and is\n"
        "designed especially for embedded, low-memory, and/or hard real-time\n"
        "systems.\n"
        "\n"
        " -h        print help\n"
        " -e        encode (compress)\n"
        " -d        decode (decompress)\n");
}

typedef enum { IO_READ, IO_WRITE, } IO_mode;
typedef enum { OP_ENC, OP_DEC, } Operation;

typedef struct {
    int fd;                     /* file descriptor */
    IO_mode mode;
    size_t fill;                /* fill index */
    size_t read;                /* read index */
    size_t size;
    size_t total;
    uint8_t buf[];
} io_handle;

typedef struct {
    uint8_t window_sz2;
    uint8_t lookahead_sz2;
    size_t decoder_input_buffer_size;
    size_t buffer_size;
    uint8_t verbose;
    Operation cmd;
    char *in_fname;
    char *out_fname;
    io_handle *in;
    io_handle *out;
} config;

static void report(config *cfg);

/* Open an IO handle. Returns RT_NULL on error. */
static io_handle *handle_open(char *fname, IO_mode m, size_t buf_sz) {
    io_handle *io = RT_NULL;
    io = rt_malloc(sizeof(*io) + buf_sz);
    if (io == RT_NULL) 
    { 
        return RT_NULL; 
    }
    rt_memset(io, 0, sizeof(*io) + buf_sz);
    io->fd = -1;
    io->size = buf_sz;
    io->mode = m;

    if (m == IO_READ)
    {
        if (0 == rt_strcmp("-", fname))
        {
            io->fd = STDIN_FILENO;
        }
        else
        {
            io->fd = open(fname, O_RDONLY | O_BINARY);
        }
    }
    else if (m == IO_WRITE)
    {
        if (0 == rt_strcmp("-", fname))
        {
            io->fd = STDOUT_FILENO;
        }
        else
        {
            io->fd = open(fname, O_WRONLY | O_BINARY | O_CREAT | O_TRUNC);
        }
    }

    if (io->fd == -1)
    {   /* failed to open */
        rt_free(io);
        LOG_E("open file failed!");
        return RT_NULL;
    }

    return io;
}

/* Read SIZE bytes from an IO handle and return a pointer to the content.
 * BUF contains at least size_t bytes. Returns 0 on EOF, -1 on error.
 * */
static ssize_t handle_read(io_handle *io, size_t size, uint8_t **buf)
{
    LOG_D("@ read %d", size);
    if (buf == RT_NULL)
    {
        return -1;
    }
    if (size > io->size)
    {
        fprintf(stderr, "size %d, io->size %d\n", size, io->size);
        return -1;
    }
    if (io->mode != IO_READ)
    {
        return -1;
    }

    size_t rem = io->fill - io->read;
    if (rem >= size)
    {
        *buf = &io->buf[io->read];
        return size;
    }
    else
    {   /* read and replenish */
        if (io->fd == -1)
        {   /* already closed, return what we've got */
            *buf = &io->buf[io->read];
            return rem;
        }

        rt_memmove(io->buf, &io->buf[io->read], rem);
        io->fill -= io->read;
        io->read = 0;
        ssize_t read_sz = read(io->fd, &io->buf[io->fill], io->size - io->fill);
        if (read_sz < 0)
        {
            LOG_E("read");
            return -1;
        }
        io->total += read_sz;
        if (read_sz == 0)
        {   /* EOF */
            if (close(io->fd) < 0)
            {
                LOG_E("close");
                return -1;
            }
            io->fd = -1;
        }
        io->fill += read_sz;
        *buf = io->buf;
        return io->fill > size ? size : io->fill;
    }
}

/* Drop the oldest SIZE bytes from the buffer. Returns <0 on error. */
static int handle_drop(io_handle *io, size_t size)
{
    LOG_D("@ drop %d", size);
    if (io->read + size <= io->fill)
    {
        io->read += size;
    }
    else
    {
        return -1;
    }
    if (io->read == io->fill)
    {
        io->read = 0;
        io->fill = 0;
    }
    return 0;
}

/* Sink SIZE bytes from INPUT into the io handle. Returns the number of
 * bytes written, or -1 on error. */
static ssize_t handle_sink(io_handle *io, size_t size, uint8_t *input)
{
    LOG_D("@ sink %d", size);
    if (size > io->size) 
    {
        return -1; 
    }
    if (io->mode != IO_WRITE) 
    { 
        return -1; 
    }

    if (io->fill + size > io->size) 
    {
        ssize_t written = write(io->fd, io->buf, io->fill);
        LOG_D("@ flushing %d, wrote %d", io->fill, written);
        io->total += written;
        if (written == -1) 
        { 
            LOG_E("write"); 
            return -1;
        }
        rt_memmove(io->buf, &io->buf[written], io->fill - written);
        io->fill -= written;
    }
    rt_memcpy(&io->buf[io->fill], input, size);
    io->fill += size;
    return size;
}

static void handle_close(io_handle *io)
{
    if (io->fd != -1)
    {
        if (io->mode == IO_WRITE)
        {
            ssize_t written = write(io->fd, io->buf, io->fill);
            io->total += written;
            LOG_D("@ close: flushing %d, wrote %d", io->fill, written);
            if (written == -1)
            {
                LOG_E("write error\n");
            }
        }
        close(io->fd);
        io->fd = -1;
    }
}

static void close_and_report(config *cfg)
{
    handle_close(cfg->in);
    handle_close(cfg->out);
    if (cfg->verbose)
    {
        report(cfg);
    }
    rt_free(cfg->in);
    rt_free(cfg->out);
}

size_t out_sz = 4096;
static uint8_t sink_out_buf[4096];

static int encoder_sink_read(config *cfg, heatshrink_encoder *hse, uint8_t *data, size_t data_sz)
{
    rt_memset(sink_out_buf, 0, out_sz);
    size_t sink_sz = 0;
    size_t poll_sz = 0;
    HSE_sink_res sres;
    HSE_poll_res pres;
    HSE_finish_res fres;
    io_handle *out = cfg->out;

    size_t sunk = 0;
    do {
        if (data_sz > 0)
        {
            sres = heatshrink_encoder_sink(hse, &data[sunk], data_sz - sunk, &sink_sz);
            if (sres < 0)
            {
                LOG_E("sink");
                return -1;
            }
            sunk += sink_sz;
        }
        
        do {
            pres = heatshrink_encoder_poll(hse, sink_out_buf, out_sz, &poll_sz);
            if (pres < 0)
            {
                LOG_E("poll");
                return -1;
            }
            if (handle_sink(out, poll_sz, sink_out_buf) < 0)
            {
                LOG_E("handle_sink");
                return -1;
            }
        } while (pres == HSER_POLL_MORE);
        
        if (poll_sz == 0 && data_sz == 0)
        {
            fres = heatshrink_encoder_finish(hse);
            if (fres < 0)
            {
                LOG_E("finish");
                return -1;
            }
            if (fres == HSER_FINISH_DONE)
            {
                return 1;
            }
        }
    } while (sunk < data_sz);
    return 0;
}

static int encode(config *cfg)
{
    uint8_t window_sz2 = cfg->window_sz2;
    size_t window_sz = 1 << window_sz2;
    heatshrink_encoder *hse = heatshrink_encoder_alloc(window_sz2, cfg->lookahead_sz2);
    if (hse == RT_NULL)
    {
        LOG_E("failed to init encoder: bad settings");
        return -1;
    }
    ssize_t read_sz = 0;
    io_handle *in = cfg->in;

    /* Process input until end of stream */
    while (1)
    {
        uint8_t *input = RT_NULL;
        read_sz = handle_read(in, window_sz, &input);
        if (input == RT_NULL)
        {
            fprintf(stderr, "handle read failure\n");
            goto __exit;
        }
        if (read_sz < 0)
        {
            break;
        }

        /* Pass read to encoder and check if input is fully processed. */
        if (encoder_sink_read(cfg, hse, input, read_sz))
        {
            break;
        }
        if (handle_drop(in, read_sz) < 0)
        {
            LOG_E("drop");
            goto __exit;
        }
    };

    if (read_sz == -1)
    {
        LOG_E("read error!\n");
    }

__exit:
    if (hse)
    {
        heatshrink_encoder_free(hse);
    }
    close_and_report(cfg);

    return 0;
}

uint8_t decode_out_buf[4096];
static int decoder_sink_read(config *cfg, heatshrink_decoder *hsd, uint8_t *data, size_t data_sz)
{
    io_handle *out = cfg->out;
    size_t sink_sz = 0;
    size_t poll_sz = 0;
    rt_memset(decode_out_buf, 0, out_sz);

    HSD_sink_res sres = {0};
    HSD_poll_res pres = {0};
    HSD_finish_res fres = {0};
    size_t sunk = 0;

    do {
        if (data_sz > 0)
        {
            sres = heatshrink_decoder_sink(hsd, &data[sunk], data_sz - sunk, &sink_sz);
            if (sres < 0)
            {
                LOG_E("sink");
                return -1;
            }
            sunk += sink_sz;
        }

        do {
            pres = heatshrink_decoder_poll(hsd, decode_out_buf, out_sz, &poll_sz);
            if (pres < 0)
            {
                LOG_E("poll");
                return -1;
            }
            if (handle_sink(out, poll_sz, decode_out_buf) < 0)
            {
                LOG_E("handle_sink");
                return -1;
            }
        } while (pres == HSDR_POLL_MORE);
        
        if (data_sz == 0 && poll_sz == 0)
        {
            fres = heatshrink_decoder_finish(hsd);
            if (fres < 0)
            {
                LOG_E("finish");
                return -1;
            }
            if (fres == HSDR_FINISH_DONE)
            {
                return 1;
            }
        }
    } while (sunk < data_sz);

    return 0;
}

static int decode(config *cfg)
{
    uint8_t window_sz2 = cfg->window_sz2;
    size_t window_sz = 1 << window_sz2;
    size_t ibs = cfg->decoder_input_buffer_size;
    heatshrink_decoder *hsd = heatshrink_decoder_alloc(ibs, window_sz2, cfg->lookahead_sz2);
    if (hsd == RT_NULL)
    {
        LOG_E("failed to init decoder");
        return -1;
    }

    ssize_t read_sz = 0;
    io_handle *in = cfg->in;
    HSD_finish_res fres;

    /* Process input until end of stream */
    while (1)
    {
        uint8_t *input = RT_NULL;
        read_sz = handle_read(in, window_sz, &input);
        if (input == RT_NULL)
        {
            fprintf(stderr, "handle read failure\n");
            goto __exit;
        }
        if (read_sz == 0)
        {
            fres = heatshrink_decoder_finish(hsd);
            if (fres < 0)
            {
                goto __exit;
            }
            if (fres == HSDR_FINISH_DONE)
            {
                break;
            }
        }
        else if (read_sz < 0)
        {
            LOG_E("read");
            goto __exit;
        }
        else
        {
            if (decoder_sink_read(cfg, hsd, input, read_sz))
            {
                break;
            }
            if (handle_drop(in, read_sz) < 0)
            {
                LOG_E("drop");
                goto __exit;
            }
        }
    }
    if (read_sz == -1)
    {
        LOG_E("read");
    }

__exit:
    if (hsd)
    {
        heatshrink_decoder_free(hsd);
    }
    close_and_report(cfg);
    return 0;
}

static void report(config *cfg)
{
    size_t inb = cfg->in->total;
    size_t outb = cfg->out->total;
    fprintf(cfg->out->fd == STDOUT_FILENO ? stderr : stdout,
        "%s %0.2f %%\t %d -> %d (-w %u -l %u)\n",
        cfg->in_fname, 100.0 - (100.0 * outb) / inb, inb, outb,
        cfg->window_sz2, cfg->lookahead_sz2);
}

int heatshrink_example(int argc, char **argv)
{
    int result = 0;
    config cfg = {0};

    cfg.window_sz2 = DEF_WINDOW_SZ2;
    cfg.lookahead_sz2 = DEF_LOOKAHEAD_SZ2;
    cfg.buffer_size = DEF_BUFFER_SIZE;
    cfg.decoder_input_buffer_size = DEF_DECODER_INPUT_BUFFER_SIZE;
    cfg.cmd = 0xFF;
    cfg.verbose = 0;

    if (rt_strcmp(argv[1], "-e") == 0)
    {
        cfg.cmd = OP_ENC;
    }
    else if (rt_strcmp(argv[1], "-d") == 0)
    {
        cfg.cmd = OP_DEC;
    }
    else
    {
        usage();
        return -1;
    }

    if (argc < 4)
    {
        usage();
        return -1;
    }

    cfg.in_fname  = argv[2];
    cfg.out_fname = argv[3];

    rt_tick_t start_tick;

    start_tick = rt_tick_get();

    if (0 == rt_strcmp(cfg.in_fname, cfg.out_fname) && (0 != rt_strcmp("-", cfg.in_fname)))
    {
        fprintf(stderr, "Refusing to overwrite file '%s' with itself.\n", cfg.in_fname);
        return -1;
    }
    /* open in file */
    cfg.in = handle_open(cfg.in_fname, IO_READ, cfg.buffer_size);
    if (cfg.in == RT_NULL)
    {
        LOG_E("Failed to open input file for read");
        result = -1;
        goto __exit;
    }

    cfg.out = handle_open(cfg.out_fname, IO_WRITE, cfg.buffer_size);
    if (cfg.out == RT_NULL)
    {
        LOG_E("Failed to open output file for write");
        result = -1;
        goto __exit;
    }

    if (cfg.cmd == OP_ENC)
    {
        result = encode(&cfg);
        rt_kprintf("encode speed tick: %d\n", rt_tick_get() - start_tick);
        return result;
    }
    else if (cfg.cmd == OP_DEC)
    {
        result = decode(&cfg);
        rt_kprintf("decode speed tick: %d\n", rt_tick_get() - start_tick);
        return result;
    }
    else
    {
        usage();
    }

__exit:
    if (cfg.in)
    {
        handle_close(cfg.in);
        rt_free(cfg.in);
    }
    if (cfg.out)
    {
        handle_close(cfg.out);
        rt_free(cfg.out);
    }

    rt_kprintf("failed\n");

    return result;
}
MSH_CMD_EXPORT(heatshrink_example, heatshrink example)