AI_Printer/components/control_plane/src/rest_server_print_image.c

#include "rest_server_internal.h"

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

#include "esp_log.h"
#include "domain.h"

static const char *TAG = "rest_print";
#define SCREEN_PREVIEW_TIMEOUT_MS 200

typedef struct {
    bool scale_to_width;
    uint8_t threshold;
    bool invert;
    uint16_t max_height;
    const char *density;
    bool direct_ignore_precheck;
} image_print_options_t;

typedef struct {
    char prompt[801];
    char size[24];
    bool has_size;
    bool prompt_extend;
    bool has_seed;
    uint32_t seed;
    uint32_t timeout_ms;
    uint32_t fetch_timeout_ms;
} image_generate_options_t;

static bool parse_bool_text(const char *text, bool *out_value) {
    if (text == NULL || out_value == NULL) {
        return false;
    }
    if (strcmp(text, "1") == 0 || strcasecmp(text, "true") == 0 || strcasecmp(text, "yes") == 0) {
        *out_value = true;
        return true;
    }
    if (strcmp(text, "0") == 0 || strcasecmp(text, "false") == 0 || strcasecmp(text, "no") == 0) {
        *out_value = false;
        return true;
    }
    return false;
}

static const char *map_density_text(const char *text) {
    if (text == NULL || text[0] == '\0') {
        return "中等";
    }
    if (strcasecmp(text, "light") == 0 || strcmp(text, "较淡") == 0) {
        return "较淡";
    }
    if (strcasecmp(text, "medium") == 0 || strcmp(text, "中等") == 0) {
        return "中等";
    }
    if (strcasecmp(text, "dark") == 0 || strcmp(text, "较浓") == 0) {
        return "较浓";
    }
    if (strcasecmp(text, "max") == 0 || strcmp(text, "最深") == 0) {
        return "最深";
    }
    return "中等";
}

static void parse_image_upload_options(httpd_req_t *req,
                                       bool *out_scale_to_width,
                                       uint8_t *out_threshold,
                                       bool *out_invert,
                                       uint16_t *out_max_height,
                                       const char **out_density,
                                       bool *out_direct_ignore_precheck) {
    if (out_scale_to_width != NULL) {
        *out_scale_to_width = true;
    }
    if (out_threshold != NULL) {
        *out_threshold = 160;
    }
    if (out_invert != NULL) {
        *out_invert = false;
    }
    if (out_max_height != NULL) {
        *out_max_height = 2200;
    }
    if (out_density != NULL) {
        *out_density = "中等";
    }
    if (out_direct_ignore_precheck != NULL) {
        *out_direct_ignore_precheck = false;
    }

    if (req == NULL) {
        return;
    }

    int qlen = httpd_req_get_url_query_len(req);
    if (qlen <= 0 || qlen > 512) {
        return;
    }

    char *query = (char *)calloc(1, (size_t)qlen + 1);
    if (query == NULL) {
        return;
    }

    if (httpd_req_get_url_query_str(req, query, (size_t)qlen + 1) != ESP_OK) {
        free(query);
        return;
    }

    char value[48] = {0};
    if (httpd_query_key_value(query, "scale_to_width", value, sizeof(value)) == ESP_OK &&
        out_scale_to_width != NULL) {
        bool parsed = *out_scale_to_width;
        if (parse_bool_text(value, &parsed)) {
            *out_scale_to_width = parsed;
        }
    }

    if (httpd_query_key_value(query, "invert", value, sizeof(value)) == ESP_OK && out_invert != NULL) {
        bool parsed = *out_invert;
        if (parse_bool_text(value, &parsed)) {
            *out_invert = parsed;
        }
    }

    if (httpd_query_key_value(query, "threshold", value, sizeof(value)) == ESP_OK && out_threshold != NULL) {
        long v = strtol(value, NULL, 10);
        if (v >= 0 && v <= 255) {
            *out_threshold = (uint8_t)v;
        }
    }

    if (httpd_query_key_value(query, "max_height", value, sizeof(value)) == ESP_OK && out_max_height != NULL) {
        long v = strtol(value, NULL, 10);
        if (v >= 64 && v <= 3000) {
            *out_max_height = (uint16_t)v;
        }
    }

    if (httpd_query_key_value(query, "density", value, sizeof(value)) == ESP_OK && out_density != NULL) {
        *out_density = map_density_text(value);
    }

    if (httpd_query_key_value(query, "ignore_precheck", value, sizeof(value)) == ESP_OK &&
        out_direct_ignore_precheck != NULL) {
        bool parsed = *out_direct_ignore_precheck;
        if (parse_bool_text(value, &parsed)) {
            *out_direct_ignore_precheck = parsed;
        }
    }

    free(query);
}

