#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <fcntl.h>
#include <unistd.h>
#include <linux/i2c.h>
#include <linux/i2c-dev.h>
#include <sys/stat.h>
#include <sys/ioctl.h>

#include "crypt.h"
#include "utils.h"
#include "i2c.h"

// Hardcoded key and nonce (that's insecure!!!!)
#define KEY   "0123456789ABCDE"
#define NONCE "0123456"

/* -----------------------------------------------------------------------------
 * @brief   Performs either writing or reading from I2C device
 *
 * @param   [in] ctx operation context
 * @param   [in/out] msg: msg contains buffer with information on
 *          how to store or retreive information.
 *          If msg->is_write is true, then data will be stored otherwise
 *          data will be read. In case of reading msg->buf must point
 *          to the buffer which is big enough to store 'msg->buf.sz' bytes,
 *          (otherwise memory can be corrupted).
 * @returns 1 on success, otherwise 0.
-------------------------------------------------------------------------------- */
static int i2c_op(const con_t* ctx, msg_t *msg) {
    // check if connected
    if (!ctx->dev) {
        return 0;
    }

    struct i2c_smbus_ioctl_data args;
    union i2c_smbus_data data;
    int res = 1;

    args.read_write = msg->is_write?I2C_SMBUS_WRITE:I2C_SMBUS_READ;
    args.command = msg->reg;
    args.size = I2C_SMBUS_BYTE_DATA;
    args.data = &data;

    // can't use msg->buf->sz as it's updated in case of read
    size_t data_len = msg->buf->sz;
    // Write/Read byte after byte. On write, bytes are stored
    // in i2c registers contiguously (first byte in reg=0,
    // second in reg=1, ...)
    for (size_t i=0; i<data_len; i++) {
        data.byte = msg->buf->b[i];
        args.command = msg->reg + i;
        if (ioctl(ctx->dev, I2C_SMBUS, &args) < 0) {
            return 0;
        }
        if (!msg->is_write) {
            msg->buf->b[i] = data.byte;
            msg->buf->sz = i + 1;
        }
    }

    return res;
}

int i2c_send(const con_t *ctx, const buf_t *plaintext) {
    uint8_t enc_data[MAX_BUFFER_SIZE] = {0};

    CHECK_BUF(plaintext);

    // Encrypted data format:
    // size (1 byte) | encrypted data | authentication tag (16 bytes)
    if (plaintext->sz > (MAX_BUFFER_SIZE - TAG_LEN - 1)) {
        // to much data
        return 0;
    }

    buf_t key = {
        .b = (uint8_t*)KEY,
        .sz = sizeof(KEY)
    };
    buf_t nonce = {
        .b = (uint8_t*)NONCE,
        .sz = sizeof(NONCE)
    };

    uint8_t tag[TAG_LEN] = {0};
    cop_t cryp = {.is_encrypt = 1};
    if (!crypt_begin(&cryp, &key, &nonce) ||
        !crypt_update(&cryp, plaintext, &enc_data[1]) ||
        !crypt_finish(&cryp, tag)) {
        return 0;
    }

    memcpy(&enc_data[plaintext->sz+1], tag, TAG_LEN);
    enc_data[0] = plaintext->sz + TAG_LEN;

    buf_t ciphertext = {
        .b = &enc_data[0],
        .sz = plaintext->sz + TAG_LEN + 1
    };

    msg_t msg = {
        .buf = &ciphertext,
        .is_write = 1,
        .reg = 0
    };
    return i2c_op(ctx, &msg);
end:
    return 0;
}

int i2c_recv(const con_t *ctx, buf_t *plaintext) {
    uint8_t enc_data[MAX_BUFFER_SIZE + TAG_LEN] = {0};
    cop_t c = {.is_encrypt = 0};
    int res = 0;

    CHECK_BUF(plaintext);

    // Setup needed buffers
    buf_t key = {
        .b = (uint8_t*)KEY,
        .sz = sizeof(KEY)
    };
    buf_t nonce = {
        .b = (uint8_t*)NONCE,
        .sz = sizeof(NONCE)
    };

    buf_t ciphertext = {
        .b = enc_data,
        .sz = 1
    };

    msg_t msg = {
        .buf = &ciphertext,
        .is_write = 0,
        .reg = 0
    };

    // Read size of encrypted data
    if (!i2c_op(ctx, &msg)) {
        goto end;
    }

    // Set size of the data to read
    ciphertext.sz = enc_data[0];
    // Skip first register as it stores size of
    // encrypted data which already was read above
    msg.reg = 1;

    // Read encrypted data with authentication tag
    if (!i2c_op(ctx, &msg)) {
        goto end;
    }

    // Authentication tag is appended after encrypted data
    msg.buf->sz -= TAG_LEN;
    uint8_t *tag = &enc_data[msg.buf->sz];
    // Decrypt. Reuse enc_data buffer to store decryption (as that's safe)
    // TODO: This could be one function
    if (!crypt_begin(&c, &key, &nonce) ||
        !crypt_update(&c, msg.buf, msg.buf->b) ||
        !crypt_finish(&c, tag)) {
        LOG("Error occured when decrypting");
        goto end;
    }
    memcpy(plaintext->b, msg.buf->b, msg.buf->sz);
    plaintext->sz = msg.buf->sz;
    // All good now
    res = 1;
end:
    if (!res) LOG("ERROR: can't receive encrypted data");
    return res;
}

// --- Operations related to connecting to the device

int i2c_dev_open(con_t *ctx, int device_num, int slave_addr) {
    char filename[64];
    sprintf(filename, "/dev/i2c-%d", device_num);
    if ((ctx->dev = open(filename, O_RDWR)) < 0) {
        LOG("Failed to open the i2c bus %d", ctx->dev);
        return 0;
    }
    if (ioctl(ctx->dev, I2C_SLAVE, slave_addr) < 0) {
        LOG("Can't acquire bus access");
        return 0;
    }
    return 1;
}

void i2c_dev_close(con_t *ctx) {
    close(ctx->dev);
    ctx->dev = 0;
}

int i2c_is_connected(con_t *ctx) {
    // Device wasn't opened
    if (!ctx || !ctx->dev) {
        return 0;
    }

    // Check if slave reponds to reads
    uint8_t tmp[1];
    buf_t b = {
        .b = tmp,
        .sz = 1
    };
    msg_t msg = {
        .buf = &b,
        .is_write = 0,
        .reg = 0
    };

    return i2c_op(ctx, &msg);
}