#include "test_common.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/time.h>

typedef struct {
    double time_ms;
    double pixels_per_sec;
    double megabytes_per_sec;
    size_t input_bytes;
    size_t output_bytes;
} BenchResult;

static double get_time_us(void) {
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return tv.tv_sec * 1000000.0 + tv.tv_usec;
}

static void calculate_bench_metrics(const char *name, int iterations, BenchResult *r) {
    if (strstr(name, "create_uniform") != NULL) {
        r->input_bytes = 0;
        r->output_bytes = 16 * 16 * 4;
    } else if (strstr(name, "create_gradient_256") != NULL) {
        r->input_bytes = 0;
        r->output_bytes = 256 * 256 * 4;
    } else if (strstr(name, "create_noise_1024") != NULL) {
        r->input_bytes = 0;
        r->output_bytes = 1024 * 1024 * 4;
    } else if (strstr(name, "downsample_uniform_16x16") != NULL) {
        r->input_bytes = 16 * 16 * 4;
        r->output_bytes = 8 * 8 * 4;
    } else if (strstr(name, "downsample_gradient_64x64") != NULL) {
        r->input_bytes = 64 * 64 * 4;
        r->output_bytes = 32 * 32 * 4;
    } else if (strstr(name, "downsample_gradient_256x256") != NULL) {
        r->input_bytes = 256 * 256 * 4;
        r->output_bytes = 128 * 128 * 4;
    } else if (strstr(name, "downsample_gradient_1024x1024") != NULL) {
        r->input_bytes = 1024 * 1024 * 4;
        r->output_bytes = 512 * 512 * 4;
    } else if (strstr(name, "downsample_noise_512x512") != NULL) {
        r->input_bytes = 512 * 512 * 4;
        r->output_bytes = 128 * 128 * 4;
    } else if (strstr(name, "downsample_noise_1024x1024") != NULL) {
        r->input_bytes = 1024 * 1024 * 4;
        r->output_bytes = 256 * 256 * 4;
    } else if (strstr(name, "downsample_noise_1920x1080") != NULL) {
        r->input_bytes = 1920 * 1080 * 4;
        r->output_bytes = 960 * 540 * 4;
    } else if (strstr(name, "downsample_noise_3840x2160") != NULL) {
        r->input_bytes = 3840 * 2160 * 4;
        r->output_bytes = 1920 * 1080 * 4;
    } else if (strstr(name, "downsample_noise_4096x4096") != NULL) {
        r->input_bytes = 4096 * 4096 * 4;
        r->output_bytes = 1024 * 1024 * 4;
    } else if (strstr(name, "downsample_checkerboard_100x100") != NULL) {
        r->input_bytes = 100 * 100 * 4;
        r->output_bytes = 50 * 50 * 4;
    } else if (strstr(name, "downsample_checkerboard_256x256") != NULL) {
        r->input_bytes = 256 * 256 * 4;
        r->output_bytes = 128 * 128 * 4;
    } else {
        r->input_bytes = 0;
        r->output_bytes = 0;
    }
    r->input_bytes *= (size_t)iterations;
    r->output_bytes *= (size_t)iterations;
}

static void run_bench(double (*fn)(void), int iterations, double *elapsed_ms) {
    double start = get_time_us();
    for (int i = 0; i < iterations; i++) {
        fn();
    }
    *elapsed_ms = (get_time_us() - start) / 1000.0;
}

static double bench_create_uniform_16x16(void) {
    for (int i = 0; i < 5000; i++) {
        uint8_t *img = create_uniform_image(16, 16, 128, 128, 128);
        free(img);
    }
    return 0;
}

static double bench_create_gradient_256x256(void) {
    for (int i = 0; i < 500; i++) {
        uint8_t *img = create_gradient_image(256, 256);
        free(img);
    }
    return 0;
}

static double bench_create_noise_1024x1024(void) {
    for (int i = 0; i < 50; i++) {
        uint8_t *img = create_noise_image(1024, 1024, 42);
        free(img);
    }
    return 0;
}

