#include "result_json.h"

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

static const char* task_name(yolo_task_type_t task) {
    switch (task) {
        case YOLO_TASK_SEG: return "seg";
        case YOLO_TASK_CLS: return "cls";
        case YOLO_TASK_OBB: return "obb";
        case YOLO_TASK_POSE: return "pose";
        default: return "detect";
    }
}

static void write_escaped(FILE* file, const char* text) {
    for (const unsigned char* cursor = (const unsigned char*)text; cursor && *cursor; ++cursor) {
        if (*cursor == '"' || *cursor == '\\') fputc('\\', file);
        if (*cursor >= 0x20) fputc(*cursor, file);
    }
}

int write_result_json(
    const char* path,
    const char* model_config,
    const yolo_result_set_t* results,
    yolo_task_type_t task) {
    uint32_t count = 0;
    yolo_performance_v1_t performance = {sizeof(performance), YOLO_PERFORMANCE_VERSION};
    if (!path || !model_config || !results ||
        yolo_result_set_get_count(results, &count) != YOLO_STATUS_OK ||
        yolo_result_set_get_performance(results, &performance) != YOLO_STATUS_OK) return -1;
    FILE* file = fopen(path, "wb");
    if (!file) return -2;
    fprintf(file, "{\n  \"schemaVersion\": \"2\",\n  \"model\": \"");
    write_escaped(file, model_config);
    fprintf(file, "\",\n  \"task\": \"%s\",\n", task_name(task));
    fprintf(file, "  \"performance\": {\"preprocessMs\": %.3f, \"inferenceMs\": %.3f, \"postprocessMs\": %.3f, \"totalMs\": %.3f},\n",
        performance.preprocess_ms, performance.inference_ms,
        performance.postprocess_ms, performance.total_ms);
    fputs("  \"results\": [", file);
    for (uint32_t index = 0; index < count; ++index) {
        yolo_result_common_v1_t common = {sizeof(common), YOLO_RESULT_COMMON_VERSION};
        yolo_bbox_v1_t bbox = {sizeof(bbox), YOLO_BBOX_VERSION};
        if (yolo_result_set_get_common(results, index, &common) != YOLO_STATUS_OK) {
            fclose(file); return -3;
        }
        fprintf(file, "%s\n    {\"task\": \"%s\", \"classId\": %d, \"label\": \"",
            index ? "," : "", task_name(common.task), common.class_id);
        write_escaped(file, common.label);
        fprintf(file, "\", \"score\": %.6f", common.score);
        if (task != YOLO_TASK_CLS) {
            if (yolo_result_set_get_bbox(results, index, &bbox) != YOLO_STATUS_OK) {
                fclose(file); return -3;
            }
            fprintf(file, ", \"bbox\": {\"x\": %.3f, \"y\": %.3f, \"width\": %.3f, \"height\": %.3f",
                bbox.x, bbox.y, bbox.width, bbox.height);
            if (task == YOLO_TASK_OBB) {
                yolo_rotated_rect_v1_t obb = {sizeof(obb), YOLO_ROTATED_RECT_VERSION};
                if (yolo_result_set_get_obb(results, index, &obb) != YOLO_STATUS_OK) {
                    fclose(file); return -3;
                }
                fprintf(file, ", \"angle\": %.6f", obb.angle_radians);
            }
            fputc('}', file);
        }
        if (task == YOLO_TASK_SEG) {
            yolo_mask_view_v1_t mask = {sizeof(mask), YOLO_MASK_VIEW_VERSION};
            yolo_seg_geometry_v1_t geometry = {sizeof(geometry), YOLO_SEG_GEOMETRY_VERSION};
            if (yolo_result_set_get_seg_mask(results, index, &mask) != YOLO_STATUS_OK ||
                yolo_result_set_get_seg_geometry(results, index, &geometry) != YOLO_STATUS_OK ||
                !mask.data || mask.width <= 0 || mask.height <= 0 || mask.row_stride_bytes < mask.width) {
                fclose(file); return -4;
            }
            fprintf(file, ", \"maskBBox\": {\"x\": %d, \"y\": %d, \"width\": %d, \"height\": %d}",
                mask.origin_x, mask.origin_y, mask.width, mask.height);
            fprintf(file, ", \"mask\": {\"width\": %d, \"height\": %d, \"strideBytes\": %d, \"originX\": %d, \"originY\": %d}",
                mask.width, mask.height, mask.row_stride_bytes, mask.origin_x, mask.origin_y);
            fprintf(file, ", \"segGeometry\": {\"type\": \"%s\", \"centerX\": %.3f, \"centerY\": %.3f",
                geometry.type == YOLO_SEG_GEOMETRY_MIN_ENCLOSING_CIRCLE
                    ? "minimumEnclosingCircle" : "minimumAreaRectangle",
                geometry.center_x, geometry.center_y);
            if (geometry.type == YOLO_SEG_GEOMETRY_MIN_ENCLOSING_CIRCLE)
                fprintf(file, ", \"radius\": %.3f", geometry.radius);
            else
                fprintf(file, ", \"width\": %.3f, \"height\": %.3f, \"angle\": %.6f",
                    geometry.width, geometry.height, geometry.angle_radians);
            fputc('}', file);
        }
        if (task == YOLO_TASK_POSE) {
            uint32_t keypoint_count = 0, copied = 0;
            if (yolo_result_set_get_pose_keypoint_count(results, index, &keypoint_count) != YOLO_STATUS_OK) {
                fclose(file); return -5;
            }
            yolo_keypoint_v1_t* keypoints = keypoint_count
                ? (yolo_keypoint_v1_t*)malloc(sizeof(*keypoints) * keypoint_count) : NULL;
            if (keypoint_count &&
                yolo_result_set_copy_pose_keypoints(
                    results, index, keypoints, keypoint_count, &copied) != YOLO_STATUS_OK) {
                free(keypoints); fclose(file); return -5;
            }
            fputs(", \"keypoints\": [", file);
            for (uint32_t point = 0; point < copied; ++point)
                fprintf(file, "%s{\"x\": %.3f, \"y\": %.3f, \"score\": %.6f}",
                    point ? "," : "", keypoints[point].x, keypoints[point].y, keypoints[point].score);
            fputc(']', file);
            free(keypoints);
        }
        fputc('}', file);
    }
    fprintf(file, "%s  ]\n}\n", count ? "\n" : "");
    fclose(file);
    return 0;
}