#include <stdlib.h>
#include <stdio.h>
#include <mpi.h>
#include <execinfo.h>
#include <signal.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include <string.h>

// 100MB is equal to 100 * 1024 * 1024 bytes, so you will need to allocate (100 * 1024 * 1024) / 4 = 25,000,000 integers. <-- chatGPT ftw
#define ARRAY_SIZE 25000000
#define DEBUG 0

void segfault_handler(int sig) {
  void *array[10];
  size_t size;

  size = backtrace(array, 10);

  fprintf(stderr, "[-] Error: signal %d:\n", sig);
  backtrace_symbols_fd(array, size, STDERR_FILENO);
  exit(1);
}

int* setup_send_array(int size) {
	int *arr = (int *)malloc(size * sizeof(int));

	if(arr == NULL) {
		printf("[-] Error: malloc failed!\n");
		exit(1);
	}

	for (int i = 0; i < size; i++) {
        if (i % 2 == 0 ){
            arr[i] = 1;
        } else {
            arr[i] = 0;
        }
	}
    // Set random bit to two
    arr[rand() % size] = 2;
	return arr;
}

int* setup_recv_array(int n_elements) {
	int *recv_arr = (int *)malloc(n_elements * sizeof(int));

	if(recv_arr == NULL) {
		printf("[-] Error: malloc failed!\n");
		exit(1);
	}

	for (int i = 0; i < n_elements; i++) {
		recv_arr[i] = 0;
	}
	return recv_arr;
}

void setup_displacements_and_counts(int *displacements, int *counts, int ntasks) {
	for(int i = 0; i < ntasks; i++) {
		displacements[i] = (int) (ARRAY_SIZE / ntasks * i); // my_start
		counts[i] = (int) (ARRAY_SIZE / ntasks * (i + 1)) - displacements[i]; // my_end - my_start
	}
}

void throw_err(int line, int err_code) {
    char err_buf[256];

    sprintf(err_buf, "[-] Error: %s -> %d", __FILE__, line);
    perror(err_buf);
    exit(err_code);
}

int main(int argc, char* argv[])
{
	int taskid, ntasks;
	int *send_arr, *recv_arr;
	int *displacements, *counts;

	double start_time = MPI_Wtime(); 

	MPI_Init(&argc, &argv);
	signal(SIGSEGV, segfault_handler);

	MPI_Comm_rank(MPI_COMM_WORLD, &taskid);
	MPI_Comm_size(MPI_COMM_WORLD, &ntasks);
    MPI_Request request;
    MPI_Status status;
    int flag = 0;
    char *message = "Hello World!\0";
    int message_len = strlen(message) + 1;
    char recv_message[message_len];

	if(DEBUG) {
		printf("[+] Task %d has started...\n", taskid);
		fflush(stdout);
	}

	if (taskid == 0) {
		//Setup array that we need to sort
		send_arr = setup_send_array(ARRAY_SIZE);
	}

	// Allocate memory for send_displacements and send_counts
	displacements = (int *)malloc(ntasks * sizeof(int));
	counts = (int *)malloc(ntasks * sizeof(int));

	if(displacements == NULL || counts == NULL) {
		printf("[-] Error: malloc failed!\n");
		exit(1);
	}

	// Setup displacements and counts
	setup_displacements_and_counts(displacements, counts, ntasks);

	if(DEBUG) {
		printf("[+] Setup displacements and counts done!\n");
		fflush(stdout);
	}

	// Setup recv array
	recv_arr = setup_recv_array(counts[taskid]);

	if(DEBUG) {
		printf("[+] Setup done!\n");
		fflush(stdout);
	}

	if(MPI_Barrier(MPI_COMM_WORLD) != MPI_SUCCESS)
		throw_err(__LINE__, errno);

	// Scatterv array to all processes
	if(MPI_Scatterv(send_arr, counts, displacements , MPI_INT, recv_arr, ARRAY_SIZE, MPI_INT, 0, MPI_COMM_WORLD) != MPI_SUCCESS)
		throw_err(__LINE__, errno);

    MPI_Irecv(recv_message, message_len, MPI_CHAR, 0, 0, MPI_COMM_WORLD, &request);


    for(int i =0; i < counts[taskid]; i++) {
        if(DEBUG) 
            printf("%d ", recv_arr[i]);
        if(i % 20 == 0) {
            MPI_Test(&request, &flag, &status);
            if(flag) {
                printf("[+] Recieved stop signal!\n");
                break;
            }
        }
        if(recv_arr[i] == 2) {
            printf("[+] Found 2 at index %d! --> taskid: %d\n", displacements[taskid] + i, taskid);
            MPI_Isend(message, message_len, MPI_CHAR, 0, 0, MPI_COMM_WORLD, &request);
            break;
        }
    }

	// Gather the array
	if(MPI_Gatherv(recv_arr, counts[taskid], MPI_INT, send_arr, counts, displacements, MPI_INT, 0, MPI_COMM_WORLD))
		throw_err(__LINE__, errno);


	if (taskid == 0) {
		if(DEBUG) {
		  	for(int i = 0; i < ARRAY_SIZE; i++) {
				printf("%d ", send_arr[i]);
			}
		}

		double total_time = MPI_Wtime() - start_time;

		printf("\nTotal time: %f\n", total_time);
	}

	MPI_Finalize();
	return 0;

    // 1:  0.616473
    // 2:  0.742917
    // 4:  0.804409
    // 8:  0.782405
    // 16: 6.045285
    // 32:  3.036735

}