#include "image_generation.h"

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

#include "cJSON.h"
#include "esp_crt_bundle.h"
#include "esp_heap_caps.h"
#include "esp_http_client.h"
#include "esp_log.h"

#ifndef CONFIG_TQ_Z_IMAGE_API_ENDPOINT
#define CONFIG_TQ_Z_IMAGE_API_ENDPOINT "https://dashscope.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation"
#endif

#ifndef CONFIG_TQ_Z_IMAGE_MODEL
#define CONFIG_TQ_Z_IMAGE_MODEL "z-image-turbo"
#endif

#ifndef CONFIG_TQ_Z_IMAGE_DEFAULT_SIZE
#define CONFIG_TQ_Z_IMAGE_DEFAULT_SIZE "1024*1536"
#endif

#ifndef CONFIG_TQ_Z_IMAGE_TIMEOUT_MS
#define CONFIG_TQ_Z_IMAGE_TIMEOUT_MS 45000
#endif

#ifndef CONFIG_TQ_Z_IMAGE_DOWNLOAD_TIMEOUT_MS
#define CONFIG_TQ_Z_IMAGE_DOWNLOAD_TIMEOUT_MS 15000
#endif

#ifndef CONFIG_TQ_Z_IMAGE_MAX_RESPONSE_BYTES
#define CONFIG_TQ_Z_IMAGE_MAX_RESPONSE_BYTES 65536
#endif

#ifndef CONFIG_TQ_Z_IMAGE_MAX_PNG_BYTES
#define CONFIG_TQ_Z_IMAGE_MAX_PNG_BYTES (4 * 1024 * 1024)
#endif

static const char *TAG = "image_generation";

typedef struct {
    uint8_t *data;
    size_t len;
    size_t cap;
    size_t max_len;
} bytes_buffer_t;

static void image_generation_fill_err(char *err, size_t err_len, const char *msg) {
    if (err != NULL && err_len > 0) {
        snprintf(err, err_len, "%s", msg != NULL ? msg : "unknown error");
    }
}

static void *malloc_prefer_psram(size_t size) {
    void *ptr = heap_caps_malloc(size, MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT);
    if (ptr == NULL) {
        ptr = malloc(size);
    }
    return ptr;
}

static void *realloc_prefer_psram(void *old, size_t size) {
    if (old == NULL) {
        return malloc_prefer_psram(size);
    }
    return realloc(old, size);
}

