#include "printer_protocol.h"
#include "printer_protocol_internal.h"

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

#include "ble_printer_client.h"
#include "esp_heap_caps.h"
#include "esp_log.h"
#include "esp_timer.h"
#include "runtime_diagnostics.h"
#include "runtime_policy.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"

SemaphoreHandle_t s_mutex;
QueueHandle_t s_job_queue;
EventGroupHandle_t s_evt;

uint32_t s_busy_refcnt;
uint32_t s_status_poll_pause_next_token;
uint32_t s_status_poll_pause_tokens[STATUS_POLL_PAUSE_SLOT_MAX];

parsed_status_t s_status = {
    .has_paper = true,
    .battery = 100,
    .temperature = 25.0f,
    .updated_ms = 0,
};

uint8_t s_last_rsp_cmd;
uint16_t s_last_rsp_payload_len;
uint8_t s_last_rsp_payload[252];

uint32_t s_next_job_id = 1;
job_slot_t s_jobs[JOB_SLOT_MAX];
TaskHandle_t s_worker_task;
TaskHandle_t s_status_poll_task_handle;
bool s_protocol_initialized;
bool s_protocol_stopping;

static const char *TAG = "printer_protocol";

#if CONFIG_FREERTOS_UNICORE
#define PRINT_WORKER_CORE_ID 0
#else
#define PRINT_WORKER_CORE_ID 1
#endif

static size_t status_poll_pause_depth_locked(void) {
    size_t depth = 0;
    for (size_t i = 0; i < STATUS_POLL_PAUSE_SLOT_MAX; ++i) {
        if (s_status_poll_pause_tokens[i] != 0) {
            ++depth;
        }
    }
    return depth;
}

static void status_poll_pause_slots_clear_locked(void) {
    memset(s_status_poll_pause_tokens, 0, sizeof(s_status_poll_pause_tokens));
}

static void reset_jobs_locked(void) {
    for (int i = 0; i < JOB_SLOT_MAX; ++i) {
        free(s_jobs[i].data);
        memset(&s_jobs[i], 0, sizeof(s_jobs[i]));
    }
    s_next_job_id = 1;
}

static void reset_runtime_state_locked(void) {
    s_busy_refcnt = 0;
    s_last_rsp_cmd = 0;
    s_last_rsp_payload_len = 0;
    memset(s_last_rsp_payload, 0, sizeof(s_last_rsp_payload));

    status_poll_pause_slots_clear_locked();
    s_status_poll_pause_next_token = 1;

    s_status.has_paper = true;
    s_status.battery = 100;
    s_status.temperature = 25.0f;
    s_status.updated_ms = 0;
}

bool printer_protocol_is_stopping(void) {
    if (!s_protocol_initialized) {
        return false;
    }
    if (s_mutex == NULL) {
        return true;
    }

    bool stopping = true;
    if (xSemaphoreTake(s_mutex, pdMS_TO_TICKS(50)) == pdTRUE) {
        stopping = s_protocol_stopping;
        xSemaphoreGive(s_mutex);
    }
    return stopping;
}

static BaseType_t printer_create_task_prefer_psram(TaskFunction_t task_fn,
                                                    const char *name,
                                                    uint32_t stack_size,
                                                    UBaseType_t priority,
                                                    TaskHandle_t *out_task) {
#if defined(CONFIG_SPIRAM_ALLOW_STACK_EXTERNAL_MEMORY) && \
    defined(CONFIG_FREERTOS_TASK_CREATE_ALLOW_EXT_MEM) && \
    (configSUPPORT_STATIC_ALLOCATION == 1)
    BaseType_t rc = xTaskCreatePinnedToCoreWithCaps(task_fn,
                                                     name,
                                                     (configSTACK_DEPTH_TYPE)stack_size,
                                                     NULL,
                                                     priority,
                                                     out_task,
                                                     PRINT_WORKER_CORE_ID,
                                                     MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT);
    if (rc == pdPASS) {
        return pdPASS;
    }
#endif
    return xTaskCreatePinnedToCore(task_fn,
                                   name,
                                   (configSTACK_DEPTH_TYPE)stack_size,
                                   NULL,
                                   priority,
                                   out_task,
                                   PRINT_WORKER_CORE_ID);
}