static void image_print_options_set_defaults(image_print_options_t *out) {
    if (out == NULL) {
        return;
    }
    out->scale_to_width = true;
    out->threshold = 160;
    out->invert = false;
    out->max_height = 2200;
    out->density = "中等";
    out->direct_ignore_precheck = false;
}

static void parse_image_print_options_from_json(cJSON *json, image_print_options_t *out) {
    if (json == NULL || out == NULL) {
        return;
    }

    cJSON *jscale = cJSON_GetObjectItemCaseSensitive(json, "scale_to_width");
    cJSON *jthreshold = cJSON_GetObjectItemCaseSensitive(json, "threshold");
    cJSON *jinvert = cJSON_GetObjectItemCaseSensitive(json, "invert");
    cJSON *jmaxh = cJSON_GetObjectItemCaseSensitive(json, "max_height");
    cJSON *jdensity = cJSON_GetObjectItemCaseSensitive(json, "density");
    cJSON *jignore = cJSON_GetObjectItemCaseSensitive(json, "ignore_precheck");

    out->scale_to_width = rest_server_json_bool_with_default(jscale, out->scale_to_width);
    out->invert = rest_server_json_bool_with_default(jinvert, out->invert);
    if (cJSON_IsNumber(jthreshold) && jthreshold->valuedouble >= 0 && jthreshold->valuedouble <= 255) {
        out->threshold = (uint8_t)jthreshold->valuedouble;
    }
    if (cJSON_IsNumber(jmaxh) && jmaxh->valuedouble >= 64 && jmaxh->valuedouble <= 3000) {
        out->max_height = (uint16_t)jmaxh->valuedouble;
    }
    if (cJSON_IsString(jdensity) && jdensity->valuestring != NULL) {
        out->density = map_density_text(jdensity->valuestring);
    }
    out->direct_ignore_precheck = rest_server_json_bool_with_default(jignore, out->direct_ignore_precheck);
}

static bool parse_image_size_text(const char *text, uint32_t *out_w, uint32_t *out_h) {
    if (text == NULL || out_w == NULL || out_h == NULL) {
        return false;
    }

    const char *sep = strchr(text, '*');
    if (sep == NULL || sep == text || *(sep + 1) == '\0') {
        return false;
    }

    char wbuf[12] = {0};
    char hbuf[12] = {0};
    size_t wlen = (size_t)(sep - text);
    size_t hlen = strlen(sep + 1);
    if (wlen == 0 || hlen == 0 || wlen >= sizeof(wbuf) || hlen >= sizeof(hbuf)) {
        return false;
    }

    memcpy(wbuf, text, wlen);
    memcpy(hbuf, sep + 1, hlen);
    wbuf[wlen] = '\0';
    hbuf[hlen] = '\0';

    char *wend = NULL;
    char *hend = NULL;
    long w = strtol(wbuf, &wend, 10);
    long h = strtol(hbuf, &hend, 10);
    if (wend == wbuf || hend == hbuf || *wend != '\0' || *hend != '\0') {
        return false;
    }
    if (w <= 0 || h <= 0) {
        return false;
    }

    *out_w = (uint32_t)w;
    *out_h = (uint32_t)h;
    return true;
}

