#include "zigbee.h"
#include "status_led.h"
#include "light_driver.h"

#include <freertos/FreeRTOS.h>
#include <freertos/task.h>
#include <esp_zigbee.h>
#include <ezbee/zha.h>
#include <esp_log.h>
#include <esp_err.h>
#include <esp_timer.h>
#include <nvs_flash.h>
#include <esp_check.h>

static const char *TAG = "zigbee";

#define ZB_STORAGE_PARTITION "zb_storage" /* must be subtype nvs in partitions.csv, not fat */
#define ZB_ROUTER_MAX_CHILDREN 10         /* v1 ZR default */
#define ZB_COMMISSIONING_RETRY_MS 1000

static char manufacturer[] = "\x05""Drath";
static char model[] = "\x15""Dimmable Zigbee Light";
static constexpr uint8_t endpoint_no = 1;

static esp_timer_handle_t retry_timer = nullptr;
static volatile ezb_bdb_comm_mode_mask_t retry_mode = 0;

static esp_err_t deferred_driver_init() {
    static bool is_inited = false;
    if (!is_inited) {
        xTaskCreate(light_driver_task, "light task", 4096, nullptr, 6, nullptr);
        is_inited = true;
    }
    return is_inited ? ESP_OK : ESP_FAIL;
}

/* caller must hold the zigbee lock (esp_zigbee_lock_acquire) unless it's running inside a stack callback */
void zigbee_report_on_off_level(const bool on, const uint8_t level) {
    bool on_val = on;
    uint8_t level_val = level;

    ezb_zcl_status_t st = ezb_zcl_set_attr_value(endpoint_no, EZB_ZCL_CLUSTER_ID_ON_OFF, EZB_ZCL_CLUSTER_SERVER,
                                                 EZB_ZCL_ATTR_ON_OFF_ON_OFF_ID, EZB_ZCL_STD_MANUF_CODE, &on_val, false);
    if (st != EZB_ZCL_STATUS_SUCCESS) { ESP_LOGW(TAG, "Failed to set on/off attribute (status 0x%02x)", st); }

    st = ezb_zcl_set_attr_value(endpoint_no, EZB_ZCL_CLUSTER_ID_LEVEL, EZB_ZCL_CLUSTER_SERVER,
                                EZB_ZCL_ATTR_LEVEL_CURRENT_LEVEL_ID, EZB_ZCL_STD_MANUF_CODE, &level_val, false);
    if (st != EZB_ZCL_STATUS_SUCCESS) { ESP_LOGW(TAG, "Failed to set level attribute (status 0x%02x)", st); }
}

static void commissioning_task_cb(void *ctx) {
    const ezb_bdb_comm_mode_mask_t mode = (ezb_bdb_comm_mode_mask_t) (uintptr_t) ctx;
    const ezb_err_t err = ezb_bdb_start_top_level_commissioning(mode);
    ESP_RETURN_ON_FALSE(err == EZB_ERR_NONE, , TAG, "Failed to start Zigbee commissioning (0x%x)", (unsigned)err);
}

static void commissioning_timer_cb(void *arg) {
    if (esp_zigbee_task_queue_post(commissioning_task_cb, (void *) (uintptr_t) retry_mode) != ESP_OK)
        ESP_LOGW(TAG, "Failed to post commissioning retry to zigbee task");
}

static void zigbee_schedule_commissioning(const ezb_bdb_comm_mode_mask_t mode, const uint32_t delay_ms) {
    if (!retry_timer) {
        const esp_timer_create_args_t args = {
            .callback = commissioning_timer_cb,
            .name = "zb_commissioning",
        };
        ESP_RETURN_ON_FALSE(esp_timer_create(&args, &retry_timer) == ESP_OK, , TAG, "Failed to create retry timer");
    }
    esp_timer_stop(retry_timer); /* errors if not running, fine */
    retry_mode = mode;
    ESP_RETURN_ON_FALSE(esp_timer_start_once(retry_timer, (uint64_t)delay_ms * 1000ULL) == ESP_OK, , TAG,
                        "Failed to start retry timer");
}

