#include <libcbd.h>

#include <cbdutil.h>

#include <lz4.h>
/* XXX */
typedef off_t off64_t;
#include <zlib.h>

struct check_state
{
    int                 fd;
    tristate_t          auto_response;
    bool                check_lblk_data;
    bool                clean;
    u64                 pblk_used;
    u64                 lblk_used;
    u8**                pbatv;

    u8*                 compress_buf;
    z_stream            zlib_dstream;
};

static void
pbat_read(int fd, const struct cbd_params* params, u32 zone, u8* data)
{
    pblk_read(fd, pblk_size(params), pbat_off(params, zone), pbat_len(params), data);
}

static void
pbat_write(int fd, const struct cbd_params* params, u32 zone, const u8* data)
{
    pblk_write(fd, pblk_size(params), pbat_off(params, zone), pbat_len(params), data);
}

static void
lbat_read(int fd, const struct cbd_params* params, u32 zone, u8* data)
{
    pblk_read(fd, pblk_size(params), lbat_off(params, zone), lbat_len(params), data);
}

static void
lbat_write(int fd, const struct cbd_params* params, u32 zone, const u8* data)
{
    pblk_write(fd, pblk_size(params), lbat_off(params, zone), lbat_len(params), data);
}

static bool
check_decompress_lz4(struct check_state* state, const struct cbd_params* params, u8* buf, u32 clen)
{
    int ret;
    u32 dlen = lblk_size(params);

    ret = LZ4_decompress_safe((const char*)buf, (char*)state->compress_buf, clen, dlen);
    if (ret != dlen) {
        return false;
    }

    return true;
}

static bool
check_decompress_zlib(struct check_state* state, const struct cbd_params* params, u8* buf, u32 clen)
{
    int ret;
    z_stream* stream;
    u32 dlen = lblk_size(params);

    stream = &state->zlib_dstream;
    ret = inflateReset(stream);
    assert(ret == Z_OK);
    stream->next_in = buf;
    stream->avail_in = clen;
    stream->next_out = state->compress_buf;
    stream->avail_out = dlen;
    ret = inflate(stream, Z_SYNC_FLUSH);
    if (ret == Z_OK && !stream->avail_in && stream->avail_out) {
        u8 zerostuff = 0;
        stream->next_in = &zerostuff;
        stream->avail_in = 1;
        ret = inflate(stream, Z_FINISH);
    }
    if (ret != Z_STREAM_END || stream->total_out != dlen) {
        return false;
    }

    return true;
}

static bool
check_decompress(struct check_state* state, const struct cbd_params* params, u8* buf, u32 clen)
{
    bool ret = false;

    switch (cbd_compression_alg_get(params)) {
    case CBD_ALG_LZ4:
        ret = check_decompress_lz4(state, params, buf, clen);
        break;
    case CBD_ALG_ZLIB:
        ret = check_decompress_zlib(state, params, buf, clen);
        break;
    default:
        ret = false;
    }

    return ret;
}

static bool
check_lblk_data(struct check_state* state,
        const struct cbd_params* params,
        u64 lblk, u8* lba)
{
    bool ret = false;
    u8* data;
    u8* buf;
    u32 len;
    u32 n_alloc;
    u32 n;
    u64 pblk;

    data = calloc(pblk_size(params), lblk_per_pblk(params));
    buf = lba;
    len = lba_len_get(params, buf);
    if (len == 0 || len == CBD_UNCOMPRESSED) {
        return true;
    }
    n_alloc = DIV_ROUND_UP(len, pblk_size(params));
    for (n = 0; n < n_alloc; ++n) {
        pblk = lba_pblk_get(params, buf, n);
        pblk_read(state->fd, pblk_size(params), pblk, 1, data + n * pblk_size(params));
    }
    ret = check_decompress(state, params, data, len);
    free(data);
    if (!ret) {
        if (ask_user_bool(state->auto_response,
                "lblk %u: failed to decompress. Clear?", lblk)) {
            memset(lba, 0, lba_len(params));
            return true;
        }
        state->clean = false;
    }

    return false;
}