static bool parse_image_generation_options(cJSON *json,
                                           image_generate_options_t *out,
                                           char *err,
                                           size_t err_len) {
    if (json == NULL || out == NULL) {
        if (err != NULL && err_len > 0) {
            snprintf(err, err_len, "invalid args");
        }
        return false;
    }

    memset(out, 0, sizeof(*out));
    out->timeout_ms = domain_policy_image_generation_timeout_default_ms();
    out->fetch_timeout_ms = domain_policy_image_download_timeout_default_ms();

    cJSON *jprompt = cJSON_GetObjectItemCaseSensitive(json, "prompt");
    if (!cJSON_IsString(jprompt) || jprompt->valuestring == NULL || jprompt->valuestring[0] == '\0') {
        if (err != NULL && err_len > 0) {
            snprintf(err, err_len, "prompt is required");
        }
        return false;
    }

    size_t prompt_len = strlen(jprompt->valuestring);
    if (prompt_len > 800) {
        if (err != NULL && err_len > 0) {
            snprintf(err, err_len, "prompt too long, max 800 chars");
        }
        return false;
    }
    strlcpy(out->prompt, jprompt->valuestring, sizeof(out->prompt));

    cJSON *jsize = cJSON_GetObjectItemCaseSensitive(json, "size");
    if (cJSON_IsString(jsize) && jsize->valuestring != NULL && jsize->valuestring[0] != '\0') {
        uint32_t w = 0;
        uint32_t h = 0;
        if (!parse_image_size_text(jsize->valuestring, &w, &h)) {
            if (err != NULL && err_len > 0) {
                snprintf(err, err_len, "size must be like 1120*1440");
            }
            return false;
        }

        uint64_t pixels = (uint64_t)w * (uint64_t)h;
        if (pixels < (512ULL * 512ULL) || pixels > (2048ULL * 2048ULL)) {
            if (err != NULL && err_len > 0) {
                snprintf(err, err_len, "size pixels must be in [512*512, 2048*2048]");
            }
            return false;
        }

        out->has_size = true;
        strlcpy(out->size, jsize->valuestring, sizeof(out->size));
    }

    cJSON *jextend = cJSON_GetObjectItemCaseSensitive(json, "prompt_extend");
    out->prompt_extend = rest_server_json_bool_with_default(jextend, false);

    cJSON *jseed = cJSON_GetObjectItemCaseSensitive(json, "seed");
    if (cJSON_IsNumber(jseed)) {
        if (jseed->valuedouble < 0 || jseed->valuedouble > 2147483647.0) {
            if (err != NULL && err_len > 0) {
                snprintf(err, err_len, "seed must be in [0, 2147483647]");
            }
            return false;
        }
        out->has_seed = true;
        out->seed = (uint32_t)jseed->valuedouble;
    }

    cJSON *jtimeout = cJSON_GetObjectItemCaseSensitive(json, "timeout_ms");
    if (cJSON_IsNumber(jtimeout)) {
        if (jtimeout->valuedouble < domain_policy_image_generation_timeout_min_ms() ||
            jtimeout->valuedouble > domain_policy_image_generation_timeout_max_ms()) {
            if (err != NULL && err_len > 0) {
                snprintf(err,
                         err_len,
                         "timeout_ms must be %u..%u",
                         (unsigned)domain_policy_image_generation_timeout_min_ms(),
                         (unsigned)domain_policy_image_generation_timeout_max_ms());
            }
            return false;
        }
        out->timeout_ms = (uint32_t)jtimeout->valuedouble;
    }

    cJSON *jfetch_timeout = cJSON_GetObjectItemCaseSensitive(json, "fetch_timeout_ms");
    if (cJSON_IsNumber(jfetch_timeout)) {
        if (jfetch_timeout->valuedouble < domain_policy_image_download_timeout_min_ms() ||
            jfetch_timeout->valuedouble > domain_policy_image_download_timeout_max_ms()) {
            if (err != NULL && err_len > 0) {
                snprintf(err,
                         err_len,
                         "fetch_timeout_ms must be %u..%u",
                         (unsigned)domain_policy_image_download_timeout_min_ms(),
                         (unsigned)domain_policy_image_download_timeout_max_ms());
            }
            return false;
        }
        out->fetch_timeout_ms = (uint32_t)jfetch_timeout->valuedouble;
    }

    return true;
}