printer_status_poll_pause_token_t printer_protocol_status_poll_pause_acquire(void) {
    if (!s_protocol_initialized || s_mutex == NULL) {
        return 0;
    }

    if (xSemaphoreTake(s_mutex, pdMS_TO_TICKS(200)) != pdTRUE) {
        return 0;
    }

    if (s_protocol_stopping) {
        xSemaphoreGive(s_mutex);
        return 0;
    }

    int slot = -1;
    for (int i = 0; i < STATUS_POLL_PAUSE_SLOT_MAX; ++i) {
        if (s_status_poll_pause_tokens[i] == 0) {
            slot = i;
            break;
        }
    }

    if (slot < 0) {
        xSemaphoreGive(s_mutex);
        ESP_LOGW(TAG, "status poll pause acquire failed: no free slot");
        return 0;
    }

    uint32_t token = s_status_poll_pause_next_token++;
    if (token == 0) {
        token = s_status_poll_pause_next_token++;
    }

    s_status_poll_pause_tokens[slot] = token;
    runtime_diag_set_gauge(RUNTIME_DIAG_GAUGE_STATUS_POLL_PAUSE_DEPTH,
                           (int32_t)status_poll_pause_depth_locked());
    xSemaphoreGive(s_mutex);
    return token;
}

void printer_protocol_status_poll_pause_release(printer_status_poll_pause_token_t *token) {
    if (token == NULL || *token == 0) {
        return;
    }

    if (!s_protocol_initialized || s_mutex == NULL) {
        *token = 0;
        return;
    }

    if (xSemaphoreTake(s_mutex, pdMS_TO_TICKS(200)) != pdTRUE) {
        return;
    }

    for (int i = 0; i < STATUS_POLL_PAUSE_SLOT_MAX; ++i) {
        if (s_status_poll_pause_tokens[i] == *token) {
            s_status_poll_pause_tokens[i] = 0;
            break;
        }
    }

    runtime_diag_set_gauge(RUNTIME_DIAG_GAUGE_STATUS_POLL_PAUSE_DEPTH,
                           (int32_t)status_poll_pause_depth_locked());
    xSemaphoreGive(s_mutex);
    *token = 0;
}

bool printer_protocol_acquire_control_lane(uint32_t timeout_ms, char *err, size_t err_len) {
    if (!s_protocol_initialized || s_mutex == NULL) {
        if (err != NULL && err_len > 0) {
            snprintf(err, err_len, "printer protocol not initialized");
        }
        return false;
    }

    if (xSemaphoreTake(s_mutex, pdMS_TO_TICKS(timeout_ms)) != pdTRUE) {
        if (err != NULL && err_len > 0) {
            snprintf(err, err_len, "lock timeout");
        }
        return false;
    }

    if (s_protocol_stopping) {
        xSemaphoreGive(s_mutex);
        if (err != NULL && err_len > 0) {
            snprintf(err, err_len, "printer protocol stopping");
        }
        return false;
    }

    if (s_busy_refcnt != 0) {
        xSemaphoreGive(s_mutex);
        if (err != NULL && err_len > 0) {
            snprintf(err, err_len, "printer busy");
        }
        return false;
    }

    s_busy_refcnt = 1;
    xSemaphoreGive(s_mutex);
    return true;
}

void printer_protocol_release_control_lane(void) {
    if (s_mutex == NULL) {
        return;
    }

    if (xSemaphoreTake(s_mutex, pdMS_TO_TICKS(1000)) != pdTRUE) {
        return;
    }

    if (s_busy_refcnt > 0) {
        --s_busy_refcnt;
    }
    xSemaphoreGive(s_mutex);
}

static uint8_t checksum8(const uint8_t *data, size_t len) {
    uint32_t sum = 0;
    for (size_t i = 0; i < len; ++i) {
        sum += data[i];
    }
    return (uint8_t)(sum & 0xFF);
}

