#include "../../unity/unity.h"
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <fcntl.h>
#include <errno.h>

/* usage(int) is defined in the program source included before this file. */

void setUp(void) {
  /* Setup code here, or leave empty */
}

void tearDown(void) {
  /* Cleanup code here, or leave empty */
}

static int read_all_fd(int fd, char **out_buf, size_t *out_len) {
  size_t cap = 1024;
  size_t len = 0;
  char *buf = (char *)malloc(cap);
  if (!buf) return -1;

  for (;;) {
    if (len == cap) {
      size_t new_cap = cap * 2;
      char *new_buf = (char *)realloc(buf, new_cap);
      if (!new_buf) {
        free(buf);
        return -1;
      }
      buf = new_buf;
      cap = new_cap;
    }
    ssize_t n = read(fd, buf + len, cap - len);
    if (n < 0) {
      if (errno == EINTR) continue;
      free(buf);
      return -1;
    } else if (n == 0) {
      break;
    } else {
      len += (size_t)n;
    }
  }

  /* NUL-terminate for convenience */
  if (len == cap) {
    char *new_buf = (char *)realloc(buf, cap + 1);
    if (!new_buf) {
      free(buf);
      return -1;
    }
    buf = new_buf;
  }
  buf[len] = '\0';

  *out_buf = buf;
  *out_len = len;
  return 0;
}

static int spawn_usage_and_capture(int status_to_pass,
                                   char **stdout_buf, size_t *stdout_len,
                                   char **stderr_buf, size_t *stderr_len,
                                   int *exit_code_out) {
  int out_pipe[2];
  int err_pipe[2];

  if (pipe(out_pipe) != 0) return -1;
  if (pipe(err_pipe) != 0) {
    close(out_pipe[0]); close(out_pipe[1]);
    return -1;
  }

  fflush(stdout);
    fflush(stderr);
    pid_t pid = fork();
  if (pid < 0) {
    close(out_pipe[0]); close(out_pipe[1]);
    close(err_pipe[0]); close(err_pipe[1]);
    return -1;
  }

  if (pid == 0) {
    /* Child: redirect stdout/stderr to pipes and call usage(). */
    /* Close read ends in child */
    close(out_pipe[0]);
    close(err_pipe[0]);

    /* Duplicate write ends onto STDOUT/ERR */
    if (dup2(out_pipe[1], STDOUT_FILENO) < 0) _exit(127);
    if (dup2(err_pipe[1], STDERR_FILENO) < 0) _exit(127);

    /* Close original write ends after dup */
    close(out_pipe[1]);
    close(err_pipe[1]);

    usage(status_to_pass);

    /* Should not reach here; if it does, return special code */
    _exit(255);
  }

  /* Parent: close write ends, read outputs, wait for child */
  close(out_pipe[1]);
  close(err_pipe[1]);

  char *outb = NULL, *errb = NULL;
  size_t outl = 0, errl = 0;

  int read_ok1 = read_all_fd(out_pipe[0], &outb, &outl);
  int read_ok2 = read_all_fd(err_pipe[0], &errb, &errl);

  close(out_pipe[0]);
  close(err_pipe[0]);

  int status;
  if (waitpid(pid, &status, 0) < 0) {
    if (outb) free(outb);
    if (errb) free(errb);
    return -1;
  }

  if (read_ok1 != 0 || read_ok2 != 0) {
    if (outb) free(outb);
    if (errb) free(errb);
    return -1;
  }

  int ec = -1;
  if (WIFEXITED(status)) {
    ec = WEXITSTATUS(status);
  } else {
    ec = -1;
  }

  *stdout_buf = outb;
  *stdout_len = outl;
  *stderr_buf = errb;
  *stderr_len = errl;
  *exit_code_out = ec;

  return 0;
}

static void free_capture(char *outb, char *errb) {
  free(outb);
  free(errb);
}

void test_usage_exit_success_prints_help_to_stdout_and_exit_zero(void) {
  char *outb = NULL, *errb = NULL;
  size_t outl = 0, errl = 0;
  int ec = -1;

  int rc = spawn_usage_and_capture(EXIT_SUCCESS, &outb, &outl, &errb, &errl, &ec);
  TEST_ASSERT_EQUAL_INT(0, rc);
  TEST_ASSERT_EQUAL_INT(0, ec);
  TEST_ASSERT_NOT_NULL(outb);
  TEST_ASSERT_TRUE_MESSAGE(outl > 0, "Expected non-empty stdout for successful usage()");
  TEST_ASSERT_NOT_NULL(errb);
  TEST_ASSERT_EQUAL_UINT64_MESSAGE(0, errl, "Expected empty stderr for successful usage()");

  free_capture(outb, errb);
}

void test_usage_exit_failure_emits_try_help_to_stderr_and_exit_failure(void) {
  char *outb = NULL, *errb = NULL;
  size_t outl = 0, errl = 0;
  int ec = -1;

  int rc = spawn_usage_and_capture(EXIT_FAILURE, &outb, &outl, &errb, &errl, &ec);
  TEST_ASSERT_EQUAL_INT(0, rc);
  TEST_ASSERT_EQUAL_INT(EXIT_FAILURE, ec);
  TEST_ASSERT_NOT_NULL(outb);
  TEST_ASSERT_EQUAL_UINT64_MESSAGE(0, outl, "Expected empty stdout for failure usage()");
  TEST_ASSERT_NOT_NULL(errb);
  TEST_ASSERT_TRUE_MESSAGE(errl > 0, "Expected non-empty stderr for failure usage()");

  free_capture(outb, errb);
}

void test_usage_nonzero_status_is_propagated_and_outputs_on_stderr(void) {
  int custom_status = 42;
  char *outb = NULL, *errb = NULL;
  size_t outl = 0, errl = 0;
  int ec = -1;

  int rc = spawn_usage_and_capture(custom_status, &outb, &outl, &errb, &errl, &ec);
  TEST_ASSERT_EQUAL_INT(0, rc);
  TEST_ASSERT_EQUAL_INT(custom_status, ec);
  TEST_ASSERT_NOT_NULL(outb);
  TEST_ASSERT_EQUAL_UINT64_MESSAGE(0, outl, "Expected empty stdout for nonzero status usage()");
  TEST_ASSERT_NOT_NULL(errb);
  TEST_ASSERT_TRUE_MESSAGE(errl > 0, "Expected non-empty stderr for nonzero status usage()");

  free_capture(outb, errb);
}

int main(void) {
  UNITY_BEGIN();

  RUN_TEST(test_usage_exit_success_prints_help_to_stdout_and_exit_zero);
  RUN_TEST(test_usage_exit_failure_emits_try_help_to_stderr_and_exit_failure);
  RUN_TEST(test_usage_nonzero_status_is_propagated_and_outputs_on_stderr);

  return UNITY_END();
}