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

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

/* Helper: capture what xwrite_stdout writes to stdout into a malloc'd buffer.
   Returns 0 on success, nonzero on failure to set up capture.
   On success, *out_buf is allocated (may be NULL if out_len==0),
   and *out_len is set to number of captured bytes. Caller must free(*out_buf). */
static int capture_xwrite_stdout_output(const char *in_buf, size_t n_bytes,
                                        char **out_buf, size_t *out_len)
{
  int rc = -1;
  int saved_stdout = -1;
  FILE *tmp = NULL;
  int tmpfd = -1;
  char *captured = NULL;
  size_t captured_len = 0;

  if (!out_buf || !out_len)
    return -1;

  *out_buf = NULL;
  *out_len = 0;

  fflush(stdout);
  saved_stdout = dup(fileno(stdout));
  if (saved_stdout < 0)
    goto cleanup;

  tmp = tmpfile();
  if (!tmp)
    goto cleanup;

  tmpfd = fileno(tmp);
  if (tmpfd < 0)
    goto cleanup;

  if (dup2(tmpfd, fileno(stdout)) < 0)
    goto cleanup;

  /* Perform the write while stdout is redirected. Do not call Unity here. */
  xwrite_stdout(in_buf, n_bytes);

  /* Flush redirected stdout and read back the data. */
  fflush(stdout);

  /* Determine size */
  long end_pos;
  if (fseek(tmp, 0, SEEK_END) != 0)
    goto cleanup;

  end_pos = ftell(tmp);
  if (end_pos < 0)
    goto cleanup;

  captured_len = (size_t)end_pos;

  if (fseek(tmp, 0, SEEK_SET) != 0)
    goto cleanup;

  if (captured_len > 0) {
    captured = (char *)malloc(captured_len + 1);
    if (!captured)
      goto cleanup;

    size_t nread = fread(captured, 1, captured_len, tmp);
    if (nread != captured_len)
      goto cleanup;

    captured[captured_len] = '\0'; /* for convenience; content may be binary */
  } else {
    captured = NULL; /* nothing written */
  }

  rc = 0;

cleanup:
  /* Restore stdout before any assertions by the caller. */
  if (saved_stdout >= 0) {
    fflush(stdout);
    dup2(saved_stdout, fileno(stdout));
    close(saved_stdout);
    saved_stdout = -1;
  }
  if (tmp) {
    fclose(tmp);
    tmp = NULL;
  }

  if (rc == 0) {
    *out_buf = captured;
    *out_len = captured_len;
  } else {
    free(captured);
  }

  return rc;
}

/* Test: writing non-zero data writes exactly those bytes */
static void test_xwrite_stdout_writes_expected_bytes(void)
{
  const char *msg = "Hello, world!";
  char *out = NULL;
  size_t out_len = 0;

  int rc = capture_xwrite_stdout_output(msg, 13, &out, &out_len);
  TEST_ASSERT_EQUAL_INT(0, rc);
  TEST_ASSERT_TRUE(out_len == 13);
  TEST_ASSERT_NOT_NULL(out);
  TEST_ASSERT_EQUAL_INT(0, memcmp(out, msg, 13));

  free(out);
}

/* Test: n_bytes == 0 should not write anything (non-null buffer) */
static void test_xwrite_stdout_zero_len_writes_nothing(void)
{
  const char *msg = "ABC";
  char *out = NULL;
  size_t out_len = 0;

  int rc = capture_xwrite_stdout_output(msg, 0, &out, &out_len);
  TEST_ASSERT_EQUAL_INT(0, rc);
  TEST_ASSERT_TRUE(out_len == 0);
  TEST_ASSERT_NULL(out);
}

/* Test: n_bytes == 0 with NULL buffer should be a no-op and not crash */
static void test_xwrite_stdout_null_buffer_and_zero_len_safe(void)
{
  char *out = NULL;
  size_t out_len = 0;

  int rc = capture_xwrite_stdout_output(NULL, 0, &out, &out_len);
  TEST_ASSERT_EQUAL_INT(0, rc);
  TEST_ASSERT_TRUE(out_len == 0);
  TEST_ASSERT_NULL(out);
}

/* Test: large write (larger than BUFSIZ) is fully written */
static void test_xwrite_stdout_large_write(void)
{
  /* Create a large buffer */
  size_t len = (size_t)BUFSIZ * 3 + 17;
  char *in = (char *)malloc(len);
  TEST_ASSERT_NOT_NULL(in);

  for (size_t i = 0; i < len; i++) {
    in[i] = (char)('A' + (i % 26));
  }

  char *out = NULL;
  size_t out_len = 0;
  int rc = capture_xwrite_stdout_output(in, len, &out, &out_len);
  TEST_ASSERT_EQUAL_INT(0, rc);
  TEST_ASSERT_TRUE(out_len == len);
  TEST_ASSERT_NOT_NULL(out);
  TEST_ASSERT_EQUAL_INT(0, memcmp(out, in, len));

  free(in);
  free(out);
}

/* Test: on write error, xwrite_stdout should terminate with EXIT_FAILURE.
   We simulate error by redirecting stdout in a child process to a non-writable FD. */
static void test_xwrite_stdout_when_write_fails_exits_with_failure(void)
{
  fflush(stdout);
    fflush(stderr);
    pid_t pid = fork();
  TEST_ASSERT_TRUE(pid >= 0);

  if (pid == 0) {
    /* Child: make stdout unwritable by dup'ing the read end of a pipe onto it. */
    int p[2];
    if (pipe(p) != 0) {
      _exit(127);
    }
    /* p[0] is read end (not writable); p[1] is write end */
    fflush(stdout);
    (void)dup2(p[0], fileno(stdout));
    close(p[0]);
    close(p[1]);

    /* This should fail to fwrite and thus call error(EXIT_FAILURE, ...) */
    xwrite_stdout("X", 1);

    /* We should never reach here; if we do, indicate unexpected success */
    _exit(0);
  } else {
    int status = 0;
    pid_t w = waitpid(pid, &status, 0);
    TEST_ASSERT_TRUE(w == pid);
    TEST_ASSERT_TRUE(WIFEXITED(status));
    TEST_ASSERT_EQUAL_INT(EXIT_FAILURE, WEXITSTATUS(status));
  }
}

int main(void)
{
  UNITY_BEGIN();
  RUN_TEST(test_xwrite_stdout_writes_expected_bytes);
  RUN_TEST(test_xwrite_stdout_zero_len_writes_nothing);
  RUN_TEST(test_xwrite_stdout_null_buffer_and_zero_len_safe);
  RUN_TEST(test_xwrite_stdout_large_write);
  RUN_TEST(test_xwrite_stdout_when_write_fails_exits_with_failure);
  return UNITY_END();
}