static bool
check_lblk_alloc(struct check_state* state,
        const struct cbd_params* params,
        u64 lblk, u8* lba)
{
    u8* buf = lba;
    u32 len;
    u32 n_alloc;
    u32 n;
    u64 pblk;
    u32 pblk_zone;
    u32 pblk_off;

    len = lba_len_get(params, buf);
    if (!len) {
        verbose(2, "  lblk[%u]: EMPTY\n", lblk);
        return false;
    }
    if (len == CBD_UNCOMPRESSED) {
        verbose(2, "  lblk[%u]: UNCOMPRESSED\n", lblk);
    }
    else {
        verbose(2, "  lblk[%u]: len=%u\n", lblk, len);
    }
    if (len > lblk_size(params)) {
        if (ask_user_bool(state->auto_response,
                "lblk %u: length %u out of bounds. Clear?", lblk, len)) {
            memset(lba, 0, lba_len(params));
            return true;
        }
        state->clean = false;
        return false;
    }
    n_alloc = (len == CBD_UNCOMPRESSED) ?
                  lblk_per_pblk(params) :
                  DIV_ROUND_UP(len, pblk_size(params));
    for (n = 0; n < n_alloc; ++n) {
        pblk = lba_pblk_get(params, buf, n);
        if (pblk < CBD_HEADER_BLOCKS) {
            verbose(2, "    [%u] :E: Alloc in header: %lu\n", n, pblk);
            if (ask_user_bool(state->auto_response,
                    "lblk %u: alloc %u in header. Clear?", lblk, n)) {
                memset(lba, 0, lba_len(params));
                return true;
            }
            state->clean = false;
            continue;
        }
        pblk_zone = zone_for_pblk(params, pblk);
        if (pblk_zone == ZONE_NONE || pblk_zone >= params->nr_zones) {
            verbose(2, "    [%u] :E: Alloc beyond end: %lu\n", n, pblk);
            if (ask_user_bool(state->auto_response,
                    "lblk %u: alloc %u beyond end. Clear?", lblk, n)) {
                memset(lba, 0, lba_len(params));
                return true;
            }
            state->clean = false;
            continue;
        }
        if (pblk < zone_data_off(params, pblk_zone)) {
            verbose(2, "    [%u] :E: Alloc in metadata: %lu\n", n, pblk);
            if (ask_user_bool(state->auto_response,
                    "lblk %u alloc in medatada. Clear?", lblk)) {
                memset(lba, 0, lba_len(params));
                return true;
            }
            state->clean = false;
            continue;
        }
        pblk_off = pblk - zone_data_off(params, pblk_zone);
        verbose(3, "    [%u] pblk=%lu\n", n, (unsigned long)pblk);
        if (cbd_bitmap_isset(state->pbatv[pblk_zone], pblk_off)) {
            verbose(2, "    [%u] :E: Duplicate allocation for pblk %lu\n", n, (unsigned long)pblk);
            if (ask_user_bool(state->auto_response,
                    "lblk %u duplicate alloc for pblk %lu. Clear?", lblk, pblk)) {
                memset(lba, 0, lba_len(params));
                return true;
            }
            state->clean = false;
            continue;
        }
        cbd_bitmap_set(state->pbatv[pblk_zone], pblk_off);
    }
    ++state->lblk_used;
    state->pblk_used += n_alloc;

    return false;
}

static void
check_lbat(struct check_state* state, const struct cbd_params* params)
{
    u32 zone;

    for (zone = 0; zone < params->init_zones; ++zone) {
        u8* lbat = calloc(pblk_size(params), lbat_len(params));
        bool zone_empty = true;
        bool changed = false;
        u32 n;

        lbat_read(state->fd, params, zone, lbat);

        verbose(2, "Zone %u: lbat=[%lu..%lu] alloc=[%lu .. %lu]\n",
                (unsigned int)zone,
                (unsigned long)lbat_off(params, zone),
                (unsigned long)(zone_data_off(params, zone) - 1),
                (unsigned long)zone_data_off(params, zone),
                (unsigned long)(zone_off(params, zone + 1) - 1));
        for (n = 0; n < params->lblk_per_zone; ++n) {
            u8* buf = lbat + n * lba_len(params);
            if (lba_len_get(params, buf) != 0) {
                zone_empty = false;
                break;
            }
        }
        if (zone_empty) {
            verbose(2, "  [empty]\n");
            continue;
        }
        for (n = 0; n < params->lblk_per_zone; ++n) {
            u64 lblk = zone * params->lblk_per_zone + n;
            u8* buf = lbat + n * lba_len(params);
            if (check_lblk_alloc(state, params, lblk, buf)) {
                changed = true;
            }
            if (state->check_lblk_data) {
                if (check_lblk_data(state, params, lblk, buf)) {
                    changed = true;
                }
            }
        }
        if (changed) {
            lbat_write(state->fd, params, zone, lbat);
        }
        free(lbat);
    }
}