esp_err_t printer_protocol_send_frame(uint8_t cmd,
                                      const uint8_t *payload,
                                      uint16_t payload_len,
                                      bool with_checksum) {
    if (!s_protocol_initialized || printer_protocol_is_stopping()) {
        return ESP_ERR_INVALID_STATE;
    }

    uint8_t frame[260];
    size_t total = 4 + payload_len + (with_checksum ? 1 : 0);
    if (total > sizeof(frame)) {
        return ESP_ERR_INVALID_SIZE;
    }

    frame[0] = PROTO_ADDR;
    frame[1] = cmd;
    frame[2] = (uint8_t)((payload_len >> 8) & 0xFF);
    frame[3] = (uint8_t)(payload_len & 0xFF);

    if (payload_len > 0 && payload != NULL) {
        memcpy(&frame[4], payload, payload_len);
    }

    if (with_checksum) {
        frame[4 + payload_len] = checksum8(frame, 4 + payload_len);
    }

    return ble_printer_client_write(frame, total);
}

static void clear_ack_signal(void) {
    if (s_evt != NULL) {
        xEventGroupClearBits(s_evt, EVT_ACK);
    }
}

bool printer_protocol_wait_response(uint8_t cmd,
                                    uint8_t *out_payload,
                                    size_t out_payload_cap,
                                    uint16_t *out_payload_len,
                                    uint32_t timeout_ms) {
    if (timeout_ms == 0 || s_evt == NULL) {
        return false;
    }

    int64_t deadline = esp_timer_get_time() / 1000 + timeout_ms;

    while (true) {
        if (printer_protocol_is_stopping()) {
            return false;
        }

        int64_t now = esp_timer_get_time() / 1000;
        if (now >= deadline) {
            return false;
        }

        uint32_t wait_ms = (uint32_t)(deadline - now);
        if (wait_ms > 100) {
            wait_ms = 100;
        }

        EventBits_t bits = xEventGroupWaitBits(s_evt,
                                               EVT_ACK | EVT_SHUTDOWN,
                                               pdFALSE,
                                               pdFALSE,
                                               pdMS_TO_TICKS(wait_ms));
        if (bits & EVT_SHUTDOWN) {
            return false;
        }

        if (!(bits & EVT_ACK)) {
            continue;
        }
        xEventGroupClearBits(s_evt, EVT_ACK);

        if (xSemaphoreTake(s_mutex, pdMS_TO_TICKS(200)) != pdTRUE) {
            continue;
        }

        uint8_t rsp_cmd = s_last_rsp_cmd;
        uint16_t rsp_len = s_last_rsp_payload_len;
        uint8_t rsp_copy[sizeof(s_last_rsp_payload)];
        if (rsp_len > sizeof(rsp_copy)) {
            rsp_len = sizeof(rsp_copy);
        }
        if (rsp_len > 0) {
            memcpy(rsp_copy, s_last_rsp_payload, rsp_len);
        }
        xSemaphoreGive(s_mutex);

        if (rsp_cmd == cmd) {
            if (out_payload != NULL && out_payload_cap > 0 && rsp_len > 0) {
                size_t copy_len = rsp_len;
                if (copy_len > out_payload_cap) {
                    copy_len = out_payload_cap;
                }
                memcpy(out_payload, rsp_copy, copy_len);
            }
            if (out_payload_len != NULL) {
                *out_payload_len = rsp_len;
            }
            return true;
        }
    }
}

static bool wait_ack(uint8_t cmd, uint32_t timeout_ms) {
    uint8_t payload[252];
    uint16_t payload_len = 0;
    if (!printer_protocol_wait_response(cmd, payload, sizeof(payload), &payload_len, timeout_ms)) {
        return false;
    }
    return payload_len >= 1 && payload[0] == 0x01;
}

bool printer_protocol_send_cmd_with_ack(uint8_t cmd,
                                        const uint8_t *payload,
                                        uint16_t payload_len,
                                        bool with_checksum,
                                        uint32_t timeout_ms) {
    clear_ack_signal();
    if (printer_protocol_send_frame(cmd, payload, payload_len, with_checksum) != ESP_OK) {
        return false;
    }
    return wait_ack(cmd, timeout_ms);
}

