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

/* Unity setup/teardown */
void setUp(void) {
  /* Ensure our parent process uses the C locale as well, though we also set it in the child. */
  setenv("LC_ALL", "C", 1);
}

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

/* Helper to read all data from an fd into a malloc'd buffer, NUL terminated.
   Returns buffer and sets *len_out to length (excluding NUL). */
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) return NULL;

  while (1) {
    if (len + 2048 > cap) {
      size_t new_cap = cap * 2;
      char *nb = (char *)realloc(buf, new_cap);
      if (!nb) {
        free(buf);
        return NULL;
      }
      buf = nb;
      cap = new_cap;
    }
    ssize_t r = read(fd, buf + len, cap - len - 1);
    if (r < 0) {
      if (errno == EINTR) continue;
      free(buf);
      return NULL;
    }
    if (r == 0) break;
    len += (size_t)r;
  }
  buf[len] = '\0';
  if (len_out) *len_out = len;
  return buf;
}

/* Run usage(status) in a child, capturing stdout and stderr.
   Returns 0 on success, -1 on failure to fork/pipe.
   On success, fills exit_code, out_stdout (malloc'd), out_stderr (malloc'd). */
static int run_usage_and_capture(int status,
                                 int *exit_code,
                                 char **out_stdout, size_t *out_stdout_len,
                                 char **out_stderr, size_t *out_stderr_len) {
  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: redirect stdout/stderr, set environment, call usage() */
    /* Locale and program name for predictable messages. */
    setenv("LC_ALL", "C", 1);
    setlocale(LC_ALL, "C");

    /* Ensure program_name is set; available via system.h/progname.h earlier. */
    set_program_name("nl");

    /* Redirect stdout and stderr to our pipes. */
    dup2(out_pipe[1], STDOUT_FILENO);
    dup2(err_pipe[1], STDERR_FILENO);

    /* Close unused fds. */
    close(out_pipe[0]); close(out_pipe[1]);
    close(err_pipe[0]); close(err_pipe[1]);

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

    /* Should never reach here, but be safe. */
    _exit(255);
  }

  /* Parent: close write ends, read from read ends. */
  close(out_pipe[1]);
  close(err_pipe[1]);

  size_t so_len = 0, se_len = 0;
  char *so = read_all_from_fd(out_pipe[0], &so_len);
  char *se = read_all_from_fd(err_pipe[0], &se_len);

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

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

  if (wait_ok < 0) {
    free(so);
    free(se);
    return -1;
  }

  int code = -1;
  if (WIFEXITED(wstatus)) {
    code = WEXITSTATUS(wstatus);
  } else if (WIFSIGNALED(wstatus)) {
    /* If it died by signal, encode as negative signal number to distinguish. */
    code = -WTERMSIG(wstatus);
  }

  if (exit_code) *exit_code = code;
  if (out_stdout) *out_stdout = so; else free(so);
  if (out_stdout_len) *out_stdout_len = so_len;
  if (out_stderr) *out_stderr = se; else free(se);
  if (out_stderr_len) *out_stderr_len = se_len;

  return 0;
}

/* Utility: check if haystack contains needle (non-NULL-safe). */
static int contains_substr(const char *haystack, const char *needle) {
  return haystack && needle && strstr(haystack, needle) != NULL;
}

/* Tests */

void test_usage_success_prints_full_usage_to_stdout_and_exits_zero(void) {
  int exit_code = -1;
  char *out = NULL;
  char *err = NULL;
  size_t out_len = 0, err_len = 0;

  int rc = run_usage_and_capture(EXIT_SUCCESS, &exit_code, &out, &out_len, &err, &err_len);
  TEST_ASSERT_EQUAL_INT(0, rc);
  TEST_ASSERT_EQUAL_INT(EXIT_SUCCESS, exit_code);

  /* Verify stdout has expected key phrases. */
  TEST_ASSERT_NOT_NULL(out);
  TEST_ASSERT_TRUE(out_len > 0);
  TEST_ASSERT_TRUE_MESSAGE(contains_substr(out, "Usage: nl [OPTION]... [FILE]..."),
                           "stdout should contain the usage synopsis line");
  TEST_ASSERT_TRUE_MESSAGE(contains_substr(out, "Default options are:"),
                           "stdout should contain the default options section");

  /* Verify stderr is empty or only whitespace. */
  if (err && err_len > 0) {
    /* Trim whitespace to check emptiness. */
    size_t i;
    int non_ws = 0;
    for (i = 0; i < err_len; i++) {
      if (!(err[i] == ' ' || err[i] == '\t' || err[i] == '\n' || err[i] == '\r' || err[i] == '\f' || err[i] == '\v'))
        { non_ws = 1; break; }
    }
    TEST_ASSERT_FALSE_MESSAGE(non_ws, "stderr should be empty for EXIT_SUCCESS");
  }

  free(out);
  free(err);
}

void test_usage_failure_prints_try_help_to_stderr_and_exits_status(void) {
  int exit_code = -1;
  char *out = NULL;
  char *err = NULL;
  size_t out_len = 0, err_len = 0;

  int nonzero_status = 2;
  int rc = run_usage_and_capture(nonzero_status, &exit_code, &out, &out_len, &err, &err_len);
  TEST_ASSERT_EQUAL_INT(0, rc);
  TEST_ASSERT_EQUAL_INT(nonzero_status, exit_code);

  /* stdout should be empty or whitespace only */
  if (out && out_len > 0) {
    size_t i;
    int non_ws = 0;
    for (i = 0; i < out_len; i++) {
      if (!(out[i] == ' ' || out[i] == '\t' || out[i] == '\n' || out[i] == '\r' || out[i] == '\f' || out[i] == '\v'))
        { non_ws = 1; break; }
    }
    TEST_ASSERT_FALSE_MESSAGE(non_ws, "stdout should be empty for non-success status");
  }

  /* stderr should contain guidance to use --help, including the program name. */
  TEST_ASSERT_NOT_NULL_MESSAGE(err, "stderr capture failed");
  TEST_ASSERT_TRUE_MESSAGE(err_len > 0, "stderr should not be empty for non-success status");
  TEST_ASSERT_TRUE_MESSAGE(contains_substr(err, "nl --help"),
                           "stderr should mention 'nl --help'");
  TEST_ASSERT_TRUE_MESSAGE(contains_substr(err, "Try '"),
                           "stderr should begin with 'Try ...' guidance");
  TEST_ASSERT_TRUE_MESSAGE(contains_substr(err, "more information."),
                           "stderr should contain 'more information.' guidance");

  free(out);
  free(err);
}

int main(void) {
  UNITY_BEGIN();
  RUN_TEST(test_usage_success_prints_full_usage_to_stdout_and_exits_zero);
  RUN_TEST(test_usage_failure_prints_try_help_to_stderr_and_exits_status);
  return UNITY_END();
}