static ezb_err_t zigbee_utility_add_ep_basic_manufacturer_info(ezb_af_ep_desc_t ep_desc,
                                                               const zcl_basic_manufacturer_info_t *info) {
    ESP_RETURN_ON_FALSE(info && info->manufacturer_name, EZB_ERR_INV_ARG, TAG, "Invalid manufacturer name");
    ESP_RETURN_ON_FALSE(info->model_identifier, EZB_ERR_INV_ARG, TAG, "Invalid model identifier");

    ezb_zcl_cluster_desc_t basic_cluster = ezb_af_endpoint_get_cluster_desc(
        ep_desc, EZB_ZCL_CLUSTER_ID_BASIC, EZB_ZCL_CLUSTER_SERVER);
    ESP_RETURN_ON_FALSE(basic_cluster != EZB_INVALID_ZCL_CLUSTER_DESC, EZB_ERR_NOT_FOUND, TAG,
                        "Failed to find basic cluster in endpoint: %d", endpoint_no);

    ezb_err_t err = ezb_zcl_basic_cluster_desc_add_attr(basic_cluster, EZB_ZCL_ATTR_BASIC_MANUFACTURER_NAME_ID,
                                                        info->manufacturer_name);
    ESP_RETURN_ON_FALSE(err == EZB_ERR_NONE, err, TAG, "Failed to add manufacturer name (0x%x)", (unsigned)err);

    err = ezb_zcl_basic_cluster_desc_add_attr(basic_cluster, EZB_ZCL_ATTR_BASIC_MODEL_IDENTIFIER_ID,
                                              info->model_identifier);
    ESP_RETURN_ON_FALSE(err == EZB_ERR_NONE, err, TAG, "Failed to add model identifier (0x%x)", (unsigned)err);
    return EZB_ERR_NONE;
}

static void zigbee_attribute_handler(ezb_zcl_set_attr_value_message_t *message) {
    ESP_RETURN_ON_FALSE(message, , TAG, "Empty message");

    /* v2: no return value, the verdict goes back to the stack through out.result */
    message->out.result = EZB_ZCL_STATUS_SUCCESS;

    if (message->info.status != EZB_ZCL_STATUS_SUCCESS) {
        ESP_LOGW(TAG, "Received message: error status(%d)", message->info.status);
        message->out.result = EZB_ZCL_STATUS_FAIL;
        return;
    }

    if (message->info.dst_ep == endpoint_no) {
        const ezb_zcl_attribute_t *attr = &message->in.attribute;

        switch (message->info.cluster_id) {
            case EZB_ZCL_CLUSTER_ID_IDENTIFY:
                if (attr->id == EZB_ZCL_ATTR_IDENTIFY_IDENTIFY_TIME_ID) {
                    ESP_LOGI(TAG, "Identify requested");
                    light_driver_identify();
                }
                break;

            case EZB_ZCL_CLUSTER_ID_ON_OFF:
                if (attr->id == EZB_ZCL_ATTR_ON_OFF_ON_OFF_ID && attr->data.type == EZB_ZCL_ATTR_TYPE_BOOL) {
                    bool light_state = false;
                    light_state = attr->data.value ? *(bool *) attr->data.value : light_state;
                    ESP_LOGI(TAG, "Light sets to %s", light_state ? "On" : "Off");
                    light_driver_set_power(light_state);
                }
                break;

            case EZB_ZCL_CLUSTER_ID_LEVEL:
                if (attr->id == EZB_ZCL_ATTR_LEVEL_CURRENT_LEVEL_ID && attr->data.type == EZB_ZCL_ATTR_TYPE_UINT8) {
                    uint8_t light_level = 0;
                    light_level = attr->data.value ? *(uint8_t *) attr->data.value : light_level;
                    light_driver_set_level(light_level);
                    ESP_LOGI(TAG, "Light level changes to %d", light_level);
                }
                break;

            default:
                ESP_LOGW(TAG, "Unhandled zigbee message, data: cluster(0x%x), attribute(0x%x)  ",
                         message->info.cluster_id, attr->id);
        }
    }
}

static void zigbee_action_handler(ezb_zcl_core_action_callback_id_t callback_id, void *message) {
    if (callback_id == EZB_ZCL_CORE_SET_ATTR_VALUE_CB_ID) {
        zigbee_attribute_handler(message);
        return;
    }

    ESP_LOGW(TAG, "Receive Zigbee action(0x%lx) callback", (unsigned long)callback_id);
}

