#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 <errno.h>
#include <locale.h>

/* usage(int) is defined in the including translation unit above.
   set_program_name is provided by gnulib and is available via the
   already-included headers from the program source. */
void usage (int status);
void set_program_name (const char *);

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

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

/* Helpers */

static char *read_all_from_fd (int fd) {
  size_t cap = 4096;
  size_t len = 0;
  char *buf = (char *) malloc(cap + 1);
  if (!buf) return NULL;

  for (;;) {
    ssize_t n = read(fd, buf + len, cap - len);
    if (n > 0) {
      len += (size_t)n;
      if (len == cap) {
        size_t new_cap = cap * 2;
        char *tmp = (char *) realloc(buf, new_cap + 1);
        if (!tmp) {
          free(buf);
          return NULL;
        }
        buf = tmp;
        cap = new_cap;
      }
    } else if (n == 0) {
      break;
    } else {
      if (errno == EINTR) continue;
      free(buf);
      return NULL;
    }
  }
  buf[len] = '\0';
  return buf;
}

static int run_usage_and_capture (int status, const char *progname,
                                  char **out_stdout, char **out_stderr,
                                  int *out_exit_code) {
  int out_pipe[2] = {-1, -1};
  int err_pipe[2] = {-1, -1};

  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: set deterministic locale and program name, redirect, and call usage */
    setenv("LC_ALL", "C", 1);
    if (progname) set_program_name(progname);

    if (dup2(out_pipe[1], STDOUT_FILENO) < 0) _exit(255);
    if (dup2(err_pipe[1], STDERR_FILENO) < 0) _exit(255);

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

    usage(status);
    _exit(254); /* Should not reach */
  }

  /* Parent */
  close(out_pipe[1]); close(err_pipe[1]);

  char *captured_out = read_all_from_fd(out_pipe[0]);
  char *captured_err = read_all_from_fd(err_pipe[0]);
  close(out_pipe[0]); close(err_pipe[0]);

  int wstatus = 0;
  int rc;
  do {
    rc = waitpid(pid, &wstatus, 0);
  } while (rc < 0 && errno == EINTR);

  if (rc < 0) {
    free(captured_out);
    free(captured_err);
    return -1;
  }

  int exit_code;
  if (WIFEXITED(wstatus)) {
    exit_code = WEXITSTATUS(wstatus);
  } else if (WIFSIGNALED(wstatus)) {
    exit_code = 128 + WTERMSIG(wstatus);
  } else {
    exit_code = -1;
  }

  if (out_stdout) *out_stdout = captured_out; else free(captured_out);
  if (out_stderr) *out_stderr = captured_err; else free(captured_err);
  if (out_exit_code) *out_exit_code = exit_code;

  return 0;
}

static int str_contains (const char *haystack, const char *needle) {
  return haystack && needle && strstr(haystack, needle) != NULL;
}

/* Tests */

void test_usage_success_prints_expected_lines(void) {
  char *out = NULL, *err = NULL;
  int code = -1;

  int rc = run_usage_and_capture(EXIT_SUCCESS, "nproc", &out, &err, &code);
  TEST_ASSERT_EQUAL_INT_MESSAGE(0, rc, "Failed to execute usage in child");
  TEST_ASSERT_NOT_NULL_MESSAGE(out, "No stdout captured");
  TEST_ASSERT_NOT_NULL_MESSAGE(err, "No stderr captured");
  TEST_ASSERT_EQUAL_INT_MESSAGE(0, code, "usage(EXIT_SUCCESS) did not exit with 0");

  /* Ensure key substrings are present in help text */
  TEST_ASSERT_TRUE_MESSAGE(str_contains(out, "Usage: nproc [OPTION]..."),
                           "Missing or incorrect Usage line");
  TEST_ASSERT_TRUE_MESSAGE(str_contains(out, "--all"),
                           "Missing '--all' option in help");
  TEST_ASSERT_TRUE_MESSAGE(str_contains(out, "print the number of installed processors"),
                           "Missing description for --all");
  TEST_ASSERT_TRUE_MESSAGE(str_contains(out, "--ignore=N"),
                           "Missing '--ignore=N' option in help");
  TEST_ASSERT_TRUE_MESSAGE(str_contains(out, "exclude N processing units"),
                           "Missing description for --ignore=N");

  /* Should not print the 'Try ... --help' line in the success path */
  TEST_ASSERT_FALSE_MESSAGE(str_contains(out, "Try '"),
                            "Unexpected 'Try' hint in stdout on success");

  /* No error output expected on success */
  TEST_ASSERT_TRUE_MESSAGE(strlen(err) == 0, "Stderr not empty on success");

  free(out);
  free(err);
}