esp_err_t printer_protocol_send_cmd_wait_response(uint8_t cmd,
                                                  const uint8_t *payload,
                                                  uint16_t payload_len,
                                                  bool with_checksum,
                                                  uint32_t timeout_ms,
                                                  bool expect_ack,
                                                  uint8_t *out_payload,
                                                  size_t out_payload_cap,
                                                  uint16_t *out_payload_len,
                                                  char *err,
                                                  size_t err_len) {
    clear_ack_signal();
    if (printer_protocol_send_frame(cmd, payload, payload_len, with_checksum) != ESP_OK) {
        if (err != NULL && err_len > 0) {
            snprintf(err, err_len, "send frame failed");
        }
        return ESP_FAIL;
    }

    uint8_t rsp[252];
    uint16_t rsp_len = 0;
    if (!printer_protocol_wait_response(cmd, rsp, sizeof(rsp), &rsp_len, timeout_ms)) {
        if (err != NULL && err_len > 0) {
            if (printer_protocol_is_stopping()) {
                snprintf(err, err_len, "protocol stopping");
            } else {
                snprintf(err, err_len, "response timeout");
            }
        }
        return printer_protocol_is_stopping() ? ESP_ERR_INVALID_STATE : ESP_ERR_TIMEOUT;
    }

    if (expect_ack && (rsp_len < 1 || rsp[0] != 0x01)) {
        if (err != NULL && err_len > 0) {
            snprintf(err, err_len, "cmd 0x%02X rejected", cmd);
        }
        return ESP_FAIL;
    }

    if (out_payload != NULL && out_payload_cap > 0 && rsp_len > 0) {
        size_t copy_len = rsp_len;
        if (copy_len > out_payload_cap) {
            copy_len = out_payload_cap;
        }
        memcpy(out_payload, rsp, copy_len);
    }
    if (out_payload_len != NULL) {
        *out_payload_len = rsp_len;
    }

    return ESP_OK;
}

static void status_poll_task(void *arg) {
    (void)arg;

    while (true) {
        bool should_stop = false;
        bool can_poll = false;

        if (xSemaphoreTake(s_mutex, pdMS_TO_TICKS(50)) == pdTRUE) {
            should_stop = s_protocol_stopping;
            can_poll = (s_busy_refcnt == 0) && (status_poll_pause_depth_locked() == 0);
            xSemaphoreGive(s_mutex);
        }

        if (should_stop) {
            break;
        }

        if (ble_printer_client_is_connected() && can_poll) {
            uint8_t dummy = 0x00;
            (void)printer_protocol_send_frame(CMD_GET_STATUS, &dummy, 0, true);
        }
        vTaskDelay(pdMS_TO_TICKS(runtime_policy_printer_status_poll_interval_ms()));
    }

    if (s_evt != NULL) {
        xEventGroupSetBits(s_evt, EVT_STATUS_POLL_EXITED);
    }
    s_status_poll_task_handle = NULL;
    vTaskDelete(NULL);
}

static void on_rx_frame(const uint8_t *data, size_t len) {
    if (data == NULL || len < 4 || s_mutex == NULL) {
        return;
    }

    if (data[0] != PROTO_ADDR) {
        return;
    }

    uint8_t cmd = data[1];
    uint16_t payload_len = ((uint16_t)data[2] << 8) | data[3];
    size_t required = (size_t)payload_len + 4;
    if (len < required) {
        return;
    }

    const uint8_t *payload = &data[4];

    if (cmd == CMD_GET_STATUS && payload_len >= 5) {
        bool has_paper = payload[0] != 0;
        uint8_t battery = payload[1];
        int sign = (payload[2] == 0x2D) ? -1 : 1;
        int temp_x10 = ((int)payload[3] << 8) | payload[4];
        float temperature = (float)(sign * temp_x10) / 10.0f;

        if (xSemaphoreTake(s_mutex, pdMS_TO_TICKS(100)) == pdTRUE) {
            s_status.has_paper = has_paper;
            s_status.battery = battery;
            s_status.temperature = temperature;
            s_status.updated_ms = esp_timer_get_time() / 1000;
            xSemaphoreGive(s_mutex);
        }
        return;
    }

    if (payload_len >= 1) {
        if (xSemaphoreTake(s_mutex, pdMS_TO_TICKS(100)) == pdTRUE) {
            s_last_rsp_cmd = cmd;
            s_last_rsp_payload_len = payload_len;
            if (s_last_rsp_payload_len > sizeof(s_last_rsp_payload)) {
                s_last_rsp_payload_len = sizeof(s_last_rsp_payload);
            }
            memcpy(s_last_rsp_payload, payload, s_last_rsp_payload_len);
            xSemaphoreGive(s_mutex);
        }
        if (s_evt != NULL) {
            xEventGroupSetBits(s_evt, EVT_ACK);
        }
    }
}

