#include <zephyr/kernel.h>
#include <zephyr/logging/log.h>
#include <zephyr/net/openthread.h>
#include <zephyr/random/random.h>
#include <openthread/thread.h>
#include <openthread/dataset.h>
#include <openthread/instance.h>
#include <openthread/udp.h>
#include <openthread/coap.h>
#include <openthread/ip6.h>
#include <string.h>
#include <stdio.h>
#include <lasertag_utils.h>
#include <thread_mgmt.h>
#include <ble_mgmt.h>

LOG_MODULE_REGISTER(thread_mgmt, CONFIG_THREAD_MGMT_LOG_LEVEL);

#if (CONFIG_THREAD_MGMT_LOG_LEVEL >= LOG_LEVEL_DBG)
static void thread_mgmt_state_notify_callback(otChangedFlags aFlags, void *aContext)
{
    ARG_UNUSED(aContext);
    if (aFlags & OT_CHANGED_THREAD_ROLE)
    {
        struct otInstance *instance = openthread_get_default_instance();
        otDeviceRole role = otThreadGetDeviceRole(instance);
        LOG_DBG("Thread Role changed: " BOLD("%s"), otThreadDeviceRoleToString(role));
    }
}

struct openthread_state_changed_callback thread_state_callback = {
    .otCallback = thread_mgmt_state_notify_callback,
    .user_data = NULL,
};

void thread_mgmt_print_dataset(const otOperationalDataset *dataset)
{
    if (dataset == NULL)
    {
        return;
    }

    // Network Name
    char net_name[OT_NETWORK_NAME_MAX_SIZE + 1] = {0};
    memcpy(net_name, dataset->mNetworkName.m8, sizeof(dataset->mNetworkName.m8));

    // Convert to hex strings using bin2hex
    char ext_pan_id_str[17];
    bin2hex(dataset->mExtendedPanId.m8, 8, ext_pan_id_str, sizeof(ext_pan_id_str));

    char net_key_str[33];
    bin2hex(dataset->mNetworkKey.m8, 16, net_key_str, sizeof(net_key_str));
    LOG_DBG("        Timestamp: " BOLD("%llu"), dataset->mActiveTimestamp.mSeconds);
    LOG_DBG("     Network Name: " BOLD("%s"), net_name);
    LOG_DBG("           PAN ID: " BOLD("0x%04X"), dataset->mPanId);
    LOG_DBG("          Channel: " BOLD("%u"), dataset->mChannel);
    LOG_DBG("  Extended PAN ID: " BOLD("%s"), ext_pan_id_str);
    LOG_DBG("      Network Key: " BOLD("%s"), net_key_str);
}
#endif

int thread_mgmt_init(void)
{
    struct otInstance *instance = openthread_get_default_instance();
    otOperationalDataset dataset;

    if (!instance)
        return -ENODEV;

#if (CONFIG_THREAD_MGMT_LOG_LEVEL >= LOG_LEVEL_DBG)
    openthread_state_changed_callback_register(&thread_state_callback);
#endif

    memset(&dataset, 0, sizeof(otOperationalDataset));
    if (otDatasetGetActive(instance, &dataset))
    {
        // No dataset found, proceed to set up a new one
        LOG_INF("No active Thread dataset found, initializing new dataset.");
        dataset.mActiveTimestamp.mSeconds = 1;
        dataset.mComponents.mIsActiveTimestampPresent = true;
        uint8_t network_name[] = "Eriks Lasertag\0";
        strncpy(dataset.mNetworkName.m8, network_name, strlen(network_name));
        dataset.mComponents.mIsNetworkNamePresent = true;
        dataset.mPanId = 0xdead;
        dataset.mComponents.mIsPanIdPresent = true;
        dataset.mChannel = 15;
        dataset.mComponents.mIsChannelPresent = true;
        uint8_t ext_pan_id[8];
        sys_csrand_get(ext_pan_id, sizeof(ext_pan_id));
        memcpy(dataset.mExtendedPanId.m8, ext_pan_id, 8);
        dataset.mComponents.mIsExtendedPanIdPresent = true;
        uint8_t network_key[16];
        sys_csrand_get(network_key, sizeof(network_key));
        memcpy(dataset.mNetworkKey.m8, network_key, 16);
        dataset.mComponents.mIsNetworkKeyPresent = true;
        memset(dataset.mMeshLocalPrefix.m8, 0, 8);
        uint8_t mesh_local_prefix[8] = {0xfd, 0x00, 0x03, 0x08, 0x20, 0x13, 0x00, 0x00};
        memcpy(dataset.mMeshLocalPrefix.m8, mesh_local_prefix, 8);
        dataset.mComponents.mIsMeshLocalPrefixPresent = true;
        otDatasetSetActive(instance, &dataset);
    }

    otIp6SetEnabled(instance, true);
    otThreadSetEnabled(instance, true);

#if (CONFIG_THREAD_MGMT_LOG_LEVEL >= LOG_LEVEL_DBG)
    memset(&dataset, 0, sizeof(otOperationalDataset));
    otDatasetGetActive(instance, &dataset);
    LOG_DBG("Thread stack dataset after Thread stack initialization:");
    thread_mgmt_print_dataset(&dataset);
#endif
    return 0;
}