static double bench_downsample_uniform_16x16(void) {
    uint8_t *src = create_uniform_image(16, 16, 128, 128, 128);
    int dw, dh;

    for (int i = 0; i < 2000; i++) {
        uint8_t *dst = downsample_image(src, 16, 16, 4, &dw, &dh, 0.5f);
        free(dst);
    }

    free(src);
    return 0;
}

static double bench_downsample_gradient_64x64(void) {
    uint8_t *src = create_gradient_image(64, 64);
    int dw, dh;

    for (int i = 0; i < 500; i++) {
        uint8_t *dst = downsample_image(src, 64, 64, 4, &dw, &dh, 0.5f);
        free(dst);
    }

    free(src);
    return 0;
}

static double bench_downsample_gradient_256x256(void) {
    uint8_t *src = create_gradient_image(256, 256);
    int dw, dh;

    for (int i = 0; i < 200; i++) {
        uint8_t *dst = downsample_image(src, 256, 256, 4, &dw, &dh, 0.5f);
        free(dst);
    }

    free(src);
    return 0;
}

static double bench_downsample_gradient_1024x1024(void) {
    uint8_t *src = create_gradient_image(1024, 1024);
    int dw, dh;

    for (int i = 0; i < 20; i++) {
        uint8_t *dst = downsample_image(src, 1024, 1024, 4, &dw, &dh, 0.5f);
        free(dst);
    }

    free(src);
    return 0;
}

static double bench_downsample_noise_512x512(void) {
    uint8_t *src = create_noise_image(512, 512, 42);
    int dw, dh;

    for (int i = 0; i < 50; i++) {
        uint8_t *dst = downsample_image(src, 512, 512, 4, &dw, &dh, 0.25f);
        free(dst);
    }

    free(src);
    return 0;
}

static double bench_downsample_noise_1024x1024(void) {
    uint8_t *src = create_noise_image(1024, 1024, 123);
    int dw, dh;

    for (int i = 0; i < 15; i++) {
        uint8_t *dst = downsample_image(src, 1024, 1024, 4, &dw, &dh, 0.25f);
        free(dst);
    }

    free(src);
    return 0;
}

static double bench_downsample_noise_1920x1080(void) {
    uint8_t *src = create_noise_image(1920, 1080, 456);
    int dw, dh;

    for (int i = 0; i < 5; i++) {
        uint8_t *dst = downsample_image(src, 1920, 1080, 4, &dw, &dh, 0.5f);
        free(dst);
    }

    free(src);
    return 0;
}

static double bench_downsample_noise_3840x2160(void) {
    uint8_t *src = create_noise_image(3840, 2160, 789);
    int dw, dh;

    for (int i = 0; i < 2; i++) {
        uint8_t *dst = downsample_image(src, 3840, 2160, 4, &dw, &dh, 0.5f);
        free(dst);
    }

    free(src);
    return 0;
}

static double bench_downsample_noise_4096x4096(void) {
    uint8_t *src = create_noise_image(4096, 4096, 456);
    int dw, dh;

    for (int i = 0; i < 2; i++) {
        uint8_t *dst = downsample_image(src, 4096, 4096, 4, &dw, &dh, 0.25f);
        free(dst);
    }

    free(src);
    return 0;
}

static double bench_downsample_checkerboard_100x100(void) {
    uint8_t *src = create_checkerboard(100, 100, 10);
    int dw, dh;

    for (int i = 0; i < 200; i++) {
        uint8_t *dst = downsample_image(src, 100, 100, 4, &dw, &dh, 0.5f);
        free(dst);
    }

    free(src);
    return 0;
}

static double bench_downsample_checkerboard_256x256(void) {
    uint8_t *src = create_checkerboard(256, 256, 16);
    int dw, dh;

    for (int i = 0; i < 100; i++) {
        uint8_t *dst = downsample_image(src, 256, 256, 4, &dw, &dh, 0.5f);
        free(dst);
    }

    free(src);
    return 0;
}

