#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/wait.h>
#include <sys/mman.h>

#include "images/images.h"
#include "complex.c"


#define MB_THREADS 7

void draw_mandelbrot(int width, int height, int iterations);
void draw_mandelbrot_auto(int width, int iterations);
int test_mandelbrot(Complex c,  int iterations);
void* create_shared_memory(size_t size);
int image_new_shared(int width, int height, Image* img);
void image_destroy_shared(Image img);


int main(int argc, char* argv[]) {
    int width = 3000;
    int iterations = 60;
    if (argc > 1) {
        width = atoi(argv[1]);
    }
    if (argc > 2) {
        iterations = atoi(argv[2]);
    }

    draw_mandelbrot_auto(width, iterations);
}

void draw_mandelbrot_auto(int width, int iterations) {
    draw_mandelbrot(width, 2.5f * width / 3, iterations);
}

void draw_mandelbrot(int width, int height, int iterations) {
    int center_x = 2 * width / 3;
    int center_y = height / 2;
    float step = 3.0f / width;

    Image img;
    image_new_shared(width, height, &img);

    int thread_id = 0;

    for (int i = 0; i < MB_THREADS - 1; i++) {
        if (fork() == 0) {
            thread_id = i + 1;
            break;
        }
    }
    printf("Thread %i reporting for duty\n", thread_id);

    for (int x = 0; x < width; x++) {
        if (x % MB_THREADS != thread_id) continue;

        for (int y = 0; y < height; y++) {
            Complex c = complex_new(
                (x - center_x) * step,
                (y - center_y) * step
            );
            int ret = test_mandelbrot(c, iterations);
            //int r = (ret * 160 / iterations);
            //int g = (ret * 9 / iterations);
            //int b = (ret * 165 / iterations);
            //image_set_px(img, x, y, r,g,b);
            int color = 255 * ret / (float) iterations;
            image_set_px(img, x, y, color, color, color);
        }
    }

    printf("Thread %i is finished\n", thread_id);
    if (thread_id != 0) {
        exit(0);
    }

    int status;
    for (int i = 0; i < MB_THREADS - 1; i++) {
        printf("Waiting for threads... %d/%d\n", i, MB_THREADS - 1);
        while(wait(&status) > 0) {}
    }

    printf("Got %d/%d threads\n", MB_THREADS - 1, MB_THREADS - 1);

    image_save_bmp(img, "mandelbrot.bmp");
    image_destroy_shared(img);
}

int test_mandelbrot(Complex c,  int iterations) {
    Complex z = complex_new(0,0);

    for (int i = 0; i < iterations; i++) {
        z = complex_add(complex_mult(z,z), c); // z^2 + c
        if (complex_abs(z) > 2) return i;
    }
    return 0;
}



int image_new_shared(int width, int height, Image* img) {
    if (img == NULL) return 0;

    img->height = height;
    img->width = width;

    if ((unsigned long int) width * height * 3 > (unsigned long int) UINT_MAX) {
        printf("Image dimensions to large for a bitmap!\n");
        return 0;
    }

    // initialize bitmap...
    img->bitmap = create_shared_memory(height * sizeof(char*));
    for (int i = 0; i < height; i++) {
        img->bitmap[i] = create_shared_memory(3 * width * sizeof(char));
    }

    return 1;
}

void image_destroy_shared(Image img) {
    for (int i = 0; i < img.height; i++) {
        munmap(img.bitmap[i], 3 * img.width * sizeof(char));
    }  
    munmap(img.bitmap, img.height * sizeof(char*));
}

void* create_shared_memory(size_t size) {
  // Our memory buffer will be readable and writable:
  int protection = PROT_READ | PROT_WRITE;

  // The buffer will be shared (meaning other processes can access it), but
  // anonymous (meaning third-party processes cannot obtain an address for it),
  // so only this process and its children will be able to use it:
  int visibility = MAP_SHARED | MAP_ANONYMOUS;

  // The remaining parameters to `mmap()` are not important for this use case,
  // but the manpage for `mmap` explains their purpose.
  return mmap(NULL, size, protection, visibility, -1, 0);
}