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

/* Helpers to capture output from a child that calls usage(int). */

/* Read all data from fd into a NUL-terminated malloc'd buffer.
   Returns buffer (caller must free). On error, returns NULL. */
static char *read_all_fd(int fd)
{
  size_t cap = 4096;
  size_t len = 0;
  char *buf = (char *)malloc(cap + 1);
  if (!buf) return NULL;

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

/* Run usage(status) in a child process, capturing stdout and stderr.
   On success, returns 0 and sets *out_stdout, *out_stderr, *out_status.
   On failure to set up, returns -1.
   The caller must free *out_stdout and *out_stderr. */
static int run_usage_and_capture(int status, char **out_stdout, char **out_stderr, int *out_status)
{
  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 and stderr and call usage. */
      /* Ensure deterministic English output and correct program_name. */
      setenv("LC_ALL", "C", 1);
      setlocale(LC_ALL, "C");
      set_program_name("head");

      /* Redirect stdout and stderr. */
      close(out_pipe[0]);
      close(err_pipe[0]);
      if (dup2(out_pipe[1], STDOUT_FILENO) < 0) _exit(127);
      if (dup2(err_pipe[1], STDERR_FILENO) < 0) _exit(127);
      close(out_pipe[1]);
      close(err_pipe[1]);

      /* Call the target; it will exit(status). */
      usage(status);

      /* Should not reach here, but just in case. */
      _exit(127);
    }

  /* Parent: close write ends, read both streams, then wait. */
  close(out_pipe[1]);
  close(err_pipe[1]);

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

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

  int wstatus = -1;
  if (waitpid(pid, &wstatus, 0) < 0)
    {
      free(captured_out);
      free(captured_err);
      return -1;
    }

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

  *out_stdout = captured_out ? captured_out : strdup("");
  *out_stderr = captured_err ? captured_err : strdup("");
  *out_status = code;
  return 0;
}

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

void setUp(void) {
  /* Nothing needed */
}

void tearDown(void) {
  /* Nothing needed */
}

void test_usage_success_outputs_expected_sections(void)
{
  char *out = NULL;
  char *err = NULL;
  int status = -1;

  int rc = run_usage_and_capture(EXIT_SUCCESS, &out, &err, &status);
  TEST_ASSERT_EQUAL_INT(0, rc);
  TEST_ASSERT_EQUAL_INT(EXIT_SUCCESS, status);

  /* Should write help to stdout and nothing (or minimal) to stderr. */
  TEST_ASSERT_NOT_NULL(out);
  TEST_ASSERT_TRUE(strlen(out) > 0);
  TEST_ASSERT_NOT_NULL(err);
  TEST_ASSERT_TRUE(strlen(err) == 0);

  /* Check for key help content. */
  TEST_ASSERT_TRUE_MESSAGE(str_contains(out, "Usage:"), "Missing 'Usage:' header");
  TEST_ASSERT_TRUE_MESSAGE(str_contains(out, "head"), "Missing program name 'head' in usage");
  TEST_ASSERT_TRUE_MESSAGE(str_contains(out, "-c, --bytes="), "Missing bytes option");
  TEST_ASSERT_TRUE_MESSAGE(str_contains(out, "-n, --lines="), "Missing lines option");
  TEST_ASSERT_TRUE_MESSAGE(str_contains(out, "-q, --quiet"), "Missing quiet option");
  TEST_ASSERT_TRUE_MESSAGE(str_contains(out, "-z, --zero-terminated"), "Missing zero-terminated option");
  TEST_ASSERT_TRUE_MESSAGE(str_contains(out, "NUM may have a multiplier suffix"), "Missing NUM suffix section");
  /* Check default number mention (10). */
  TEST_ASSERT_TRUE_MESSAGE(str_contains(out, "first 10"), "Missing default number of lines mention");

  free(out);
  free(err);
}

void test_usage_error_shows_try_help_and_status(void)
{
  char *out = NULL;
  char *err = NULL;
  int status = -1;

  int rc = run_usage_and_capture(EXIT_FAILURE, &out, &err, &status);
  TEST_ASSERT_EQUAL_INT(0, rc);
  TEST_ASSERT_EQUAL_INT(EXIT_FAILURE, status);

  /* Should print nothing to stdout, 'Try ... --help ...' to stderr. */
  TEST_ASSERT_NOT_NULL(out);
  TEST_ASSERT_TRUE_MESSAGE(strlen(out) == 0, "Expected no stdout on error usage");
  TEST_ASSERT_NOT_NULL(err);
  TEST_ASSERT_TRUE_MESSAGE(strlen(err) > 0, "Expected stderr on error usage");

  TEST_ASSERT_TRUE_MESSAGE(str_contains(err, "--help"), "Expected '--help' hint in stderr");
  TEST_ASSERT_TRUE_MESSAGE(str_contains(err, "head"), "Expected program name 'head' in stderr");

  free(out);
  free(err);
}

void test_usage_nonzero_status_code_propagated(void)
{
  char *out = NULL;
  char *err = NULL;
  int status = -1;

  int desired_status = 2;
  int rc = run_usage_and_capture(desired_status, &out, &err, &status);
  TEST_ASSERT_EQUAL_INT(0, rc);
  TEST_ASSERT_EQUAL_INT(desired_status, status);

  /* As with other non-success statuses, expect try-help on stderr. */
  TEST_ASSERT_TRUE_MESSAGE(strlen(out) == 0, "Expected no stdout for non-success status");
  TEST_ASSERT_TRUE_MESSAGE(str_contains(err, "--help"), "Expected '--help' hint in stderr for non-success status");
  TEST_ASSERT_TRUE_MESSAGE(str_contains(err, "head"), "Expected program name 'head' in stderr for non-success status");

  free(out);
  free(err);
}

int main(void)
{
  UNITY_BEGIN();
  RUN_TEST(test_usage_success_outputs_expected_sections);
  RUN_TEST(test_usage_error_shows_try_help_and_status);
  RUN_TEST(test_usage_nonzero_status_code_propagated);
  return UNITY_END();
}