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

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>
#include <locale.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/select.h>
#include <errno.h>

#include "progname.h" /* for set_program_name */

/* Prototype of the function under test */
void usage (int status);

/* Helper: read all available data from fd into a malloc'd buffer.
   Uses select in the parent to multiplex stdout/stderr of the child. */

typedef struct {
  char *data;
  size_t size;
  size_t cap;
  int fd;
  bool closed;
} capture_buf;

static void buf_init(capture_buf *b, int fd) {
  b->data = NULL;
  b->size = 0;
  b->cap = 0;
  b->fd = fd;
  b->closed = false;
}

static void buf_reserve(capture_buf *b, size_t add) {
  if (b->size + add + 1 > b->cap) {
    size_t new_cap = b->cap ? b->cap : 1024;
    while (b->size + add + 1 > new_cap) {
      new_cap *= 2;
    }
    char *p = (char*)realloc(b->data, new_cap);
    if (!p) {
      /* In tests, aborting is acceptable */
      perror("realloc");
      exit(1);
    }
    b->data = p;
    b->cap = new_cap;
  }
}

static int read_into_buf(capture_buf *b) {
  char tmp[4096];
  ssize_t n = read(b->fd, tmp, sizeof tmp);
  if (n > 0) {
    buf_reserve(b, (size_t)n);
    memcpy(b->data + b->size, tmp, (size_t)n);
    b->size += (size_t)n;
    b->data[b->size] = '\0';
    return 1;
  } else if (n == 0) {
    close(b->fd);
    b->closed = true;
    return 0;
  } else {
    if (errno == EINTR) return 1;
    /* Treat other errors as stream closed */
    close(b->fd);
    b->closed = true;
    return 0;
  }
}

/* Run usage(status) in a child, capturing stdout and stderr. */
static int run_usage_and_capture(int status,
                                 char **out_str, size_t *out_len,
                                 char **err_str, size_t *err_len,
                                 int *exit_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/stderr, then call usage */
    /* Ensure C locale in child as well */
    setlocale(LC_ALL, "C");
    /* program_name should already be set by parent; inherited on fork */

    dup2(out_pipe[1], STDOUT_FILENO);
    dup2(err_pipe[1], STDERR_FILENO);
    close(out_pipe[0]); close(out_pipe[1]);
    close(err_pipe[0]); close(err_pipe[1]);

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

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

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

  capture_buf bout, berr;
  buf_init(&bout, out_pipe[0]);
  buf_init(&berr, err_pipe[0]);

  while (!bout.closed || !berr.closed) {
    fd_set rfds;
    FD_ZERO(&rfds);
    int maxfd = -1;
    if (!bout.closed) { FD_SET(bout.fd, &rfds); if (bout.fd > maxfd) maxfd = bout.fd; }
    if (!berr.closed) { FD_SET(berr.fd, &rfds); if (berr.fd > maxfd) maxfd = berr.fd; }

    int sel = select(maxfd + 1, &rfds, NULL, NULL, NULL);
    if (sel < 0) {
      if (errno == EINTR) continue;
      break;
    }
    if (!bout.closed && FD_ISSET(bout.fd, &rfds)) {
      read_into_buf(&bout);
    }
    if (!berr.closed && FD_ISSET(berr.fd, &rfds)) {
      read_into_buf(&berr);
    }
  }

  int status_code = -1;
  int wstatus = 0;
  if (waitpid(pid, &wstatus, 0) >= 0) {
    if (WIFEXITED(wstatus)) status_code = WEXITSTATUS(wstatus);
    else if (WIFSIGNALED(wstatus)) status_code = 128 + WTERMSIG(wstatus);
    else status_code = -1;
  }

  /* Null-terminate buffers */
  if (bout.data == NULL) {
    bout.data = (char*)malloc(1);
    if (bout.data) bout.data[0] = '\0';
  } else {
    bout.data[bout.size] = '\0';
  }
  if (berr.data == NULL) {
    berr.data = (char*)malloc(1);
    if (berr.data) berr.data[0] = '\0';
  } else {
    berr.data[berr.size] = '\0';
  }

  if (out_str) *out_str = bout.data; else free(bout.data);
  if (out_len) *out_len = bout.size;
  if (err_str) *err_str = berr.data; else free(berr.data);
  if (err_len) *err_len = berr.size;
  if (exit_status) *exit_status = status_code;

  return 0;
}

void setUp(void) {
  /* Ensure consistent locale for translatable strings */
  setlocale(LC_ALL, "C");
  /* Ensure program_name is set for correct message formatting */
  set_program_name("numfmt");
}

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

static void assert_substring(const char *haystack, const char *needle, const char *ctx) {
  if (!haystack || !needle) {
    TEST_FAIL_MESSAGE("Invalid input to assert_substring");
  }
  if (strstr(haystack, needle) == NULL) {
    char msg[512];
    snprintf(msg, sizeof msg, "Expected to find substring '%s' in %s output", needle, ctx);
    TEST_FAIL_MESSAGE(msg);
  }
}

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

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

  /* Expect non-empty stdout, empty stderr */
  TEST_ASSERT_NOT_NULL(out);
  TEST_ASSERT_TRUE(out_len > 0);
  TEST_ASSERT_NOT_NULL(err);
  TEST_ASSERT_EQUAL_UINT64(0, (unsigned long long)err_len);

  /* Check for some key tokens in help output */
  assert_substring(out, "Usage:", "help");
  assert_substring(out, "numfmt", "help");
  /* Options and examples are shown; check a couple of representative strings */
  assert_substring(out, "--to=si", "help");
  assert_substring(out, "--from=iec", "help");

  free(out);
  free(err);
}

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

  int wanted_status = 1;
  int rc = run_usage_and_capture(wanted_status, &out, &out_len, &err, &err_len, &code);
  TEST_ASSERT_EQUAL_INT(0, rc);
  TEST_ASSERT_EQUAL_INT(wanted_status, code);

  /* Expect empty stdout, non-empty stderr containing a 'Try ... --help' message */
  TEST_ASSERT_NOT_NULL(out);
  TEST_ASSERT_EQUAL_UINT64(0, (unsigned long long)out_len);
  TEST_ASSERT_NOT_NULL(err);
  TEST_ASSERT_TRUE(err_len > 0);

  assert_substring(err, "Try", "try-help");
  assert_substring(err, "--help", "try-help");
  assert_substring(err, "numfmt", "try-help");

  free(out);
  free(err);
}

int main(void) {
  UNITY_BEGIN();
  RUN_TEST(test_usage_success_prints_help_and_exits_zero);
  RUN_TEST(test_usage_nonzero_status_prints_try_help_and_exits_status);
  return UNITY_END();
}