typedef struct {
    const char *name;
    double (*fn)(void);
} BenchDef;

static BenchDef benchmarks[] = {
    {"create_uniform_16x16", bench_create_uniform_16x16},
    {"create_gradient_256x256", bench_create_gradient_256x256},
    {"create_noise_1024x1024", bench_create_noise_1024x1024},
    {"downsample_uniform_16x16_0.5x", bench_downsample_uniform_16x16},
    {"downsample_gradient_64x64_0.5x", bench_downsample_gradient_64x64},
    {"downsample_gradient_256x256_0.5x", bench_downsample_gradient_256x256},
    {"downsample_gradient_1024x1024_0.5x", bench_downsample_gradient_1024x1024},
    {"downsample_noise_512x512_0.25x", bench_downsample_noise_512x512},
    {"downsample_noise_1024x1024_0.25x", bench_downsample_noise_1024x1024},
    {"downsample_noise_1920x1080_0.5x", bench_downsample_noise_1920x1080},
    {"downsample_noise_3840x2160_0.5x", bench_downsample_noise_3840x2160},
    {"downsample_noise_4096x4096_0.25x", bench_downsample_noise_4096x4096},
    {"downsample_checkerboard_100x100_0.5x", bench_downsample_checkerboard_100x100},
    {"downsample_checkerboard_256x256_0.5x", bench_downsample_checkerboard_256x256},
};

static int benchmark_iterations[] = {
    5000, 500, 50, 2000, 500, 200, 20, 50, 15, 5, 2, 2, 200, 100
};

int main(int argc, char **argv) {
    int csv_output = 0;

    for (int i = 1; i < argc; i++) {
        if (strcmp(argv[i], "--csv") == 0) {
            csv_output = 1;
        }
    }

    if (csv_output) {
        printf("name,time_ms,pixels_per_sec,megabytes_per_sec,iterations\n");
    } else {
        printf("Chroma Performance Benchmarks\n");
        printf("=============================\n\n");
        printf("  %-42s %16s %22s %19s\n", "Benchmark", "Time (ms)", "Pixels/sec", "MB/sec");
        printf("  %-42s %16s %22s %19s\n", "-----------------------------------------", "--------------", "-----------------", "--------------");
    }

    int num_benchmarks = sizeof(benchmarks) / sizeof(benchmarks[0]);
    int max_name_len = 0;
    for (int i = 0; i < num_benchmarks; i++) {
        int len = strlen(benchmarks[i].name);
        if (len > max_name_len) max_name_len = len;
    }

    for (int i = 0; i < num_benchmarks; i++) {
        BenchResult result = {0};
        calculate_bench_metrics(benchmarks[i].name, benchmark_iterations[i], &result);

        double elapsed_ms;
        run_bench(benchmarks[i].fn, benchmark_iterations[i], &elapsed_ms);

        size_t total_pixels = result.input_bytes > 0 ? result.input_bytes / 4 : 0;
        result.time_ms = elapsed_ms;
        if (total_pixels > 0 && elapsed_ms > 0) {
            result.pixels_per_sec = total_pixels / (elapsed_ms / 1000.0);
        }
        double total_mb = (result.input_bytes + result.output_bytes) / (1024.0 * 1024.0);
        if (total_mb > 0 && elapsed_ms > 0) {
            result.megabytes_per_sec = total_mb / (elapsed_ms / 1000.0);
        }

        if (csv_output) {
            printf("%s,%.3f,%.0f,%.2f,%d\n",
                   benchmarks[i].name, result.time_ms, result.pixels_per_sec,
                   result.megabytes_per_sec, benchmark_iterations[i]);
        } else {
            printf("  %-42s %16.3f %22.0f %19.2f\n",
                   benchmarks[i].name, result.time_ms,
                   result.pixels_per_sec, result.megabytes_per_sec);
        }
    }

    if (!csv_output) {
        printf("\n");
    }

    (void)argc;
    return 0;
}