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

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

/* usage(int) is defined in the included program source.
   system.h (already included by the program) declares set_program_name. */

/* Unity setup/teardown */
void setUp(void) {
  /* no-op */
}
void tearDown(void) {
  /* no-op */
}

/* Helper to read all data from fd into a malloc'd buffer.
   Returns a NUL-terminated buffer and length via out parameters.
   On error, returns NULL and sets *out_len to 0. */
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 + 1);
  if (!buf) {
    *out_len = 0;
    return NULL;
  }

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

/* Run usage(status) in a child process, capturing stdout/stderr and exit code.
   If progname is non-NULL, set program_name to progname in child via set_program_name().
   Returns 0 on success capturing; nonzero on infrastructure error. */
static int run_usage_and_capture(int status_to_pass,
                                 const char* progname,
                                 char** out_buf, size_t* out_len,
                                 char** err_buf, size_t* err_len,
                                 int* exit_code, int* term_sig)
{
  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 -2;
  }

  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 -3;
  }

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

    /* Duplicate write ends to STDOUT/ERR */
    if (dup2(out_pipe[1], STDOUT_FILENO) < 0) _exit(127);
    if (dup2(err_pipe[1], STDERR_FILENO) < 0) _exit(127);

    /* Close original write ends */
    close(out_pipe[1]);
    close(err_pipe[1]);

    /* Set program name if requested */
    if (progname) {
      /* Declared in system.h (already included by program source) */
      extern void set_program_name (const char *);
      set_program_name(progname);
    }

    /* Call target function; it will call exit(). */
    usage(status_to_pass);

    /* Should not reach here; ensure child exits even if usage doesn't exit. */
    _exit(126);
  }

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

  char* out = read_all_from_fd(out_pipe[0], out_len);
  char* err = read_all_from_fd(err_pipe[0], err_len);

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

  int wstatus = 0;
  if (waitpid(pid, &wstatus, 0) < 0) {
    if (out) free(out);
    if (err) free(err);
    return -4;
  }

  *exit_code = -1;
  *term_sig = 0;
  if (WIFEXITED(wstatus)) {
    *exit_code = WEXITSTATUS(wstatus);
  } else if (WIFSIGNALED(wstatus)) {
    *term_sig = WTERMSIG(wstatus);
  }

  *out_buf = out ? out : strdup("");
  *err_buf = err ? err : strdup("");

  if (!*out_buf || !*err_buf) {
    if (*out_buf) free(*out_buf);
    if (*err_buf) free(*err_buf);
    return -5;
  }

  return 0;
}

/* Utility: check substring presence safely. */
static bool contains_substr(const char* haystack, const char* needle) {
  if (!haystack || !needle) return false;
  return strstr(haystack, needle) != NULL;
}

/* Test 1: EXIT_SUCCESS should print full help to stdout, nothing to stderr, and exit(0). */
static void test_usage_success_outputs_help_and_exits_zero(void) {
  char *out = NULL, *err = NULL;
  size_t out_len = 0, err_len = 0;
  int code = -1, sig = 0;

  int rc = run_usage_and_capture(EXIT_SUCCESS, "dd",
                                 &out, &out_len, &err, &err_len,
                                 &code, &sig);

  TEST_ASSERT_EQUAL_INT_MESSAGE(0, rc, "Harness failed to run usage(EXIT_SUCCESS)");
  TEST_ASSERT_EQUAL_INT_MESSAGE(0, sig, "Child terminated by signal unexpectedly");
  TEST_ASSERT_EQUAL_INT(EXIT_SUCCESS, code);

  /* Expect non-empty stdout, empty stderr. */
  TEST_ASSERT_MESSAGE(out_len > 0, "Expected help text on stdout");
  TEST_ASSERT_EQUAL_UINT_MESSAGE(0u, err_len, "Expected no stderr output for success usage");

  /* Help should likely mention program name and --help option somewhere. */
  TEST_ASSERT_TRUE_MESSAGE(contains_substr(out, "dd"),
                           "Expected program name 'dd' in help output");
  TEST_ASSERT_TRUE_MESSAGE(contains_substr(out, "--help"),
                           "Expected '--help' mention in full help output");

  free(out);
  free(err);
}

/* Test 2: Non-zero status should print brief hint to stderr, nothing to stdout, and exit(status). */
static void test_usage_failure_emits_try_help_and_exits_status(void) {
  char *out = NULL, *err = NULL;
  size_t out_len = 0, err_len = 0;
  int code = -1, sig = 0;

  int rc = run_usage_and_capture(EXIT_FAILURE, "dd",
                                 &out, &out_len, &err, &err_len,
                                 &code, &sig);

  TEST_ASSERT_EQUAL_INT_MESSAGE(0, rc, "Harness failed to run usage(EXIT_FAILURE)");
  TEST_ASSERT_EQUAL_INT_MESSAGE(0, sig, "Child terminated by signal unexpectedly");
  TEST_ASSERT_EQUAL_INT(EXIT_FAILURE, code);

  /* Expect empty stdout, non-empty stderr. */
  TEST_ASSERT_EQUAL_UINT_MESSAGE(0u, out_len, "Expected no stdout output for failure usage");
  TEST_ASSERT_MESSAGE(err_len > 0, "Expected message on stderr for failure usage");

  /* Expect hint to include program name and --help. */
  TEST_ASSERT_TRUE_MESSAGE(contains_substr(err, "dd"),
                           "Expected program name 'dd' in error hint");
  TEST_ASSERT_TRUE_MESSAGE(contains_substr(err, "--help"),
                           "Expected '--help' mention in error hint");

  free(out);
  free(err);
}

/* Test 3: The message should incorporate the current program_name. */
static void test_usage_uses_program_name_variable(void) {
  char *out = NULL, *err = NULL;
  size_t out_len = 0, err_len = 0;
  int code = -1, sig = 0;

  const char* custom = "mydd";
  int rc = run_usage_and_capture(2 /* arbitrary non-zero */,
                                 custom,
                                 &out, &out_len, &err, &err_len,
                                 &code, &sig);

  TEST_ASSERT_EQUAL_INT_MESSAGE(0, rc, "Harness failed to run usage(non-zero) with custom program name");
  TEST_ASSERT_EQUAL_INT_MESSAGE(0, sig, "Child terminated by signal unexpectedly");
  TEST_ASSERT_EQUAL_INT(2, code);

  /* Should be error path: stderr non-empty, stdout empty. */
  TEST_ASSERT_EQUAL_UINT_MESSAGE(0u, out_len, "Expected no stdout for failure usage");
  TEST_ASSERT_MESSAGE(err_len > 0, "Expected stderr output for failure usage");

  /* stderr should contain the program name we set */
  TEST_ASSERT_TRUE_MESSAGE(contains_substr(err, custom),
                           "Expected custom program name in failure message");
  TEST_ASSERT_TRUE_MESSAGE(contains_substr(err, "--help"),
                           "Expected '--help' mention in failure message");

  free(out);
  free(err);
}

int main(void) {
  UNITY_BEGIN();

  RUN_TEST(test_usage_success_outputs_help_and_exits_zero);
  RUN_TEST(test_usage_failure_emits_try_help_and_exits_status);
  RUN_TEST(test_usage_uses_program_name_variable);

  return UNITY_END();
}