#include "types.h"
#include "stat.h"
#include "user.h"

#define MAX_CHILDREN 16
        
int test_simple_crash(void (*test_func)(), const char *crash_message, const char *no_crash_message) {
    int fds[2];
    pipe(fds);
    int pid = fork();
    if (pid == 0) {
        /* child process */
        close(1);
        dup(fds[1]);
        test_func();
        write(1, "X", 1);
        exit();
    } else {
        char text[1];
        close(fds[1]);
        int size = read(fds[0], text, 1);
        wait();
        close(fds[0]);
        if (size == 1) {
            printf(1, "%s\n", no_crash_message);
            return 0;
        } else {
            printf(1, "%s\n", crash_message);
            return 1;
        }
    }
}

static unsigned out_of_bounds_offset = 1;
void test_out_of_bounds_internal() {
    volatile char *end_of_heap = sbrk(0);
    (void) end_of_heap[out_of_bounds_offset];
}

int test_out_of_bounds(int offset, const char *crash_message, const char *no_crash_message) {
    out_of_bounds_offset = offset;
    return test_simple_crash(test_out_of_bounds_internal, crash_message, no_crash_message);
}

int test_allocation(int size, const char *describe_size, const char *describe_amount, int offset1, int count1, int offset2, int count2) {
    int fds[2];
    pipe(fds);
    int pid = fork();
    if (pid == 0) {
        /* child process */
        char *old_end_of_heap = sbrk(size);
        char *new_end_of_heap = sbrk(0);
        if (old_end_of_heap == (char*) -1) {
            write(fds[1], "NA", 2);
        } else if (new_end_of_heap - old_end_of_heap != size) {
            write(fds[1], "NS", 2);
        } else {
            int failed = 0;
            char *place_one = &new_end_of_heap[-offset1 - count1];
            char *place_two = &new_end_of_heap[-offset2 - count2];
            int i;
            for (i = 0; i < count1; ++i) {
                place_one[i] = 'A';
            }
            for (i = 0; i < count2; ++i) {
                place_two[i] = 'B';
            }
            for (i = 0; i < count1; ++i) {
                if (place_one[i] != 'A')
                    failed = 1;
            } 
            for (i = 0; i < count2; ++i) {
                if (place_two[i] != 'B')
                    failed = 1;
            }
            write(fds[1], failed ? "NR" : "YY", 2);
        }
        exit();
    } else if (pid > 0) {
        char text[10];
        close(fds[1]);
        int size = read(fds[0], text, 10);
        wait();
        close(fds[0]);
        if (size == 2 && text[0] == 'N') {
            if (text[1] == 'A') {
                printf(1, "allocating (but not using) %s with sbrk() returned error\n", describe_size);
            } else if (text[1] == 'R') {
                printf(1, "using %s parts of %s allocation read wrong value\n", describe_amount, describe_size);
            } else if (text[1] == 'S') {
                printf(1, "wrong size allocated by %s allocation\n", describe_size);
            } else {
                printf(1, "unknown error using %s parts of %s allocation\n", describe_amount, describe_size);
            }
            return 0;
        } else if (size == 0) {
            printf(1, "allocating %s and using %s parts of allocation crashed\n", describe_size, describe_amount);
            return 0;
        } else if (size >= 1 && text[0] == 'Y') {
            printf(1, "allocating %s and using %s parts of allocation passed\n", describe_size, describe_amount );
            return 1;
        } else {
            printf(1, "unknown error\n");
            return 0;
        }
    } else {
        printf(1, "allocation test: first fork failed\n");
        return 0;
    }
}

void wait_forever() {
  while (1) { sleep(1000); }
}

