#include "../../unity/unity.h"

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

/* Access the program_name used by usage() for printing. */
extern char const *program_name;

/* Declare the target function so tests can call it. */
void usage (int status);

/* Helper to read all data from an fd into a malloc'ed buffer. */
static char *read_all_from_fd(int fd, size_t *out_len)
{
  size_t cap = 4096;
  size_t len = 0;
  char *buf = (char *)malloc(cap);
  if (!buf)
    return NULL;

  for (;;)
    {
      if (len == 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 n = read(fd, buf + len, cap - len);
      if (n > 0)
        {
          len += (size_t)n;
          continue;
        }
      if (n == 0)
        break;
      if (errno == EINTR)
        continue;
      /* Some read error; return what we have so far. */
      break;
    }

  /* NUL-terminate for convenience in strstr, though length is authoritative. */
  if (len == cap)
    {
      char *nb = (char *)realloc(buf, cap + 1);
      if (!nb)
        {
          free(buf);
          return NULL;
        }
      buf = nb;
    }
  buf[len] = '\0';
  if (out_len)
    *out_len = len;
  return buf;
}

typedef struct {
  char *out_buf;
  size_t out_len;
  char *err_buf;
  size_t err_len;
  int exit_code;    /* -1 if terminated by signal or abnormal. */
  bool exited_normally;
} usage_result_t;

static usage_result_t run_usage_and_capture(int status_to_call)
{
  usage_result_t r = {0};
  int out_pipe[2] = {-1, -1};
  int err_pipe[2] = {-1, -1};

  if (pipe(out_pipe) != 0 || pipe(err_pipe) != 0)
    {
      /* On pipe failure, return empty result with exit_code -1. */
      r.exit_code = -1;
      r.exited_normally = false;
      if (out_pipe[0] != -1) close(out_pipe[0]);
      if (out_pipe[1] != -1) close(out_pipe[1]);
      if (err_pipe[0] != -1) close(err_pipe[0]);
      if (err_pipe[1] != -1) close(err_pipe[1]);
      return r;
    }

  fflush(stdout);
    fflush(stderr);
    pid_t pid = fork();
  if (pid == -1)
    {
      /* Fork failed. */
      r.exit_code = -1;
      r.exited_normally = false;
      close(out_pipe[0]); close(out_pipe[1]);
      close(err_pipe[0]); close(err_pipe[1]);
      return r;
    }

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

      /* Duplicate write ends to STDOUT and STDERR. */
      if (dup2(out_pipe[1], STDOUT_FILENO) == -1) _exit(127);
      if (dup2(err_pipe[1], STDERR_FILENO) == -1) _exit(127);

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

      /* Call the function; it will call exit(status_to_call). */
      usage(status_to_call);

      /* Should never reach here. */
      _exit(125);
    }

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

  r.out_buf = read_all_from_fd(out_pipe[0], &r.out_len);
  r.err_buf = read_all_from_fd(err_pipe[0], &r.err_len);

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

  int status;
  pid_t w = waitpid(pid, &status, 0);
  if (w == pid)
    {
      if (WIFEXITED(status))
        {
          r.exited_normally = true;
          r.exit_code = WEXITSTATUS(status);
        }
      else if (WIFSIGNALED(status))
        {
          r.exited_normally = false;
          r.exit_code = -1;
        }
      else
        {
          r.exited_normally = false;
          r.exit_code = -1;
        }
    }
  else
    {
      r.exited_normally = false;
      r.exit_code = -1;
    }

  return r;
}

void setUp(void) {
  /* Ensure program_name is non-null and recognizable in outputs. */
  program_name = "ln";
}

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

/* Tests */

static void free_usage_result(usage_result_t *r)
{
  if (!r) return;
  free(r->out_buf);
  free(r->err_buf);
  memset(r, 0, sizeof(*r));
}

void test_usage_exits_success_and_writes_to_stdout_only(void)
{
  program_name = "ln-success-test";
  usage_result_t r = run_usage_and_capture(EXIT_SUCCESS);

  TEST_ASSERT_TRUE_MESSAGE(r.exited_normally, "Child did not exit normally");
  TEST_ASSERT_EQUAL_INT_MESSAGE(EXIT_SUCCESS, r.exit_code, "Exit code mismatch for success");
  TEST_ASSERT_NOT_NULL_MESSAGE(r.out_buf, "Expected stdout buffer");
  TEST_ASSERT_TRUE_MESSAGE(r.out_len > 0, "Expected non-empty stdout on success");
  TEST_ASSERT_TRUE_MESSAGE(r.err_len == 0, "Expected empty stderr on success");
  /* Program name should appear somewhere in the help text. */
  TEST_ASSERT_NOT_NULL_MESSAGE(strstr(r.out_buf, "ln-success-test"),
                               "program_name not found in success help output");

  free_usage_result(&r);
}

void test_usage_exits_with_given_error_status_and_writes_to_stderr_only(void)
{
  program_name = "ln-error-test";
  int exp = 2;
  usage_result_t r = run_usage_and_capture(exp);

  TEST_ASSERT_TRUE_MESSAGE(r.exited_normally, "Child did not exit normally");
  TEST_ASSERT_EQUAL_INT_MESSAGE(exp, r.exit_code, "Exit code mismatch for error");
  TEST_ASSERT_TRUE_MESSAGE(r.out_len == 0, "Expected empty stdout on error");
  TEST_ASSERT_NOT_NULL_MESSAGE(r.err_buf, "Expected stderr buffer on error");
  TEST_ASSERT_TRUE_MESSAGE(r.err_len > 0, "Expected non-empty stderr on error");
  /* Program name should appear in the 'try help' message. */
  TEST_ASSERT_NOT_NULL_MESSAGE(strstr(r.err_buf, "ln-error-test"),
                               "program_name not found in error help output");

  free_usage_result(&r);
}

void test_usage_handles_program_name_change_between_calls(void)
{
  /* Call with one name, then another, ensuring each appears as expected. */
  program_name = "ln-first-name";
  usage_result_t r1 = run_usage_and_capture(EXIT_SUCCESS);

  TEST_ASSERT_TRUE(r1.exited_normally);
  TEST_ASSERT_EQUAL_INT(EXIT_SUCCESS, r1.exit_code);
  TEST_ASSERT_NOT_NULL(r1.out_buf);
  TEST_ASSERT_TRUE(r1.out_len > 0);
  TEST_ASSERT_NOT_NULL(strstr(r1.out_buf, "ln-first-name"));
  free_usage_result(&r1);

  program_name = "ln-second-name";
  usage_result_t r2 = run_usage_and_capture(1);

  TEST_ASSERT_TRUE(r2.exited_normally);
  TEST_ASSERT_EQUAL_INT(1, r2.exit_code);
  TEST_ASSERT_TRUE(r2.out_len == 0);
  TEST_ASSERT_NOT_NULL(r2.err_buf);
  TEST_ASSERT_TRUE(r2.err_len > 0);
  TEST_ASSERT_NOT_NULL(strstr(r2.err_buf, "ln-second-name"));

  free_usage_result(&r2);
}

int main(void)
{
  UNITY_BEGIN();
  RUN_TEST(test_usage_exits_success_and_writes_to_stdout_only);
  RUN_TEST(test_usage_exits_with_given_error_status_and_writes_to_stderr_only);
  RUN_TEST(test_usage_handles_program_name_change_between_calls);
  return UNITY_END();
}