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

/* usage(int) is defined in the program under test. We just use it here. */
extern void usage (int status);

/* set_program_name is provided by gnulib (included via system.h in the program). */
extern void set_program_name (const char *argv0);

/* Helper structure to collect results from a child invocation of usage(). */
typedef struct RunResult {
  int exit_status;      /* -1 if not exited normally */
  char *out;            /* NUL-terminated, may be empty string */
  size_t out_len;
  char *err;            /* NUL-terminated, may be empty string */
  size_t err_len;
} RunResult;

static char* read_all_from_fd (int fd, size_t *len_out) {
  size_t cap = 4096;
  size_t len = 0;
  char *buf = (char*) malloc (cap);
  if (!buf) {
    /* On allocation failure, return empty string to avoid NULL deref in tests. */
    buf = (char*) malloc (1);
    if (!buf) _exit(255);
    buf[0] = '\0';
    if (len_out) *len_out = 0;
    return buf;
  }

  for (;;) {
    char tmp[2048];
    ssize_t n = read (fd, tmp, sizeof tmp);
    if (n < 0) {
      if (errno == EINTR) continue;
      break;
    }
    if (n == 0) break;
    if (len + (size_t) n + 1 > cap) {
      size_t new_cap = (len + (size_t) n + 1) * 2;
      char *nb = (char*) realloc (buf, new_cap);
      if (!nb) {
        free (buf);
        buf = (char*) malloc (1);
        if (!buf) _exit(255);
        buf[0] = '\0';
        if (len_out) *len_out = 0;
        return buf;
      }
      buf = nb;
      cap = new_cap;
    }
    memcpy (buf + len, tmp, (size_t) n);
    len += (size_t) n;
  }
  buf[len] = '\0';
  if (len_out) *len_out = len;
  return buf;
}

static void free_run_result (RunResult *r) {
  if (!r) return;
  free (r->out);
  free (r->err);
  r->out = NULL;
  r->err = NULL;
}

/* Run usage(status) in a child process, capturing stdout and stderr. */
static void run_usage_capture (int status, RunResult *res) {
  int out_pipe[2] = {-1, -1};
  int err_pipe[2] = {-1, -1};

  if (pipe (out_pipe) != 0) {
    TEST_FAIL_MESSAGE("Failed to create stdout pipe");
  }
  if (pipe (err_pipe) != 0) {
    close (out_pipe[0]); close (out_pipe[1]);
    TEST_FAIL_MESSAGE("Failed to create stderr pipe");
  }

  pid_t pid = fork ();
  if (pid < 0) {
    close (out_pipe[0]); close (out_pipe[1]);
    close (err_pipe[0]); close (err_pipe[1]);
    TEST_FAIL_MESSAGE("fork() failed");
  }

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

    /* Duplicate write ends to stdout/stderr. */
    if (dup2 (out_pipe[1], STDOUT_FILENO) < 0) _exit (253);
    if (dup2 (err_pipe[1], STDERR_FILENO) < 0) _exit (252);

    /* Close original write fds after dup2. */
    close (out_pipe[1]);
    close (err_pipe[1]);

    /* Call the function under test. This will exit(status). */
    usage (status);

    /* Should not reach here. */
    _exit (251);
  }

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

  res->out = read_all_from_fd (out_pipe[0], &res->out_len);
  res->err = read_all_from_fd (err_pipe[0], &res->err_len);

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

  int wstatus;
  for (;;) {
    if (waitpid (pid, &wstatus, 0) < 0) {
      if (errno == EINTR) continue;
      res->exit_status = -1;
      break;
    }
    if (WIFEXITED (wstatus)) {
      res->exit_status = WEXITSTATUS (wstatus);
    } else if (WIFSIGNALED (wstatus)) {
      res->exit_status = 128 + WTERMSIG (wstatus);
    } else {
      res->exit_status = -1;
    }
    break;
  }
}

void setUp(void) {
  /* Ensure program_name is initialized to a sensible value. */
  set_program_name ("nohup");
}

void tearDown(void) {
  /* Nothing to clean up. */
}

static void assert_buffer_empty (const char *label, const char *buf, size_t len) {
  /* Ensure both logical and actual length agree on emptiness. */
  TEST_ASSERT_NOT_NULL_MESSAGE(buf, "Captured buffer is NULL");
  if (len == 0) {
    TEST_ASSERT_EQUAL_UINT64_MESSAGE(0u, (unsigned long long)strlen(buf), label);
  } else {
    TEST_ASSERT_MESSAGE(len > 0, "Expected non-zero length");
  }
}

/* Test: usage(nonzero) exits with that status, writes nothing to stdout and something to stderr. */
void test_usage_failure_status_one(void) {
  RunResult r = {0};
  run_usage_capture (1, &r);

  TEST_ASSERT_EQUAL_INT_MESSAGE(1, r.exit_status, "usage(1) exit status mismatch");
  assert_buffer_empty ("stdout should be empty", r.out, r.out_len);
  TEST_ASSERT_TRUE_MESSAGE(r.err_len > 0, "stderr should be non-empty for failure");
  TEST_ASSERT_NOT_EQUAL_MESSAGE('\0', r.err[0], "stderr should contain data");

  free_run_result (&r);
}

/* Test: usage(127) (POSIX_NOHUP_FAILURE) behaves similarly: no stdout, some stderr, exits 127. */
void test_usage_failure_status_127(void) {
  RunResult r = {0};
  run_usage_capture (127, &r);

  TEST_ASSERT_EQUAL_INT_MESSAGE(127, r.exit_status, "usage(127) exit status mismatch");
  assert_buffer_empty ("stdout should be empty", r.out, r.out_len);
  TEST_ASSERT_TRUE_MESSAGE(r.err_len > 0, "stderr should be non-empty for failure");
  TEST_ASSERT_NOT_EQUAL_MESSAGE('\0', r.err[0], "stderr should contain data");

  free_run_result (&r);
}

/* Test: usage(0) prints full help to stdout, nothing to stderr, and exits 0. */
void test_usage_success_status_zero(void) {
  RunResult r = {0};
  run_usage_capture (0, &r);

  TEST_ASSERT_EQUAL_INT_MESSAGE(0, r.exit_status, "usage(0) exit status mismatch");
  TEST_ASSERT_TRUE_MESSAGE(r.out_len > 0, "stdout should be non-empty for success");
  TEST_ASSERT_NOT_EQUAL_MESSAGE('\0', r.out[0], "stdout should contain data");
  assert_buffer_empty ("stderr should be empty", r.err, r.err_len);

  free_run_result (&r);
}

int main(void) {
  UNITY_BEGIN();
  RUN_TEST(test_usage_failure_status_one);
  RUN_TEST(test_usage_failure_status_127);
  RUN_TEST(test_usage_success_status_zero);
  return UNITY_END();
}