static void destroy_runtime_objects(void) {
    if (s_job_queue != NULL) {
        vQueueDelete(s_job_queue);
        s_job_queue = NULL;
    }
    if (s_evt != NULL) {
        vEventGroupDelete(s_evt);
        s_evt = NULL;
    }
    if (s_mutex != NULL) {
        vSemaphoreDelete(s_mutex);
        s_mutex = NULL;
    }
}

esp_err_t printer_protocol_init(void) {
    if (s_protocol_initialized) {
        return ESP_OK;
    }

    s_mutex = xSemaphoreCreateMutex();
    if (s_mutex == NULL) {
        return ESP_ERR_NO_MEM;
    }

    s_evt = xEventGroupCreate();
    if (s_evt == NULL) {
        destroy_runtime_objects();
        return ESP_ERR_NO_MEM;
    }

    s_job_queue = xQueueCreate(JOB_QUEUE_LEN, sizeof(uint32_t));
    if (s_job_queue == NULL) {
        destroy_runtime_objects();
        return ESP_ERR_NO_MEM;
    }

    if (xSemaphoreTake(s_mutex, pdMS_TO_TICKS(200)) == pdTRUE) {
        reset_jobs_locked();
        reset_runtime_state_locked();
        s_protocol_stopping = false;
        xSemaphoreGive(s_mutex);
    }

    runtime_diag_set_gauge(RUNTIME_DIAG_GAUGE_STATUS_POLL_PAUSE_DEPTH, 0);
    runtime_diag_set_gauge(RUNTIME_DIAG_GAUGE_PRINTER_QUEUE_DEPTH, 0);
    s_protocol_initialized = true;
    s_protocol_stopping = false;

    esp_err_t err = ble_printer_client_init(on_rx_frame);
    if (err != ESP_OK) {
        s_protocol_initialized = false;
        destroy_runtime_objects();
        return err;
    }

    BaseType_t ok = printer_create_task_prefer_psram(printer_protocol_worker_task,
                                                      "print_worker",
                                                      6144,
                                                      6,
                                                      &s_worker_task);
    if (ok != pdPASS) {
        s_protocol_initialized = false;
        destroy_runtime_objects();
        return ESP_ERR_NO_MEM;
    }

    ok = printer_create_task_prefer_psram(status_poll_task,
                                          "status_poll",
                                          4096,
                                          4,
                                          &s_status_poll_task_handle);
    if (ok != pdPASS) {
        if (s_worker_task != NULL) {
            vTaskDelete(s_worker_task);
            s_worker_task = NULL;
        }
        s_protocol_initialized = false;
        destroy_runtime_objects();
        return ESP_ERR_NO_MEM;
    }

    ESP_LOGI(TAG, "printer protocol initialized");
    return ESP_OK;
}

