#include <zephyr/kernel.h>
#include <zephyr/drivers/gpio.h>
#include <zephyr/logging/log.h>
#include <zephyr/drivers/uart.h>
#include "../inc/uart.h"
#include "../inc/gpio_demo.h"

LOG_MODULE_REGISTER(aram_uart, LOG_LEVEL_INF);

/* ---------------------------------------------------------------
 * UART2 - Modem (u-blox SARA)
 * --------------------------------------------------------------- */
#define UART2_NODE DT_NODELABEL(uart2)
static const struct device *const uart2_dev = DEVICE_DT_GET(UART2_NODE);

#define MSG_SIZE 128
K_MSGQ_DEFINE(uart_msgq, MSG_SIZE, 100, 4);

static char uart2_rx_buf[MSG_SIZE];
static int  uart2_rx_pos;

void serial_cb(const struct device *dev, void *user_data)
{
    uint8_t c;
    if (!uart_irq_update(uart2_dev))    { return; }
    if (!uart_irq_rx_ready(uart2_dev))  { return; }

    while (uart_fifo_read(uart2_dev, &c, 1) == 1) {
        if ((c == '\n' || c == '\r') && uart2_rx_pos > 0) {
            uart2_rx_buf[uart2_rx_pos] = '\0';
            k_msgq_put(&uart_msgq, &uart2_rx_buf, K_NO_WAIT);
            uart2_rx_pos = 0;
        } else if (uart2_rx_pos < (sizeof(uart2_rx_buf) - 1)) {
            uart2_rx_buf[uart2_rx_pos++] = c;
        }
    }
}

void print_uart(unsigned char *buf)
{
    int msg_len = strlen(buf);
    for (int i = 0; i < msg_len; i++) {
        uart_poll_out(uart2_dev, buf[i]);
    }
}

void send_command(uint8_t *cmd)
{
    print_uart(cmd);
}

int wait_for_command_response(uint8_t *expected_response, k_timeout_t timeout)
{
    uint8_t rx_buf[MSG_SIZE];
    int ret;

    while (1) {
        ret = k_msgq_get(&uart_msgq, rx_buf, timeout);
        if (ret == 0) {
            if (strstr(rx_buf, expected_response) != NULL) {
                return 0;
            }
        } else if (ret == -EAGAIN) {
            return -EAGAIN;
        } else {
            return -EINVAL;
        }
    }
    return -1;
}

int send_wait_resp(uint8_t *cmd, uint8_t *resp)
{
    send_command(cmd);
    print_uart("\r\n");
    return (wait_for_command_response(resp, K_MSEC(100)) == 0) ? 0 : 1;
}

int get_msg_nbr(char *str)
{
    int ret = -1;
    char *token = strtok(str, ":");
    if (token != NULL) {
        token = strtok(NULL, ",");
        sscanf(token, "%d", &ret);
    }
    return ret;
}

/* ---------------------------------------------------------------
 * UART3 - Console / host command interface
 * --------------------------------------------------------------- */
#define UART3_NODE DT_NODELABEL(uart3)
static const struct device *const uart3_dev = DEVICE_DT_GET(UART3_NODE);

#define UART3_MSG_SIZE 128
K_MSGQ_DEFINE(uart3_msgq, UART3_MSG_SIZE, 10, 4);

static char uart3_rx_buf[UART3_MSG_SIZE];
static int  uart3_rx_pos;

static void uart3_serial_cb(const struct device *dev, void *user_data)
{
    uint8_t c;
    if (!uart_irq_update(uart3_dev))   { return; }
    if (!uart_irq_rx_ready(uart3_dev)) { return; }

    while (uart_fifo_read(uart3_dev, &c, 1) == 1) {
        if ((c == '\n' || c == '\r') && uart3_rx_pos > 0) {
            uart3_rx_buf[uart3_rx_pos] = '\0';
            LOG_INF("UART3 enqueue: %s", uart3_rx_buf);
            k_msgq_put(&uart3_msgq, &uart3_rx_buf, K_NO_WAIT);
            uart3_rx_pos = 0;
        } else if (uart3_rx_pos < (UART3_MSG_SIZE - 1)) {
            uart3_rx_buf[uart3_rx_pos++] = c;
        }
    }
}

void print_uart3(const char *buf)
{
    int msg_len = strlen(buf);
    for (int i = 0; i < msg_len; i++) {
        uart_poll_out(uart3_dev, buf[i]);
    }
}

/* ---------------------------------------------------------------
 * uart_init - initialise both UART2 and UART3
 * --------------------------------------------------------------- */
int uart_init(void)
{
    /* UART2 - modem */
    if (!device_is_ready(uart2_dev)) {
        LOG_ERR("UART2 device not ready");
        return -1;
    }
    int ret = uart_irq_callback_user_data_set(uart2_dev, serial_cb, NULL);
    if (ret < 0) {
        LOG_ERR("UART2 IRQ callback set failed: %d", ret);
        return ret;
    }
    uart_irq_rx_enable(uart2_dev);
    LOG_INF("UART2 (modem) initialised");

    /* UART3 - console */
    if (!device_is_ready(uart3_dev)) {
        LOG_ERR("UART3 device not ready");
        return -1;
    }
    ret = uart_irq_callback_user_data_set(uart3_dev, uart3_serial_cb, NULL);
    if (ret < 0) {
        LOG_ERR("UART3 IRQ callback set failed: %d", ret);
        return ret;
    }
    uart_irq_rx_enable(uart3_dev);
    LOG_INF("UART3 (console) initialised");

    return 0;
}

int modem_startup_sequence(void)
{
    int ret;
    
    // Array of commands to simplify the loop, or call them sequentially:
    const char *cmds[] = {
        CMD_KUSBCOMP,
        CMD_TRACE_QUERY,
        CMD_TRACE_AT,
        CMD_CFUN_RESET,
        CMD_TRACE_CUST,
        CMD_CFUN_RESET
    };

    for (int i = 0; i < ARRAY_SIZE(cmds); i++) {
        LOG_INF("Sending: %s", cmds[i]);
        
        // Using a 2-second timeout for modem processing
        ret = send_wait_resp((uint8_t *)cmds[i], (uint8_t *)RESP_OK);
        
        if (ret != 0) {
            LOG_WRN("Command %s failed or timed out", cmds[i]);
            // Depending on requirements, you might want to 'return ret' here to abort
        } else {
            LOG_INF("Command %s successful", cmds[i]);
        }
        
        // Optional: Small delay between commands to allow modem buffer to clear
        k_msleep(100); 
    }

    return 0;
}