#include <zephyr/kernel.h>
#include <zephyr/sys/printk.h>
#include <zephyr/net/socket.h>
#include <zephyr/net/net_ip.h>
#include <zephyr/net/openthread.h>
#include <zephyr/sys/byteorder.h>
#include <zephyr/drivers/adc.h>
#include <zephyr/drivers/sensor.h>
#include <zephyr/drivers/hwinfo.h>
#include <zephyr/drivers/regulator.h>
#include <zephyr/logging/log.h>
#include <zephyr/settings/settings.h>
#if defined(CONFIG_SHELL)
#include <zephyr/shell/shell.h>
#endif

#include <openthread/thread.h>
#include <openthread/platform/radio.h>

#include <errno.h>
#include <stdlib.h>
#include <string.h>

// #define FAST_TIMES

#define TARGET_IP "fd00:10:0:10:be24:11ff:fe12:6196"
#define TARGET_PORT 6969
#define LISTEN_PORT 6969
#define MEASURE_INTERVAL_M 5
#define SEND_DELTA_MC 10
#define MAX_SEND_INTERVAL_M 20
#define BATTERY_MEASURE_INTERVAL_H 24
#define UDP_RX_POLL_TIMEOUT_MS 100
#define IEEE802154_MAX_TX_POWER_DBM 8

#define S_TO_MS(s) ((uint64_t)(s) * 1000ULL)
#define M_TO_MS(m) ((uint64_t)(m) * 60ULL * 1000ULL)
#define H_TO_MS(h) ((uint64_t)(h) * 60ULL * 60ULL * 1000ULL)

LOG_MODULE_REGISTER(main, LOG_LEVEL_DBG);

enum payload_type
{
    PAYLOAD_TYPE_TEMP = 0x00,
    PAYLOAD_TYPE_BATTERY = 0x01,
    PAYLOAD_TYPE_CONFIG = 0x10,
};

struct header
{
    uint8_t proto_version;
    uint8_t type;
    uint8_t device_id[8];
    uint16_t seq_num;
} __attribute__((packed));

int32_t temp_mcelsius;
int32_t filtered_temperature_mc = 0xFFFF;
int16_t adc_buffer;

uint16_t measure_interval_m = MEASURE_INTERVAL_M;
uint16_t send_delta_mc = SEND_DELTA_MC;
uint16_t max_send_interval_m = MAX_SEND_INTERVAL_M;
uint16_t battery_measure_interval_h = BATTERY_MEASURE_INTERVAL_H;

int sock = -1;
struct sockaddr_in6 dest_addr;
struct header header = {.proto_version = 1, .device_id = {0}, .seq_num = 0};

static const struct adc_dt_spec adc_channel = ADC_DT_SPEC_GET(DT_PATH(zephyr_user));
static const struct device *const sensor_power_regulator = DEVICE_DT_GET(DT_NODELABEL(sensor_power));
static const struct device *const temp_dev = DEVICE_DT_GET(DT_ALIAS(pooltemp));

struct adc_sequence sequence = {
    .buffer = &adc_buffer,
    .buffer_size = sizeof(adc_buffer),
    .calibrate = true,
};

static int sensor_power_init(void)
{
    if (!device_is_ready(sensor_power_regulator))
    {
        LOG_ERR("Sensor power regulator is not ready");
        return -ENODEV;
    }

    int rc = regulator_enable(sensor_power_regulator);
    if (rc < 0)
    {
        LOG_ERR("Failed to enable sensor power regulator (%d)", rc);
        return rc;
    }

    k_msleep(10);
    return 0;
}

static int sensor_power_off(void)
{
    int rc = regulator_disable(sensor_power_regulator);
    if (rc < 0)
    {
        LOG_WRN("Failed to disable sensor power regulator (%d)", rc);
    }

    return rc;
}

static int radio_set_max_tx_power(void)
{
    otInstance *instance;
    otError err;

    openthread_mutex_lock();
    instance = openthread_get_default_instance();
    if (instance == NULL)
    {
        openthread_mutex_unlock();
        LOG_ERR("OpenThread instance not available");
        return -ENODEV;
    }

    err = otPlatRadioSetTransmitPower(instance, IEEE802154_MAX_TX_POWER_DBM);
    openthread_mutex_unlock();

    if (err != OT_ERROR_NONE)
    {
        LOG_ERR("Failed to set IEEE 802.15.4 TX power to %d dBm (%d)", IEEE802154_MAX_TX_POWER_DBM, err);
        return -EIO;
    }

    LOG_INF("IEEE 802.15.4 TX power set to %d dBm", IEEE802154_MAX_TX_POWER_DBM);
    return 0;
}