static void
check_pbat(struct check_state* state, const struct cbd_params* params)
{
    u32 zone;
    u8* pbat;

    pbat = malloc(pblk_size(params) * pbat_len(params));
    for (zone = 0; zone < params->init_zones; ++zone) {
        bool changed = false;
        pbat_read(state->fd, params, zone, pbat);
        if (memcmp(pbat, state->pbatv[zone], pblk_size(params) * pbat_len(params)) != 0) {
            if (ask_user_bool(state->auto_response,
                    "zone %u has incorrect pbat. Fix?", zone)) {
                memcpy(pbat, state->pbatv[zone], pblk_size(params) * pbat_len(params));
                changed = true;
            }
            else {
                state->clean = false;
            }
            if (changed) {
                pbat_write(state->fd, params, zone, state->pbatv[zone]);
            }
        }
    }
    free(pbat);
}

int
cbd_check(const char* dev,
        bool force,
        tristate_t auto_response,
        bool full_check)
{
    struct check_state state;
    struct cbd_header header;
    uint8_t buf[SECTOR_SIZE];
    u32 n;

    memset(&state, 0, sizeof(state));
    state.fd = open(dev, O_RDWR);
    if (state.fd < 0) {
        error("Cannot open device\n");
    }

    verbose(1, "Reading header\n");
    pblk_read(state.fd, SECTOR_SIZE, 0, 1, buf);
    cbd_header_get(buf, &header);
    verbose(1, "Checking header\n");
    cbd_check_header(&header);
    if (header.params.init_zones > header.params.nr_zones) {
        verbose(1, "init_zones incorrect, fixing\n");
        header.params.init_zones = header.params.nr_zones;
    }
    if (!force && !(header.params.flags & CBD_FLAG_DIRTY)) {
        printf("%s: clean\n", dev);
        close(state.fd);
        return 0;
    }

    state.auto_response = auto_response;
    state.check_lblk_data = full_check;
    state.clean = true;
    state.pbatv = calloc(header.params.nr_zones, sizeof(u8*));
    for (n = 0; n < header.params.nr_zones; ++n) {
        state.pbatv[n] = calloc(pblk_size(&header.params), pbat_len(&header.params));
    }
    if (full_check) {
        state.compress_buf = malloc(lblk_size(&header.params));
        memset(state.compress_buf, 0, lblk_size(&header.params));
        memset(&state.zlib_dstream, 0, sizeof(z_stream));
        inflateInit2(&state.zlib_dstream, MAX_WBITS);
    }

    verbose(1, "Checking lbat\n");
    check_lbat(&state, &header.params);
    verbose(1, "Checking pbat\n");
    check_pbat(&state, &header.params);

    free(state.compress_buf);
    for (n = 0; n < header.params.nr_zones; ++n) {
        free(state.pbatv[n]);
    }
    free(state.pbatv);

    if (state.clean) {
        if (state.pblk_used != header.stats.pblk_used) {
            verbose(1, "pblk used incorrect (%lu expected, %lu found), fixing\n",
                    (unsigned long)state.pblk_used,
                    (unsigned long)header.stats.pblk_used);
            header.stats.pblk_used = state.pblk_used;
        }
        if (state.lblk_used != header.stats.lblk_used) {
            verbose(1, "lblk used incorrect (%lu expected, %lu found), fixing\n",
                    (unsigned long)state.lblk_used,
                    (unsigned long)header.stats.lblk_used);
            header.stats.lblk_used = state.lblk_used;
        }
        header.params.flags &= ~(CBD_FLAG_ERROR | CBD_FLAG_DIRTY);
        cbd_header_put(buf, &header);
        pblk_write(state.fd, SECTOR_SIZE, 0, 1, buf);
    }
    close(state.fd);

    return 0;
}