static esp_err_t bytes_buffer_reserve(bytes_buffer_t *buf, size_t need_extra) {
    if (buf == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    if (need_extra > SIZE_MAX - buf->len) {
        return ESP_ERR_INVALID_SIZE;
    }
    size_t required = buf->len + need_extra;
    if (required > buf->max_len) {
        return ESP_ERR_INVALID_SIZE;
    }
    if (required <= buf->cap) {
        return ESP_OK;
    }

    size_t new_cap = buf->cap == 0 ? 1024 : buf->cap;
    while (new_cap < required) {
        if (new_cap > (SIZE_MAX / 2)) {
            new_cap = required;
            break;
        }
        new_cap *= 2;
    }

    uint8_t *new_data = (uint8_t *)realloc_prefer_psram(buf->data, new_cap);
    if (new_data == NULL) {
        return ESP_ERR_NO_MEM;
    }
    buf->data = new_data;
    buf->cap = new_cap;
    return ESP_OK;
}

static esp_err_t bytes_buffer_append(bytes_buffer_t *buf, const void *data, size_t len) {
    if (buf == NULL || data == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    if (len == 0) {
        return ESP_OK;
    }
    esp_err_t rc = bytes_buffer_reserve(buf, len);
    if (rc != ESP_OK) {
        return rc;
    }
    memcpy(buf->data + buf->len, data, len);
    buf->len += len;
    return ESP_OK;
}

static void bytes_buffer_free(bytes_buffer_t *buf) {
    if (buf == NULL) {
        return;
    }
    free(buf->data);
    buf->data = NULL;
    buf->len = 0;
    buf->cap = 0;
}

static esp_err_t bytes_buffer_terminate(bytes_buffer_t *buf) {
    if (buf == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    esp_err_t rc = bytes_buffer_reserve(buf, 1);
    if (rc != ESP_OK) {
        return rc;
    }
    buf->data[buf->len] = '\0';
    return ESP_OK;
}

static const char *image_generation_api_key(void) {
    if (strlen(CONFIG_TQ_Z_IMAGE_API_KEY) > 0) {
        return CONFIG_TQ_Z_IMAGE_API_KEY;
    }
    if (strlen(CONFIG_TQ_VOICE_API_KEY) > 0) {
        return CONFIG_TQ_VOICE_API_KEY;
    }
    return NULL;
}

static esp_err_t http_read_full_response(esp_http_client_handle_t client,
                                         bytes_buffer_t *out_body,
                                         char *err,
                                         size_t err_len) {
    if (client == NULL || out_body == NULL) {
        image_generation_fill_err(err, err_len, "invalid http read args");
        return ESP_ERR_INVALID_ARG;
    }

    char tmp[1024];
    while (true) {
        int n = esp_http_client_read(client, tmp, sizeof(tmp));
        if (n < 0) {
            image_generation_fill_err(err, err_len, "http read failed");
            return ESP_FAIL;
        }
        if (n == 0) {
            if (esp_http_client_is_complete_data_received(client)) {
                return ESP_OK;
            }
            image_generation_fill_err(err, err_len, "http response incomplete");
            return ESP_FAIL;
        }

        esp_err_t append_rc = bytes_buffer_append(out_body, tmp, (size_t)n);
        if (append_rc != ESP_OK) {
            image_generation_fill_err(err, err_len, "http response too large");
            return append_rc;
        }
    }
}

static esp_err_t http_request_collect(const esp_http_client_config_t *config,
                                      const char *auth_header,
                                      const char *body,
                                      int *out_status_code,
                                      bytes_buffer_t *out_body,
                                      char *err,
                                      size_t err_len) {
    if (config == NULL || out_body == NULL) {
        image_generation_fill_err(err, err_len, "invalid http request args");
        return ESP_ERR_INVALID_ARG;
    }

    esp_http_client_handle_t client = esp_http_client_init(config);
    if (client == NULL) {
        image_generation_fill_err(err, err_len, "http client init failed");
        return ESP_FAIL;
    }

    if (auth_header != NULL) {
        esp_http_client_set_header(client, "Authorization", auth_header);
    }
    if (body != NULL) {
        esp_http_client_set_header(client, "Content-Type", "application/json");
    }

    size_t body_len = body != NULL ? strlen(body) : 0;
    esp_err_t open_rc = esp_http_client_open(client, (int)body_len);
    if (open_rc != ESP_OK) {
        esp_http_client_cleanup(client);
        image_generation_fill_err(err, err_len, "http open failed");
        return open_rc;
    }

    if (body != NULL && body_len > 0) {
        int written = esp_http_client_write(client, body, (int)body_len);
        if (written < 0 || (size_t)written != body_len) {
            esp_http_client_close(client);
            esp_http_client_cleanup(client);
            image_generation_fill_err(err, err_len, "http request write failed");
            return ESP_FAIL;
        }
    }

    int64_t content_len = esp_http_client_fetch_headers(client);
    if (content_len > 0 && (size_t)content_len > out_body->max_len) {
        esp_http_client_close(client);
        esp_http_client_cleanup(client);
        image_generation_fill_err(err, err_len, "http response too large");
        return ESP_ERR_INVALID_SIZE;
    }
    if (out_status_code != NULL) {
        *out_status_code = esp_http_client_get_status_code(client);
    }

    esp_err_t read_rc = http_read_full_response(client, out_body, err, err_len);
    esp_http_client_close(client);
    esp_http_client_cleanup(client);
    return read_rc;
}

static const char *json_get_string(cJSON *obj, const char *name) {
    if (obj == NULL || name == NULL) {
        return NULL;
    }
    cJSON *item = cJSON_GetObjectItemCaseSensitive(obj, name);
    if (cJSON_IsString(item) && item->valuestring != NULL) {
        return item->valuestring;
    }
    return NULL;
}

static bool json_extract_image_url(cJSON *root,
                                   char *out_url,
                                   size_t out_url_len,
                                   char *out_prompt,
                                   size_t out_prompt_len) {
    if (root == NULL || out_url == NULL || out_url_len == 0) {
        return false;
    }

    cJSON *output = cJSON_GetObjectItemCaseSensitive(root, "output");
    cJSON *choices = cJSON_GetObjectItemCaseSensitive(output, "choices");
    if (!cJSON_IsArray(choices)) {
        return false;
    }

    cJSON *choice = cJSON_GetArrayItem(choices, 0);
    if (!cJSON_IsObject(choice)) {
        return false;
    }

    cJSON *message = cJSON_GetObjectItemCaseSensitive(choice, "message");
    cJSON *content = cJSON_GetObjectItemCaseSensitive(message, "content");
    if (!cJSON_IsArray(content)) {
        return false;
    }

    out_url[0] = '\0';
    if (out_prompt != NULL && out_prompt_len > 0) {
        out_prompt[0] = '\0';
    }

    int count = cJSON_GetArraySize(content);
    for (int i = 0; i < count; ++i) {
        cJSON *item = cJSON_GetArrayItem(content, i);
        if (!cJSON_IsObject(item)) {
            continue;
        }
        const char *image = json_get_string(item, "image");
        if (image != NULL && out_url[0] == '\0') {
            strlcpy(out_url, image, out_url_len);
        }

        if (out_prompt != NULL && out_prompt_len > 0 && out_prompt[0] == '\0') {
            const char *text = json_get_string(item, "text");
            if (text != NULL) {
                strlcpy(out_prompt, text, out_prompt_len);
            }
        }
    }

    return out_url[0] != '\0';
}

static void extract_remote_error_message(const char *json_text,
                                         int status_code,
                                         char *err,
                                         size_t err_len) {
    if (json_text == NULL) {
        image_generation_fill_err(err, err_len, "upstream request failed");
        return;
    }

    cJSON *root = cJSON_Parse(json_text);
    if (root == NULL) {
        if (status_code > 0) {
            snprintf(err, err_len, "upstream request failed, http_status=%d", status_code);
        } else {
            image_generation_fill_err(err, err_len, "upstream request failed");
        }
        return;
    }

    const char *code = json_get_string(root, "code");
    const char *message = json_get_string(root, "message");

    if (code != NULL && message != NULL) {
        snprintf(err, err_len, "%s: %s", code, message);
    } else if (message != NULL) {
        image_generation_fill_err(err, err_len, message);
    } else if (status_code > 0) {
        snprintf(err, err_len, "upstream request failed, http_status=%d", status_code);
    } else {
        image_generation_fill_err(err, err_len, "upstream request failed");
    }

    cJSON_Delete(root);
}

static char *build_generation_request_json(const image_generation_request_t *req) {
    cJSON *root = cJSON_CreateObject();
    if (root == NULL) {
        return NULL;
    }

    cJSON_AddStringToObject(root, "model", CONFIG_TQ_Z_IMAGE_MODEL);

    cJSON *input = cJSON_AddObjectToObject(root, "input");
    cJSON *messages = cJSON_AddArrayToObject(input, "messages");
    cJSON *message = cJSON_CreateObject();
    cJSON_AddItemToArray(messages, message);
    cJSON_AddStringToObject(message, "role", "user");
    cJSON *content = cJSON_AddArrayToObject(message, "content");
    cJSON *content_item = cJSON_CreateObject();
    cJSON_AddItemToArray(content, content_item);
    cJSON_AddStringToObject(content_item, "text", req->prompt);

    cJSON *parameters = cJSON_AddObjectToObject(root, "parameters");
    cJSON_AddBoolToObject(parameters, "prompt_extend", req->prompt_extend);

    const char *size = (req->size != NULL && req->size[0] != '\0')
                           ? req->size
                           : CONFIG_TQ_Z_IMAGE_DEFAULT_SIZE;
    cJSON_AddStringToObject(parameters, "size", size);
    if (req->has_seed) {
        cJSON_AddNumberToObject(parameters, "seed", (double)req->seed);
    }

    char *body = cJSON_PrintUnformatted(root);
    cJSON_Delete(root);
    return body;
}

static esp_err_t image_generation_invoke_model(const image_generation_request_t *req,
                                               image_generation_result_t *out_result,
                                               char *err,
                                               size_t err_len) {
    char *req_body = build_generation_request_json(req);
    if (req_body == NULL) {
        image_generation_fill_err(err, err_len, "build request json failed");
        return ESP_ERR_NO_MEM;
    }

    const char *api_key = image_generation_api_key();
    if (api_key == NULL) {
        cJSON_free(req_body);
        image_generation_fill_err(err, err_len, "z-image api key not configured");
        return ESP_ERR_INVALID_STATE;
    }

    char auth_header[192] = {0};
    snprintf(auth_header, sizeof(auth_header), "Bearer %s", api_key);

    esp_http_client_config_t config = {
        .url = CONFIG_TQ_Z_IMAGE_API_ENDPOINT,
        .method = HTTP_METHOD_POST,
        .timeout_ms = (int)req->generation_timeout_ms,
        .keep_alive_enable = false,
        .buffer_size = 2048,
        .buffer_size_tx = 2048,
        .crt_bundle_attach = esp_crt_bundle_attach,
    };

    bytes_buffer_t resp = {
        .data = NULL,
        .len = 0,
        .cap = 0,
        .max_len = CONFIG_TQ_Z_IMAGE_MAX_RESPONSE_BYTES,
    };
    int http_status = 0;
    esp_err_t rc = http_request_collect(&config,
                                        auth_header,
                                        req_body,
                                        &http_status,
                                        &resp,
                                        err,
                                        err_len);
    cJSON_free(req_body);
    if (rc != ESP_OK) {
        bytes_buffer_free(&resp);
        return rc;
    }

    rc = bytes_buffer_terminate(&resp);
    if (rc != ESP_OK) {
        bytes_buffer_free(&resp);
        image_generation_fill_err(err, err_len, "generation response too large");
        return rc;
    }

    if (http_status < 200 || http_status >= 300) {
        extract_remote_error_message((const char *)resp.data, http_status, err, err_len);
        bytes_buffer_free(&resp);
        return ESP_FAIL;
    }

    cJSON *root = cJSON_Parse((const char *)resp.data);
    bytes_buffer_free(&resp);
    if (root == NULL) {
        image_generation_fill_err(err, err_len, "invalid generation response");
        return ESP_FAIL;
    }

    const char *request_id = json_get_string(root, "request_id");
    if (request_id != NULL) {
        strlcpy(out_result->request_id, request_id, sizeof(out_result->request_id));
    }

    cJSON *usage = cJSON_GetObjectItemCaseSensitive(root, "usage");
    cJSON *usage_w = cJSON_GetObjectItemCaseSensitive(usage, "width");
    cJSON *usage_h = cJSON_GetObjectItemCaseSensitive(usage, "height");
    if (cJSON_IsNumber(usage_w) && usage_w->valuedouble >= 1 && usage_w->valuedouble <= UINT16_MAX) {
        out_result->width = (uint16_t)usage_w->valuedouble;
    }
    if (cJSON_IsNumber(usage_h) && usage_h->valuedouble >= 1 && usage_h->valuedouble <= UINT16_MAX) {
        out_result->height = (uint16_t)usage_h->valuedouble;
    }

    bool has_url = json_extract_image_url(root,
                                          out_result->image_url,
                                          sizeof(out_result->image_url),
                                          out_result->output_prompt,
                                          sizeof(out_result->output_prompt));
    if (!has_url) {
        const char *code = json_get_string(root, "code");
        const char *message = json_get_string(root, "message");
        if (code != NULL && message != NULL) {
            snprintf(err, err_len, "%s: %s", code, message);
        } else {
            image_generation_fill_err(err, err_len, "generation response has no image url");
        }
        cJSON_Delete(root);
        return ESP_FAIL;
    }

    cJSON_Delete(root);
    return ESP_OK;
}

static esp_err_t image_generation_download_png(image_generation_result_t *out_result,
                                               uint32_t timeout_ms,
                                               char *err,
                                               size_t err_len) {
    esp_http_client_config_t config = {
        .url = out_result->image_url,
        .method = HTTP_METHOD_GET,
        .timeout_ms = (int)timeout_ms,
        .keep_alive_enable = false,
        .buffer_size = 2048,
        .crt_bundle_attach = esp_crt_bundle_attach,
    };

    bytes_buffer_t resp = {
        .data = NULL,
        .len = 0,
        .cap = 0,
        .max_len = CONFIG_TQ_Z_IMAGE_MAX_PNG_BYTES,
    };
    int http_status = 0;
    esp_err_t rc = http_request_collect(&config,
                                        NULL,
                                        NULL,
                                        &http_status,
                                        &resp,
                                        err,
                                        err_len);
    if (rc != ESP_OK) {
        bytes_buffer_free(&resp);
        return rc;
    }

    if (http_status < 200 || http_status >= 300) {
        rc = bytes_buffer_terminate(&resp);
        if (rc != ESP_OK) {
            bytes_buffer_free(&resp);
            image_generation_fill_err(err, err_len, "image download failed");
            return ESP_FAIL;
        }
        extract_remote_error_message((const char *)resp.data, http_status, err, err_len);
        bytes_buffer_free(&resp);
        return ESP_FAIL;
    }

    static const uint8_t png_header[8] = {
        0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A,
    };
    if (resp.len < sizeof(png_header) || memcmp(resp.data, png_header, sizeof(png_header)) != 0) {
        bytes_buffer_free(&resp);
        image_generation_fill_err(err, err_len, "downloaded image is not png");
        return ESP_FAIL;
    }

    out_result->png = resp.data;
    out_result->png_len = resp.len;
    return ESP_OK;
}

void image_generation_result_reset(image_generation_result_t *result) {
    if (result == NULL) {
        return;
    }
    memset(result, 0, sizeof(*result));
}

void image_generation_result_free(image_generation_result_t *result) {
    if (result == NULL) {
        return;
    }
    free(result->png);
    result->png = NULL;
    result->png_len = 0;
}

esp_err_t image_generation_generate_png(const image_generation_request_t *req,
                                        image_generation_result_t *out_result,
                                        char *err,
                                        size_t err_len) {
    if (req == NULL || out_result == NULL || req->prompt == NULL || req->prompt[0] == '\0') {
        image_generation_fill_err(err, err_len, "prompt is required");
        return ESP_ERR_INVALID_ARG;
    }

    image_generation_result_reset(out_result);

    image_generation_request_t actual = *req;
    if (actual.generation_timeout_ms == 0) {
        actual.generation_timeout_ms = CONFIG_TQ_Z_IMAGE_TIMEOUT_MS;
    }
    if (actual.download_timeout_ms == 0) {
        actual.download_timeout_ms = CONFIG_TQ_Z_IMAGE_DOWNLOAD_TIMEOUT_MS;
    }

    ESP_LOGI(TAG, "invoke z-image, prompt_len=%u", (unsigned)strlen(actual.prompt));
    esp_err_t gen_rc = image_generation_invoke_model(&actual, out_result, err, err_len);
    if (gen_rc != ESP_OK) {
        image_generation_result_free(out_result);
        return gen_rc;
    }

    ESP_LOGI(TAG, "download image from model output");
    esp_err_t dl_rc = image_generation_download_png(out_result, actual.download_timeout_ms, err, err_len);
    if (dl_rc != ESP_OK) {
        image_generation_result_free(out_result);
        return dl_rc;
    }

    return ESP_OK;
}