static int pool_settings_set(const char *name, size_t len, settings_read_cb read_cb, void *cb_arg)
{
    const char *next;
    size_t name_len = settings_name_next(name, &next);
    int rc;

    if (name_len == 0)
    {
        return -ENOENT;
    }

    LOG_DBG("settings_set: name='%.*s', len=%zu", (int)name_len, name, len);

    if (strncmp(name, "meas_int", name_len) == 0)
    {
        rc = read_cb(cb_arg, &measure_interval_m, sizeof(measure_interval_m));
        if (rc < 0)
        {
            LOG_ERR("Failed to read meas_int: %d", rc);
            return rc;
        }
        LOG_DBG("Read meas_int=%u", measure_interval_m);
        return 0;
    }

    if (strncmp(name, "delta_mc", name_len) == 0)
    {
        rc = read_cb(cb_arg, &send_delta_mc, sizeof(send_delta_mc));
        if (rc < 0)
        {
            LOG_ERR("Failed to read delta_mc: %d", rc);
            return rc;
        }
        LOG_DBG("Read delta_mc=%u", send_delta_mc);
        return 0;
    }

    if (strncmp(name, "max_int", name_len) == 0)
    {
        rc = read_cb(cb_arg, &max_send_interval_m, sizeof(max_send_interval_m));
        if (rc < 0)
        {
            LOG_ERR("Failed to read max_int: %d", rc);
            return rc;
        }
        LOG_DBG("Read max_int=%u", max_send_interval_m);
        return 0;
    }

    if (strncmp(name, "batt_h", name_len) == 0)
    {
        rc = read_cb(cb_arg, &battery_measure_interval_h, sizeof(battery_measure_interval_h));
        if (rc < 0)
        {
            LOG_ERR("Failed to read batt_h: %d", rc);
            return rc;
        }
        LOG_DBG("Read batt_h=%u", battery_measure_interval_h);
        return 0;
    }

    return -ENOENT;
}

struct settings_handler pool_conf_handler = {
    .name = "pool",
    .h_set = pool_settings_set,
};

static bool thread_is_attached(void)
{
    otInstance *instance;
    otDeviceRole role;

    openthread_mutex_lock();
    instance = openthread_get_default_instance();
    if (instance == NULL)
    {
        openthread_mutex_unlock();
        return false;
    }

    role = otThreadGetDeviceRole(instance);
    openthread_mutex_unlock();

    return role == OT_DEVICE_ROLE_CHILD || role == OT_DEVICE_ROLE_ROUTER || role == OT_DEVICE_ROLE_LEADER;
}

int adc_init(void)
{
    int rc;

    if (!adc_is_ready_dt(&adc_channel))
    {
        printk("Error: ADC device is not ready.\n");
        return -ENODEV;
    }

    rc = adc_channel_setup_dt(&adc_channel);
    if (rc < 0)
    {
        printk("Failed to set up ADC channel (%d)\n", rc);
        return rc;
    }

    rc = adc_sequence_init_dt(&adc_channel, &sequence);
    if (rc < 0)
    {
        printk("Failed to initialize ADC sequence (%d)\n", rc);
        return rc;
    }

    return 0;
}

int send_frame(uint8_t type, uint8_t *data, size_t data_len)
{
    int rc;
    uint8_t buffer[sizeof(header) + data_len];

    header.type = type;
    buffer[0] = header.proto_version;
    buffer[1] = header.type;
    memcpy(&buffer[2], header.device_id, sizeof(header.device_id));
    sys_put_be16(header.seq_num, &buffer[10]);
    memcpy(&buffer[12], data, data_len);

    rc = zsock_sendto(sock, buffer, sizeof(header) + data_len, 0, (struct sockaddr *)&dest_addr, sizeof(dest_addr));
    if (rc < 0)
    {
        LOG_ERR("Send failed: %d", rc);
    }
    else
    {
        LOG_INF("Packet sent | type: 0x%02x | seq: %u", header.type, header.seq_num);
        LOG_HEXDUMP_DBG(buffer, sizeof(header) + data_len, "Sent data:");
    }

    header.seq_num++;
    return rc;
}