static esp_err_t rest_server_send_image_generation_error(httpd_req_t *req,
                                                         esp_err_t rc,
                                                         const char *msg) {
    const char *message = (msg != NULL && msg[0] != '\0') ? msg : "image generation failed";
    if (rc == ESP_ERR_INVALID_ARG) {
        return rest_server_send_error(req, "400 Bad Request", message);
    }
    if (rc == ESP_ERR_INVALID_STATE) {
        return rest_server_send_error(req, "503 Service Unavailable", message);
    }
    if (rc == ESP_ERR_TIMEOUT) {
        return rest_server_send_error(req, "504 Gateway Timeout", message);
    }
    if (rc == ESP_ERR_NO_MEM) {
        return rest_server_send_error(req, "500 Internal Server Error", message);
    }
    return rest_server_send_error(req, "502 Bad Gateway", message);
}

static void rest_server_try_preview_raster(const char *source,
                                           const uint8_t *raster,
                                           size_t raster_len,
                                           uint16_t width,
                                           uint16_t height) {
    if (raster == NULL || raster_len == 0 || width == 0 || height == 0) {
        return;
    }

    char preview_err[96] = {0};
    esp_err_t preview_rc = screen_preview_show_raster(raster,
                                                      raster_len,
                                                      width,
                                                      height,
                                                      SCREEN_PREVIEW_TIMEOUT_MS,
                                                      preview_err,
                                                      sizeof(preview_err));
    if (preview_rc == ESP_OK) {
        ESP_LOGI(TAG,
                 "screen preview updated (%s), size=%ux%u",
                 source != NULL ? source : "unknown",
                 (unsigned)width,
                 (unsigned)height);
        return;
    }
    if (preview_rc == ESP_ERR_NOT_SUPPORTED) {
        return;
    }

    ESP_LOGW(TAG,
             "screen preview skipped (%s), rc=0x%x, msg=%s",
             source != NULL ? source : "unknown",
             (unsigned)preview_rc,
             preview_err[0] != '\0' ? preview_err : "preview failed");
}

static esp_err_t rest_server_print_image_binary(httpd_req_t *req) {
    bool scale_to_width = true;
    uint8_t threshold = 160;
    bool invert = false;
    uint16_t max_height = 2200;
    const char *density = "中等";
    bool direct_ignore_precheck = false;
    parse_image_upload_options(req,
                               &scale_to_width,
                               &threshold,
                               &invert,
                               &max_height,
                               &density,
                               &direct_ignore_precheck);

    ESP_LOGI(TAG,
             "image process start (upload), png_bytes=%u, scale_to_width=%d, threshold=%u, invert=%d, max_height=%u, density=%s, ignore_precheck=%d",
             (unsigned)req->content_len,
             scale_to_width,
             (unsigned)threshold,
             invert,
             (unsigned)max_height,
             density != NULL ? density : "中等",
             direct_ignore_precheck);

    char *body = NULL;
    esp_err_t body_err = rest_server_read_body(req, &body);
    if (body_err != ESP_OK) {
        return rest_server_send_error(req, "400 Bad Request", "invalid request body");
    }

    uint16_t width = 0;
    uint16_t height = 0;
    uint8_t *raster = NULL;
    size_t raster_len = 0;
    char decode_err[128] = {0};
    esp_err_t decode_rc = raster_tools_convert_png_to_raster_384((const uint8_t *)body,
                                                                  (size_t)req->content_len,
                                                                  scale_to_width,
                                                                  threshold,
                                                                  invert,
                                                                  max_height,
                                                                  &width,
                                                                  &height,
                                                                  &raster,
                                                                  &raster_len,
                                                                  decode_err,
                                                                  sizeof(decode_err));
    free(body);
    if (decode_rc != ESP_OK) {
        ESP_LOGW(TAG,
                 "image process failed (upload), rc=0x%x, msg=%s",
                 (unsigned)decode_rc,
                 decode_err[0] != '\0' ? decode_err : "decode png failed");
        return rest_server_send_error(req,
                                      "400 Bad Request",
                                      decode_err[0] != '\0' ? decode_err : "decode png failed");
    }

    ESP_LOGI(TAG,
             "image process done (upload), raster_size=%ux%u, raster_bytes=%u",
             (unsigned)width,
             (unsigned)height,
             (unsigned)raster_len);
    rest_server_try_preview_raster("upload", raster, raster_len, width, height);

    esp_err_t submit_rc = rest_server_print_submit_raster_job_and_reply(req,
                                                                         raster,
                                                                         raster_len,
                                                                         width,
                                                                         height,
                                                                         density,
                                                                         direct_ignore_precheck,
                                                                         NULL,
                                                                         NULL);
    free(raster);
    return submit_rc;
}