void test_copy_on_write_main_child(int result_fd, int size, const char *describe_size, int forks) {
  char *old_end_of_heap = sbrk(size);
  char *new_end_of_heap = sbrk(0);
  for (char *p = old_end_of_heap; p < new_end_of_heap; ++p) {
      *p = 'A';
  }
  int children[MAX_CHILDREN] = {0};
  if (forks > MAX_CHILDREN) {
    printf(2, "unsupported number of children in test_copy_on_write\n");
  }
  int failed = 0;
  char failed_code = ' ';
  for (int i = 0; i < forks; ++i) {
    int child_fds[2];
    pipe(child_fds);
    children[i] = fork();
    if (children[i] == -1) {
      printf(2, "fork failed\n");
      failed = 1;
      failed_code = 'f';
      break;
    } else if (children[i] == 0) {
      int found_wrong_memory = 0;
      for (char *p = old_end_of_heap; p < new_end_of_heap; ++p) {
        if (*p != 'A') {
          found_wrong_memory = 1;
        }
      }
      old_end_of_heap[size / 2] = 'B';
      old_end_of_heap[4096 * i] = 'C';
      if (old_end_of_heap[4096 * i + 1] != 'A' || old_end_of_heap[size/2+1] != 'A') {
          found_wrong_memory = 1;
      }
      write(child_fds[1], found_wrong_memory ? "-" : "+", 1);
      wait_forever();
    } else {
      char buffer[1] = {'X'};
      read(child_fds[0], buffer, 1);
      if (buffer[0] != '+') {
        failed = 1;
        failed_code = 'c';
      }
      close(child_fds[0]); close(child_fds[1]);
    }
  }
  for (int i = 0; i < forks; ++i) {
    kill(children[i]);
    wait();
  }
  for (char *p = old_end_of_heap; p < new_end_of_heap; ++p) {
    if (*p != 'A') {
      failed = 1;
      failed_code = 'p';
    }
  }
  if (failed) {
    char buffer[2] = {'N', ' '};
    buffer[1] = failed_code;
    write(result_fd, buffer, 2);
  } else {
    write(result_fd, "YY", 2);
  }
  exit();
}

int test_copy_on_write(int size,  const char *describe_size, int forks) {
  int fds[2];
  pipe(fds);
  int pid = fork();
  if (pid == 0) {
    test_copy_on_write_main_child(fds[1], size, describe_size, forks);
  } else if (pid > 0) {
    char text[10] = {'X', 'X'};
    close(fds[1]);
    read(fds[0], text, 10);
    wait();
    close(fds[0]);
    if (text[0] == 'X') {
      printf(1, "copy on write test failed --- crash?\n");
      return 0;
    } else if (text[0] == 'N') {
      switch (text[1]) {
      case 'f':
        printf(1, "copy on write test failed --- fork failed\n");
        break;
      case 'p':
        printf(1, "copy on write test failed --- wrong value for memory in parent\n");
        break;
      case 'c':
        printf(1, "copy on write test failed --- wrong value for memory in child\n");
        break;
      default:
        printf(1, "copy on write test failed --- unknown reason\n");
        break;
      }
      return 0;
    } else {
      printf(1, "copy on write test passed --- allocate %s; "
             "fork %d children; read+write small parts in each child\n",
             describe_size, forks);
      return 1;
    }
  } else if (pid == -1) {
     printf(1, "copy on write test failed --- first fork failed\n");
  }
  return 0;
}

int
main(int argc, char *argv[])
{
    int passed = 0;
    passed += test_out_of_bounds(1, 
        "reading immediately after heap end crashes",
        "reading immediately after heap end does NOT crash");
    passed += test_out_of_bounds(1024 * 1024 * 1024 * 2U, 
        "reading 2GB after heap end crashes",
        "reading 2GB after heap end does NOT crash");
    passed += test_allocation(32 * 1024, "32KB", "two 4KB", 4 * 1024,  4096, 16 * 1024, 4096); 
    passed += test_allocation(768 * 1024 * 1024, "768MB", "two 8KB", 0, 8192, 256 * 1024 * 1024, 8192); 
    passed += test_allocation(768 * 1024 * 1024, "768MB", "two 32KB", 128 * 1024 * 1024, 32 * 1024, 384 * 1024 * 1024, 32 * 1024); 
    passed += test_copy_on_write(32 * 1024, "32KB", 2); 
    passed += test_copy_on_write(64 * 1024 * 1024, "64MB", 7); 
    printf(1, "running repeated copy-on-write tests to make sure "
              "memory is eventually reclaimed\n");
    for (int i = 0; i < 10; ++i) {
        passed += test_copy_on_write(100 * 1024 * 1024, "100MB", 2); 
    }
    printf(1, "passed %d tests of 17\n", passed);
    exit();
}