int send_config(void)
{
    uint8_t send_buf[8];

    sys_put_be16(measure_interval_m, &send_buf[0]);
    sys_put_be16(send_delta_mc, &send_buf[2]);
    sys_put_be16(max_send_interval_m, &send_buf[4]);
    sys_put_be16(battery_measure_interval_h, &send_buf[6]);

    return send_frame(PAYLOAD_TYPE_CONFIG, send_buf, sizeof(send_buf));
}

int recv_udp(void)
{
    uint8_t recv_buf[32];
    struct sockaddr_in6 from_addr;
    socklen_t from_addr_len = sizeof(from_addr);
    struct zsock_pollfd fds[1];

    fds[0].fd = sock;
    fds[0].events = ZSOCK_POLLIN;

    LOG_DBG("Waiting briefly for incoming config data...");

    int ret = zsock_poll(fds, 1, UDP_RX_POLL_TIMEOUT_MS);
    if (ret > 0)
    {
        if (fds[0].revents & ZSOCK_POLLIN)
        {
            int len = zsock_recvfrom(sock, recv_buf, sizeof(recv_buf), 0, (struct sockaddr *)&from_addr, &from_addr_len);
            if (len < 0)
            {
                LOG_ERR("recvfrom failed: %d (errno: %d)", len, errno);
                return len;
            }

            if (len >= 12 && recv_buf[0] == header.proto_version && recv_buf[1] == PAYLOAD_TYPE_CONFIG)
            {
                uint16_t rx_seq = sys_get_be16(&recv_buf[2]);
                uint16_t new_meas_int = sys_get_be16(&recv_buf[4]);
                uint16_t new_delta_mc = sys_get_be16(&recv_buf[6]);
                uint16_t new_max_int = sys_get_be16(&recv_buf[8]);
                uint16_t new_batt_int = sys_get_be16(&recv_buf[10]);

                LOG_HEXDUMP_DBG(recv_buf, len, "Received config data:");
                LOG_INF("New config received! seq: %u", rx_seq);
                LOG_INF(" -> measurement interval: %u min", new_meas_int);
                LOG_INF(" -> temperature delta: %u mC", new_delta_mc);
                LOG_INF(" -> max interval: %u min", new_max_int);
                LOG_INF(" -> battery interval: %u h", new_batt_int);

                if (new_max_int < new_meas_int)
                {
                    LOG_WRN("Received max interval (%u) is smaller than measurement interval (%u). Ignoring max interval.", new_max_int, new_meas_int);
                    new_max_int = new_meas_int;
                }

                if (new_meas_int <= 60 && new_meas_int != measure_interval_m)
                {
                    measure_interval_m = new_meas_int;
                    settings_save_one("pool/meas_int", &measure_interval_m, sizeof(measure_interval_m));
                    LOG_INF("Measurement interval set to %u min and saved to flash.", measure_interval_m);
                }

                if (new_delta_mc > 0 && new_delta_mc <= 2000 && new_delta_mc != send_delta_mc)
                {
                    send_delta_mc = new_delta_mc;
                    settings_save_one("pool/delta_mc", &send_delta_mc, sizeof(send_delta_mc));
                    LOG_INF("Temperature delta set to %u mC and saved to flash.", send_delta_mc);
                }

                if (new_max_int >= 5 && new_max_int <= 120 && new_max_int != max_send_interval_m)
                {
                    max_send_interval_m = new_max_int;
                    settings_save_one("pool/max_int", &max_send_interval_m, sizeof(max_send_interval_m));
                    LOG_INF("Max interval set to %u min and saved to flash.", max_send_interval_m);
                }

                if (new_batt_int >= 2 && new_batt_int <= 72 && new_batt_int != battery_measure_interval_h)
                {
                    battery_measure_interval_h = new_batt_int;
                    settings_save_one("pool/batt_h", &battery_measure_interval_h, sizeof(battery_measure_interval_h));
                    LOG_INF("Battery interval set to %u h and saved to flash.", battery_measure_interval_h);
                }

                LOG_INF("Settings saved to flash.");
                send_config();
            }
            else
            {
                LOG_DBG("UDP packet ignored (len=%d, type=0x%02x)", len, len > 1 ? recv_buf[1] : 0xff);
            }
        }
    }
    else if (ret == 0)
    {
        LOG_DBG("No config data pending from the gateway (timeout).");
    }
    else
    {
        LOG_ERR("Error polling socket: %d (errno: %d)", ret, errno);
    }

    return 0;
}