/*
 * ZHA helper builds the whole endpoint now: basic, identify, groups, scenes, on/off, level
 * (device id 0x0101 / HA profile are set by it). ota is still left out. maybe one day
 */
static ezb_af_ep_desc_t zigbee_endpoint_create() {
    ezb_zha_dimmable_light_config_t light_cfg = EZB_ZHA_DIMMABLE_LIGHT_CONFIG();
    light_cfg.basic_cfg.power_source = EZB_ZCL_BASIC_POWER_SOURCE_DC_SOURCE;

    ezb_af_ep_desc_t ep_desc = ezb_zha_create_dimmable_light(endpoint_no, &light_cfg);
    ESP_RETURN_ON_FALSE(ep_desc != EZB_INVALID_AF_EP_DESC, EZB_INVALID_AF_EP_DESC, TAG,
                        "Failed to create dimmable light endpoint");

    const zcl_basic_manufacturer_info_t mfg_info = {
        .manufacturer_name = manufacturer,
        .model_identifier = model,
    };
    if (zigbee_utility_add_ep_basic_manufacturer_info(ep_desc, &mfg_info) != EZB_ERR_NONE) {
        ezb_af_free_endpoint_desc(ep_desc);
        return EZB_INVALID_AF_EP_DESC;
    }
    return ep_desc;
}

/* v2: not a magic symbol anymore, registered with ezb_app_signal_add_handler(). return true = handled */
static bool zigbee_app_signal_handler(const ezb_app_signal_t *signal_s) {
    const ezb_app_signal_type_t sig_type = ezb_app_signal_get_type(signal_s);
    const void *params = ezb_app_signal_get_params(signal_s);

    switch (sig_type) {
        case EZB_ZDO_SIGNAL_SKIP_STARTUP:
            /* v2 doesn't run initialization (rejoin) by itself anymore, kick it here */
            ESP_LOGI(TAG, "Initialize Zigbee stack");
            ezb_bdb_start_top_level_commissioning(EZB_BDB_MODE_INITIALIZATION);
            return true;

        case EZB_BDB_SIGNAL_DEVICE_FIRST_START:
        case EZB_BDB_SIGNAL_DEVICE_REBOOT: {
            const uint8_t status = params
                                       ? ((const ezb_bdb_signal_simple_params_t *) params)->status
                                       : (uint8_t) EZB_BDB_STATUS_SUCCESS;
            if (status == EZB_BDB_STATUS_SUCCESS) {
                ESP_LOGI(TAG, "Deferred driver initialization %s", deferred_driver_init() ? "failed" : "successful");
                ESP_LOGI(TAG, "Device started up in%s factory-reset mode", ezb_bdb_is_factory_new() ? "" : " non");
                if (ezb_bdb_is_factory_new()) {
                    ESP_LOGI(TAG, "Start network steering");
                    ezb_bdb_start_top_level_commissioning(EZB_BDB_MODE_NETWORK_STEERING);
                } else {
                    ESP_LOGI(TAG, "Device rebooted, steering anyway");
                    ezb_bdb_start_top_level_commissioning(EZB_BDB_MODE_NETWORK_STEERING);
                }
            } else {
                ESP_LOGW(TAG, "%s failed with status: 0x%02x, retrying", ezb_app_signal_to_string(sig_type), status);
                zigbee_schedule_commissioning(EZB_BDB_MODE_INITIALIZATION, ZB_COMMISSIONING_RETRY_MS);
                set_status_error_or_resetting();
            }
            return true;
        }

        case EZB_BDB_SIGNAL_STEERING: {
            const uint8_t status = params
                                       ? ((const ezb_bdb_signal_simple_params_t *) params)->status
                                       : (uint8_t) EZB_BDB_STATUS_SUCCESS;
            if (status == EZB_BDB_STATUS_SUCCESS) {
                ezb_extpanid_t extended_pan_id;
                ezb_get_extended_panid(&extended_pan_id);
                ESP_LOGI(
                    TAG,
                    "Joined network successfully (Extended PAN ID: %02x:%02x:%02x:%02x:%02x:%02x:%02x:%02x, PAN ID: 0x%04hx, Channel:%d, Short Address: 0x%04hx)",
                    extended_pan_id.u8[7], extended_pan_id.u8[6], extended_pan_id.u8[5], extended_pan_id.u8[4],
                    extended_pan_id.u8[3], extended_pan_id.u8[2], extended_pan_id.u8[1], extended_pan_id.u8[0],
                    ezb_get_panid(), ezb_get_current_channel(), ezb_get_short_address());
                zigbee_report_on_off_level(false, 254); /* we're in a stack callback, lock is already held */
                set_status_joined();
            } else {
                ESP_LOGI(TAG, "Network steering was not successful (status: 0x%02x)", status);
                zigbee_schedule_commissioning(EZB_BDB_MODE_NETWORK_STEERING, ZB_COMMISSIONING_RETRY_MS);
                set_status_error_or_resetting();
            }
            return true;
        }

        case EZB_NWK_SIGNAL_PERMIT_JOIN_STATUS: {
            /* no err_status in v2, the signal just fires on open/close */
            const uint8_t duration = params
                                         ? ((const ezb_nwk_signal_permit_join_status_params_t *) params)->duration
                                         : 0;
            if (duration) {
                ESP_LOGI(TAG, "Network(0x%04hx) is open for %d seconds", ezb_get_panid(), duration);
            } else { ESP_LOGW(TAG, "Network(0x%04hx) closed, devices joining not allowed.", ezb_get_panid()); }
            return true;
        }

        default:
            ESP_LOGI(TAG, "ZDO signal: %s (0x%x)", ezb_app_signal_to_string(sig_type), sig_type);
            return false;
    }
}