void test_usage_error_prints_try_help_to_stderr_only(void) {
  char *out = NULL, *err = NULL;
  int code = -1;

  int rc = run_usage_and_capture(1, "nproc", &out, &err, &code);
  TEST_ASSERT_EQUAL_INT_MESSAGE(0, rc, "Failed to execute usage in child");
  TEST_ASSERT_NOT_NULL_MESSAGE(out, "No stdout captured");
  TEST_ASSERT_NOT_NULL_MESSAGE(err, "No stderr captured");

  TEST_ASSERT_EQUAL_INT_MESSAGE(1, code, "Exit status should match provided nonzero status");

  /* On error, stdout should be empty, stderr should have the 'Try ... --help' hint */
  TEST_ASSERT_TRUE_MESSAGE(strlen(out) == 0, "Stdout should be empty on error path");
  TEST_ASSERT_TRUE_MESSAGE(str_contains(err, "Try '"), "Missing 'Try' hint in stderr");
  TEST_ASSERT_TRUE_MESSAGE(str_contains(err, "nproc"), "Program name not present in hint");
  TEST_ASSERT_TRUE_MESSAGE(str_contains(err, "--help"), "Missing '--help' in hint");
  TEST_ASSERT_TRUE_MESSAGE(str_contains(err, "more information"),
                           "Missing 'more information' in hint");

  free(out);
  free(err);
}

void test_usage_respects_program_name(void) {
  char *out = NULL, *err = NULL;
  int code = -1;

  const char *custom = "my-nproc";
  int rc = run_usage_and_capture(EXIT_SUCCESS, custom, &out, &err, &code);
  TEST_ASSERT_EQUAL_INT_MESSAGE(0, rc, "Failed to execute usage in child");
  TEST_ASSERT_EQUAL_INT_MESSAGE(0, code, "usage(EXIT_SUCCESS) did not exit with 0");
  TEST_ASSERT_NOT_NULL(out);
  TEST_ASSERT_NOT_NULL(err);

  TEST_ASSERT_TRUE_MESSAGE(str_contains(out, "Usage: my-nproc [OPTION]..."),
                           "Usage line did not reflect custom program_name");
  TEST_ASSERT_TRUE_MESSAGE(strlen(err) == 0, "Stderr not empty on success");

  free(out);
  free(err);
}

void test_usage_nonzero_exit_code_is_preserved(void) {
  char *out = NULL, *err = NULL;
  int code = -1;
  int expected = 7;

  int rc = run_usage_and_capture(expected, "nproc", &out, &err, &code);
  TEST_ASSERT_EQUAL_INT_MESSAGE(0, rc, "Failed to execute usage in child");
  TEST_ASSERT_EQUAL_INT_MESSAGE(expected, code, "Nonzero exit code not preserved");

  /* On error path, stdout is empty, stderr has some hint */
  TEST_ASSERT_TRUE_MESSAGE(strlen(out) == 0, "Stdout should be empty on error path");
  TEST_ASSERT_TRUE_MESSAGE(strlen(err) > 0, "Stderr should contain hint on error path");

  free(out);
  free(err);
}

/* Unity runner */
int main(void) {
  /* Initialize program_name for the test process itself. Children override as needed. */
  set_program_name("nproc");

  UNITY_BEGIN();
  RUN_TEST(test_usage_success_prints_expected_lines);
  RUN_TEST(test_usage_error_prints_try_help_to_stderr_only);
  RUN_TEST(test_usage_respects_program_name);
  RUN_TEST(test_usage_nonzero_exit_code_is_preserved);
  return UNITY_END();
}