static esp_err_t rest_server_print_image_generate(httpd_req_t *req) {
    char *body = NULL;
    cJSON *json = NULL;
    image_generation_result_t *gen_result = NULL;
    uint8_t *raster = NULL;
    char request_id[80] = {0};
    printer_status_poll_pause_token_t status_poll_pause_token = 0;
    esp_err_t ret = ESP_FAIL;

    esp_err_t body_err = rest_server_read_body(req, &body);
    if (body_err != ESP_OK) {
        return rest_server_send_error(req, "400 Bad Request", "invalid request body");
    }

    json = cJSON_Parse(body);
    free(body);
    body = NULL;
    if (json == NULL) {
        return rest_server_send_error(req, "400 Bad Request", "invalid json");
    }

    image_generate_options_t gen_opt = {0};
    char parse_err[128] = {0};
    if (!parse_image_generation_options(json, &gen_opt, parse_err, sizeof(parse_err))) {
        ret = rest_server_send_error(req,
                                     "400 Bad Request",
                                     parse_err[0] != '\0' ? parse_err : "invalid image options");
        goto cleanup;
    }

    image_print_options_t print_opt = {0};
    image_print_options_set_defaults(&print_opt);
    parse_image_print_options_from_json(json, &print_opt);
    cJSON_Delete(json);
    json = NULL;

    ESP_LOGI(TAG,
             "image generate start, prompt_len=%u, size=%s, prompt_extend=%d, seed=%s, timeout_ms=%u, fetch_timeout_ms=%u",
             (unsigned)strlen(gen_opt.prompt),
             gen_opt.has_size ? gen_opt.size : "default",
             gen_opt.prompt_extend,
             gen_opt.has_seed ? "set" : "auto",
             (unsigned)gen_opt.timeout_ms,
             (unsigned)gen_opt.fetch_timeout_ms);
    ESP_LOGI(TAG,
             "image process config, scale_to_width=%d, threshold=%u, invert=%d, max_height=%u, density=%s, ignore_precheck=%d",
             print_opt.scale_to_width,
             (unsigned)print_opt.threshold,
             print_opt.invert,
             (unsigned)print_opt.max_height,
             print_opt.density != NULL ? print_opt.density : "中等",
             print_opt.direct_ignore_precheck);

    image_generation_request_t gen_req = {
        .prompt = gen_opt.prompt,
        .size = gen_opt.has_size ? gen_opt.size : NULL,
        .prompt_extend = gen_opt.prompt_extend,
        .has_seed = gen_opt.has_seed,
        .seed = gen_opt.seed,
        .generation_timeout_ms = gen_opt.timeout_ms,
        .download_timeout_ms = gen_opt.fetch_timeout_ms,
    };

    gen_result = (image_generation_result_t *)calloc(1, sizeof(*gen_result));
    if (gen_result == NULL) {
        ret = rest_server_send_error(req, "500 Internal Server Error", "no memory");
        goto cleanup;
    }
    image_generation_result_reset(gen_result);

    char model_err[160] = {0};
    status_poll_pause_token = printer_protocol_status_poll_pause_acquire();
    esp_err_t gen_rc = image_generation_generate_png(&gen_req,
                                                     gen_result,
                                                     model_err,
                                                     sizeof(model_err));
    printer_protocol_status_poll_pause_release(&status_poll_pause_token);
    if (gen_rc != ESP_OK) {
        ESP_LOGW(TAG,
                 "image generate failed, rc=0x%x, msg=%s",
                 (unsigned)gen_rc,
                 model_err[0] != '\0' ? model_err : "image generation failed");
        ret = rest_server_send_image_generation_error(req, gen_rc, model_err);
        goto cleanup;
    }

    if (gen_result->request_id[0] != '\0') {
        strlcpy(request_id, gen_result->request_id, sizeof(request_id));
    }

    ESP_LOGI(TAG,
             "image generate done, request_id=%s, model_size=%ux%u, png_bytes=%u",
             request_id[0] != '\0' ? request_id : "-",
             (unsigned)gen_result->width,
             (unsigned)gen_result->height,
             (unsigned)gen_result->png_len);

    uint16_t width = 0;
    uint16_t height = 0;
    size_t raster_len = 0;
    char decode_err[128] = {0};
    esp_err_t decode_rc = raster_tools_convert_png_to_raster_384(gen_result->png,
                                                                  gen_result->png_len,
                                                                  print_opt.scale_to_width,
                                                                  print_opt.threshold,
                                                                  print_opt.invert,
                                                                  print_opt.max_height,
                                                                  &width,
                                                                  &height,
                                                                  &raster,
                                                                  &raster_len,
                                                                  decode_err,
                                                                  sizeof(decode_err));
    if (decode_rc != ESP_OK) {
        ESP_LOGW(TAG,
                 "image process failed (generated), rc=0x%x, msg=%s",
                 (unsigned)decode_rc,
                 decode_err[0] != '\0' ? decode_err : "decode generated image failed");
        ret = rest_server_send_error(req,
                                     "400 Bad Request",
                                     decode_err[0] != '\0' ? decode_err : "decode generated image failed");
        goto cleanup;
    }

    ESP_LOGI(TAG,
             "image process done (generated), raster_size=%ux%u, raster_bytes=%u",
             (unsigned)width,
             (unsigned)height,
             (unsigned)raster_len);
    rest_server_try_preview_raster("generated", raster, raster_len, width, height);

    char warning[160] = {0};
    if (request_id[0] != '\0') {
        snprintf(warning, sizeof(warning), "z-image request_id=%s", request_id);
    }

    ret = rest_server_print_submit_raster_job_and_reply(req,
                                                         raster,
                                                         raster_len,
                                                         width,
                                                         height,
                                                         print_opt.density,
                                                         print_opt.direct_ignore_precheck,
                                                         warning,
                                                         request_id);

cleanup:
    printer_protocol_status_poll_pause_release(&status_poll_pause_token);
    if (body != NULL) {
        free(body);
    }
    if (json != NULL) {
        cJSON_Delete(json);
    }
    if (raster != NULL) {
        free(raster);
    }
    if (gen_result != NULL) {
        image_generation_result_free(gen_result);
        free(gen_result);
    }
    printer_runtime_status_t runtime = {0};
    printer_protocol_get_runtime_status(&runtime);
    if (!runtime.busy) {
        image_generation_schedule_prewarm();
    }
    return ret;
}

esp_err_t rest_server_print_image_post(httpd_req_t *req) {
    if (!rest_server_auth_ok(req)) {
        return rest_server_send_error(req, "401 Unauthorized", "unauthorized");
    }

    char content_type[96] = {0};
    if (httpd_req_get_hdr_value_str(req, "Content-Type", content_type, sizeof(content_type)) == ESP_OK) {
        if (strstr(content_type, "application/json") != NULL) {
            return rest_server_print_image_generate(req);
        }
        if (strstr(content_type, "image/png") == NULL) {
            return rest_server_send_error(req,
                                          "415 Unsupported Media Type",
                                          "Content-Type must be application/json or image/png");
        }
    }

    return rest_server_print_image_binary(req);
}