void zigbee_task(void *pvParameters) {
    /* v2 keeps its persistent data in nvs, not fat. old v1 data in this partition is not migrated */
    esp_err_t nvs_err = nvs_flash_init_partition(ZB_STORAGE_PARTITION);
    if (nvs_err == ESP_ERR_NVS_NO_FREE_PAGES || nvs_err == ESP_ERR_NVS_NEW_VERSION_FOUND) {
        ESP_LOGW(TAG, "%s is not valid nvs (old v1 fat data?), erasing, device will need to re-pair",
                 ZB_STORAGE_PARTITION);
        ESP_ERROR_CHECK(nvs_flash_erase_partition(ZB_STORAGE_PARTITION));
        nvs_err = nvs_flash_init_partition(ZB_STORAGE_PARTITION);
    }
    ESP_ERROR_CHECK(nvs_err);

    const esp_zigbee_config_t zb_config = {
        .device_config = {
            .device_type = EZB_NWK_DEVICE_TYPE_ROUTER,
            .install_code_policy = false,
            .zczr_config = {
                .max_children = ZB_ROUTER_MAX_CHILDREN,
            },
        },
        .platform_config = {
            .storage_partition_name = ZB_STORAGE_PARTITION,
            .radio_config = {
                .radio_mode = ESP_ZIGBEE_RADIO_MODE_NATIVE,
            },
        },
    };

    ESP_ERROR_CHECK(esp_zigbee_init(&zb_config));
    set_status_connecting();

    ESP_ERROR_CHECK(ezb_app_signal_add_handler(zigbee_app_signal_handler) == EZB_ERR_NONE ? ESP_OK : ESP_FAIL);

    ezb_af_device_desc_t device_desc = ezb_af_create_device_desc();
    ESP_ERROR_CHECK(device_desc != EZB_INVALID_AF_DEVICE_DESC ? ESP_OK : ESP_FAIL);

    ezb_af_ep_desc_t ep_desc = zigbee_endpoint_create();
    ESP_ERROR_CHECK(ep_desc != EZB_INVALID_AF_EP_DESC ? ESP_OK : ESP_FAIL);

    ESP_ERROR_CHECK(ezb_af_device_add_endpoint_desc(device_desc, ep_desc) == EZB_ERR_NONE ? ESP_OK : ESP_FAIL);
    ESP_ERROR_CHECK(ezb_af_device_desc_register(device_desc) == EZB_ERR_NONE ? ESP_OK : ESP_FAIL);

    ezb_zcl_core_action_handler_register(zigbee_action_handler);

    /* channel mask: v2 scans all channels by default, ESP_ZB_TRANSCEIVER_ALL_CHANNELS_MASK is gone */

    ESP_ERROR_CHECK(esp_zigbee_start(false));

    esp_zigbee_launch_mainloop(); // blocking, holds the task here until esp_zigbee_stop()

    esp_zigbee_deinit();
    vTaskDelete(nullptr);
}