void thread_mgmt_restart_thread_stack(device_config_payload_t *pending_config, bool force)
{

    struct otInstance *instance = openthread_get_default_instance();
    otOperationalDataset dataset;
    bool changed = false;

    memset(&dataset, 0, sizeof(otOperationalDataset));

    // Get current active dataset. Force restart if unable to get it.
    if (otDatasetGetActive(instance, &dataset) != OT_ERROR_NONE)
    {
        LOG_WRN("Failed to get active dataset, forcing Thread stack restart.");
        force = true;
    }

    // Compare each relevant field and set 'changed' if any differ
    if (force || dataset.mChannel != pending_config->channel)
    {
        LOG_DBG("Thread channel change detected: %u -> %u", dataset.mChannel, pending_config->channel);
        dataset.mChannel = pending_config->channel;
        dataset.mComponents.mIsChannelPresent = true;
        changed = true;
    }
    if (force || dataset.mPanId != pending_config->pan_id)
    {
        LOG_DBG("Thread PAN ID change detected: 0x%04X -> 0x%04X", dataset.mPanId, pending_config->pan_id);
        dataset.mPanId = pending_config->pan_id;
        dataset.mComponents.mIsPanIdPresent = true;
        changed = true;
    }
    if (force || memcmp(dataset.mExtendedPanId.m8, pending_config->ext_pan_id, 8) != 0)
    {
        LOG_DBG("Thread extended PAN ID change detected.");
        memcpy(dataset.mExtendedPanId.m8, pending_config->ext_pan_id, 8);
        dataset.mComponents.mIsExtendedPanIdPresent = true;
        changed = true;
    }

    if (force || memcmp(dataset.mNetworkKey.m8, pending_config->network_key, 16) != 0)
    {
        LOG_DBG("Thread network key change detected.");
        memcpy(dataset.mNetworkKey.m8, pending_config->network_key, 16);
        dataset.mComponents.mIsNetworkKeyPresent = true;
        changed = true;
    }

    if (force || strncmp(dataset.mNetworkName.m8, pending_config->network_name, 16) != 0)
    {
        LOG_DBG("Thread network name change detected: %s -> %s", dataset.mNetworkName.m8, pending_config->network_name);
        strncpy(dataset.mNetworkName.m8, pending_config->network_name, 16);
        dataset.mComponents.mIsNetworkNamePresent = true;
        changed = true;
    }

    if (changed)
    {
        LOG_DBG("Thread stack restart required; dataset changed.");
        
        dataset.mActiveTimestamp.mSeconds++; 
        dataset.mComponents.mIsActiveTimestampPresent = true;

        otThreadSetEnabled(instance, false);
        otIp6SetEnabled(instance, false);

        // Set the network key BEFORE updating the dataset
        // This ensures OpenThread's crypto layer has the correct key before we update the operational dataset
        otNetworkKey stored_network_key;
        if (dataset.mComponents.mIsNetworkKeyPresent) {
            otError key_err = otThreadSetNetworkKey(instance, &dataset.mNetworkKey);
            if (key_err != OT_ERROR_NONE) {
                LOG_ERR("Error setting network key: %d", key_err);
            }
            
            // Read the key back from the thread stack to ensure it's properly stored
            otThreadGetNetworkKey(instance, &stored_network_key);
            memcpy(dataset.mNetworkKey.m8, stored_network_key.m8, 16);
            // Keep the flag set so the key is included in the persistent dataset
        }

        otError err = otDatasetSetActive(instance, &dataset);
        if (err != OT_ERROR_NONE) {
            LOG_ERR("Error writing dataset: %d", err);
        }

        otIp6SetEnabled(instance, true);
        otThreadSetEnabled(instance, true);

#if (CONFIG_THREAD_MGMT_LOG_LEVEL >= LOG_LEVEL_DBG)
        otOperationalDataset new_active_dataset;
        memset(&new_active_dataset, 0, sizeof(otOperationalDataset));
        if (otDatasetGetActive(instance, &new_active_dataset) == OT_ERROR_NONE) {
            LOG_DBG("Thread stack dataset after Thread stack update:");
            thread_mgmt_print_dataset(&new_active_dataset);
        }
#endif
        LOG_DBG("Thread stack restarted successfully.");
    }
    else
    {
        LOG_DBG("Thread stack restart not required; dataset unchanged.");
    }
}

uint16_t thread_mgmt_get_pan_id(void)
{
    struct otInstance *instance = openthread_get_default_instance();
    otOperationalDataset dataset;

    if (otDatasetGetActive(instance, &dataset) == OT_ERROR_NONE)
    {
        return dataset.mPanId;
    }
    return 0;
}

uint8_t thread_mgmt_get_channel(void)
{
    struct otInstance *instance = openthread_get_default_instance();
    otOperationalDataset dataset;

    if (otDatasetGetActive(instance, &dataset) == OT_ERROR_NONE)
    {
        return dataset.mChannel;
    }
    return 0;
}

void thread_mgmt_get_ext_pan_id(uint8_t *dest_8byte)
{
    struct otInstance *instance = openthread_get_default_instance();
    otOperationalDataset dataset;

    if (otDatasetGetActive(instance, &dataset) == OT_ERROR_NONE) {
        memcpy(dest_8byte, dataset.mExtendedPanId.m8, 8);
    }
}

void thread_mgmt_get_network_key(uint8_t *dest_16byte)
{
    struct otInstance *instance = openthread_get_default_instance();
    otOperationalDataset dataset;

    if (otDatasetGetActive(instance, &dataset) == OT_ERROR_NONE) {
        memcpy(dest_16byte, dataset.mNetworkKey.m8, 16);
    }
}

void thread_mgmt_get_network_name(char *dest_str, size_t max_len)
{
    struct otInstance *instance = openthread_get_default_instance();
    otOperationalDataset dataset;

    if (otDatasetGetActive(instance, &dataset) == OT_ERROR_NONE) {
        // Safe copy with forced null-termination
        snprintf(dest_str, max_len, "%s", dataset.mNetworkName.m8);
    }
}