#if defined(CONFIG_SHELL)
static int cmd_udp_receive_pending(const struct shell *sh, size_t argc, char **argv)
{
    ARG_UNUSED(argc);
    ARG_UNUSED(argv);

    if (sock < 0)
    {
        shell_error(sh, "UDP socket is not initialized yet.");
        return -ENOTCONN;
    }

    int rc = send_frame(0xFF, NULL, 0);
    if (rc < 0)
    {
        shell_error(sh, "Failed to send UDP ping frame (%d)", rc);
        return rc;
    }

    k_msleep(300);

    rc = recv_udp();
    if (rc < 0)
    {
        shell_error(sh, "UDP receive failed (%d)", rc);
        return rc;
    }

    shell_print(sh, "UDP ping sent, waited 300 ms, receive poll done.");
    return 0;
}

SHELL_CMD_REGISTER(udp_recv, NULL,
                   "Poll pending UDP config frames once",
                   cmd_udp_receive_pending);
#endif

int main(void)
{
    static int rc;
    static struct sensor_value temp_val;
    static uint32_t battery_mv;
    static uint8_t send_buf[8];
    static bool was_connected = false;
    static uint64_t last_temp_send_time = 0;
    static int32_t last_temp_send_mc = 0;
    static uint64_t last_battery_send_time = 0;
    uint32_t temperature_mc;

    rc = sensor_power_init();
    if (rc < 0)
    {
        return rc;
    }
    sensor_power_off();

    if (!device_is_ready(temp_dev))
    {
        LOG_ERR("Error: Temperature sensor device is not ready.");
        return -ENODEV;
    }

    LOG_INF("Temperature sensor (%s) successfully initialized.", temp_dev->name);

    rc = adc_init();
    if (rc < 0)
    {
        printk("Failed to initialize ADC (%d)\n", rc);
        return rc;
    }

    rc = settings_subsys_init();
    if (rc < 0)
    {
        LOG_ERR("Settings subsystem init failed: %d", rc);
    }
    else
    {
        rc = settings_register(&pool_conf_handler);
        if (rc < 0)
        {
            LOG_ERR("Settings register failed: %d", rc);
        }

        rc = settings_load();
        if (rc < 0)
        {
            LOG_ERR("Settings load failed: %d", rc);
        }
        else
        {
            LOG_INF("Settings successfully loaded from flash.");
        }
    }

    LOG_INF("Starting UDP sensor node...");
    k_sleep(K_SECONDS(1));

    rc = radio_set_max_tx_power();
    if (rc < 0)
    {
        return rc;
    }

    LOG_INF("Waiting for OpenThread attach...");
    for (int i = 0; i < 15; i++)
    {
        if (thread_is_attached())
        {
            was_connected = true;
            LOG_INF("OpenThread connected. UDP will start.");
            break;
        }

        k_sleep(K_SECONDS(2));
    }

    if (!was_connected)
    {
        LOG_WRN("Not attached yet; starting in pause mode and waiting for re-attach.");
    }

    sock = zsock_socket(AF_INET6, SOCK_DGRAM, IPPROTO_UDP);
    if (sock < 0)
    {
        LOG_ERR("Failed to create socket: %d", sock);
        return -1;
    }

    memset(&dest_addr, 0, sizeof(dest_addr));
    dest_addr.sin6_family = AF_INET6;
    dest_addr.sin6_port = htons(TARGET_PORT);
    rc = zsock_inet_pton(AF_INET6, TARGET_IP, &dest_addr.sin6_addr);
    if (rc != 1)
    {
        LOG_ERR("Invalid destination IP: %s", TARGET_IP);
        zsock_close(sock);
        return -EINVAL;
    }

    struct sockaddr_in6 local_addr;
    memset(&local_addr, 0, sizeof(local_addr));
    local_addr.sin6_family = AF_INET6;
    local_addr.sin6_port = htons(LISTEN_PORT);
    local_addr.sin6_addr = in6addr_any;

    rc = zsock_bind(sock, (struct sockaddr *)&local_addr, sizeof(local_addr));
    if (rc < 0)
    {
        LOG_ERR("Failed to bind socket: %d", errno);
        zsock_close(sock);
        return -1;
    }

    LOG_INF("Socket successfully bound to local port %d.", LISTEN_PORT);
    LOG_INF("Target: [%s]:%d", TARGET_IP, TARGET_PORT);

    hwinfo_get_device_id(header.device_id, sizeof(header.device_id));
    LOG_INF("Device ID: %02x%02x%02x%02x%02x%02x%02x%02x",
            header.device_id[0], header.device_id[1], header.device_id[2], header.device_id[3],
            header.device_id[4], header.device_id[5], header.device_id[6], header.device_id[7]);

#ifdef FAST_TIMES
    last_temp_send_time = -S_TO_MS(measure_interval_m);
    last_battery_send_time = -S_TO_MS(battery_measure_interval_h);
#else
    last_temp_send_time = -M_TO_MS(measure_interval_m);
    last_battery_send_time = -H_TO_MS(battery_measure_interval_h);
#endif

    send_config();

    while (1)
    {
        bool connected = thread_is_attached();
        uint64_t current_time = k_uptime_get();

        if (connected != was_connected)
        {
            if (connected)
            {
                LOG_INF("OpenThread reconnected, continuing to send.");
                send_config();
            }
            else
            {
                LOG_WRN("OpenThread disconnected, sending paused.");
            }

            was_connected = connected;
        }

        if (!connected)
        {
            k_sleep(K_SECONDS(2));
            continue;
        }

        rc = sensor_power_init();
        if (rc < 0)
        {
            k_sleep(K_SECONDS(2));
            continue;
        }

        rc = sensor_sample_fetch(temp_dev);
        if (rc < 0)
        {
            LOG_ERR("Error fetching temperature data (%d)", rc);
            sensor_power_off();
        }
        else
        {
            rc = sensor_channel_get(temp_dev, SENSOR_CHAN_AMBIENT_TEMP, &temp_val);
            if (rc < 0)
            {
                LOG_ERR("Error reading temperature channel (%d)", rc);
                sensor_power_off();
                k_msleep(2000);
                continue;
            }

            temp_mcelsius = temp_val.val1 * 1000 + temp_val.val2 / 1000;
            temperature_mc = (uint32_t)temp_mcelsius;

            if (filtered_temperature_mc == 0xFFFF)
            {
                filtered_temperature_mc = temperature_mc;
            }
            else
            {
                filtered_temperature_mc = (temperature_mc + (filtered_temperature_mc * 3)) / 4;
            }

            LOG_DBG("Temperature: %u.%03u C (filtered: %u.%03u C)",
                    temperature_mc / 1000, temperature_mc % 1000,
                    filtered_temperature_mc / 1000, filtered_temperature_mc % 1000);

#ifdef FAST_TIMES
            if (abs((int32_t)filtered_temperature_mc - last_temp_send_mc) >= send_delta_mc ||
                current_time - last_temp_send_time >= S_TO_MS(max_send_interval_m))
#else
            if (abs((int32_t)filtered_temperature_mc - last_temp_send_mc) >= send_delta_mc ||
                current_time - last_temp_send_time >= M_TO_MS(max_send_interval_m))
#endif
            {
                sys_put_be32(filtered_temperature_mc, &send_buf[0]);
                rc = send_frame(PAYLOAD_TYPE_TEMP, send_buf, sizeof(filtered_temperature_mc));

                if (rc >= 0)
                {
                    recv_udp();
                    last_temp_send_mc = (int32_t)filtered_temperature_mc;
                    last_temp_send_time = current_time;
                }
            }

            sensor_power_off();
        }

#ifdef FAST_TIMES
        if (current_time - last_battery_send_time >= S_TO_MS(battery_measure_interval_h))
#else
        if (current_time - last_battery_send_time >= H_TO_MS(battery_measure_interval_h))
#endif
        {
            last_battery_send_time = current_time;

            rc = adc_read(adc_channel.dev, &sequence);
            if (rc < 0)
            {
                LOG_ERR("Error reading ADC channel (%d)", rc);
                break;
            }

            battery_mv = adc_buffer;
            rc = adc_raw_to_millivolts_dt(&adc_channel, &battery_mv);
            if (rc < 0)
            {
                LOG_ERR("Error in mV conversion (%d)", rc);
            }
            else
            {
                LOG_DBG("VDD voltage: %d mV", battery_mv);
            }

            uint16_t battery_mv_u16 = (uint16_t)battery_mv;
            sys_put_be16(battery_mv_u16, &send_buf[0]);

            rc = send_frame(PAYLOAD_TYPE_BATTERY, send_buf, sizeof(battery_mv_u16));
            if (rc >= 0)
            {
                last_battery_send_time = current_time;
            }
        }

#ifdef FAST_TIMES
        k_sleep(K_SECONDS(measure_interval_m));
#else
        k_sleep(K_MINUTES(measure_interval_m));
#endif
    }

    return 0;
}