esp_err_t printer_protocol_stop(uint32_t timeout_ms) {
    if (!s_protocol_initialized) {
        return ESP_OK;
    }

    if (timeout_ms == 0) {
        timeout_ms = runtime_policy_printer_stop_timeout_ms();
    }

    if (s_mutex != NULL && xSemaphoreTake(s_mutex, pdMS_TO_TICKS(1000)) == pdTRUE) {
        s_protocol_stopping = true;
        s_busy_refcnt = 0;
        status_poll_pause_slots_clear_locked();
        xSemaphoreGive(s_mutex);
    }

    runtime_diag_set_gauge(RUNTIME_DIAG_GAUGE_STATUS_POLL_PAUSE_DEPTH, 0);

    if (s_evt != NULL) {
        xEventGroupSetBits(s_evt, EVT_SHUTDOWN);
    }

    if (s_job_queue != NULL) {
        const uint32_t sentinel = PRINTER_JOB_SENTINEL_STOP;
        (void)xQueueSend(s_job_queue, &sentinel, 0);
    }

    ble_printer_client_disconnect();

    bool worker_done = (s_worker_task == NULL);
    bool poll_done = (s_status_poll_task_handle == NULL);

    if (s_evt != NULL && (!worker_done || !poll_done)) {
        EventBits_t bits = xEventGroupWaitBits(s_evt,
                                               EVT_WORKER_EXITED | EVT_STATUS_POLL_EXITED,
                                               pdFALSE,
                                               pdTRUE,
                                               pdMS_TO_TICKS(timeout_ms));
        worker_done = worker_done || ((bits & EVT_WORKER_EXITED) != 0) || s_worker_task == NULL;
        poll_done = poll_done || ((bits & EVT_STATUS_POLL_EXITED) != 0) || s_status_poll_task_handle == NULL;
    }

    if (!worker_done && s_worker_task != NULL) {
        ESP_LOGW(TAG, "worker task stop timeout, force delete");
        vTaskDelete(s_worker_task);
        s_worker_task = NULL;
    }
    if (!poll_done && s_status_poll_task_handle != NULL) {
        ESP_LOGW(TAG, "status poll task stop timeout, force delete");
        vTaskDelete(s_status_poll_task_handle);
        s_status_poll_task_handle = NULL;
    }

    if (s_mutex != NULL && xSemaphoreTake(s_mutex, pdMS_TO_TICKS(500)) == pdTRUE) {
        reset_jobs_locked();
        reset_runtime_state_locked();
        xSemaphoreGive(s_mutex);
    }

    destroy_runtime_objects();

    s_busy_refcnt = 0;
    s_status_poll_pause_next_token = 1;
    memset(s_status_poll_pause_tokens, 0, sizeof(s_status_poll_pause_tokens));
    memset(s_last_rsp_payload, 0, sizeof(s_last_rsp_payload));
    s_last_rsp_cmd = 0;
    s_last_rsp_payload_len = 0;
    s_worker_task = NULL;
    s_status_poll_task_handle = NULL;
    s_protocol_initialized = false;
    s_protocol_stopping = false;

    runtime_diag_set_gauge(RUNTIME_DIAG_GAUGE_STATUS_POLL_PAUSE_DEPTH, 0);
    runtime_diag_set_gauge(RUNTIME_DIAG_GAUGE_PRINTER_QUEUE_DEPTH, 0);

    ESP_LOGI(TAG, "printer protocol stopped");
    return (worker_done && poll_done) ? ESP_OK : ESP_ERR_TIMEOUT;
}

esp_err_t printer_protocol_connect(const char *target_name, uint32_t timeout_ms) {
    if (!s_protocol_initialized || s_protocol_stopping) {
        return ESP_ERR_INVALID_STATE;
    }
    return ble_printer_client_connect(target_name, timeout_ms);
}

void printer_protocol_disconnect(void) {
    ble_printer_client_disconnect();
}

void printer_protocol_get_runtime_status(printer_runtime_status_t *out_status) {
    if (out_status == NULL) {
        return;
    }

    memset(out_status, 0, sizeof(*out_status));

    ble_link_state_t link = {0};
    ble_printer_client_get_link_state(&link);

    out_status->connected = link.connected;
    out_status->notify_ready = link.notify_ready;
    out_status->mtu = link.mtu;

    if (s_mutex != NULL && xSemaphoreTake(s_mutex, pdMS_TO_TICKS(200)) == pdTRUE) {
        out_status->busy = s_busy_refcnt > 0;
        out_status->has_paper = s_status.has_paper;
        out_status->battery_percent = s_status.battery;
        out_status->temperature = s_status.temperature;
        out_status->last_status_ms = s_status.updated_ms;
        out_status->queue_depth = (s_job_queue != NULL)
                                      ? (uint32_t)uxQueueMessagesWaiting(s_job_queue)
                                      : 0;
        xSemaphoreGive(s_mutex);
    }

    runtime_diag_set_gauge(RUNTIME_DIAG_GAUGE_PRINTER_QUEUE_DEPTH, (int32_t)out_status->queue_depth);
}