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

/* Note: We rely on internal symbols from dd.c being visible here,
   since this file is included into the same translation unit. */

/* Helpers to manage stdout redirection in the parent process only. */
static int saved_stdout_fd = -1;

static int redirect_stdout_to_fd(int newfd)
{
  if (saved_stdout_fd != -1)
    return -1; /* already redirected */
  int dupfd = dup(STDOUT_FILENO);
  if (dupfd < 0)
    return -1;
  if (dup2(newfd, STDOUT_FILENO) < 0)
    {
      int e = errno;
      close(dupfd);
      errno = e;
      return -1;
    }
  saved_stdout_fd = dupfd;
  return 0;
}

static int redirect_stdout_to_devnull(void)
{
  int devnull = open("/dev/null", O_WRONLY);
  if (devnull < 0)
    return -1;
  int rc = redirect_stdout_to_fd(devnull);
  int e = errno;
  close(devnull);
  errno = e;
  return rc;
}

static int restore_stdout(void)
{
  if (saved_stdout_fd == -1)
    return -1;
  int rc = dup2(saved_stdout_fd, STDOUT_FILENO);
  int e = errno;
  close(saved_stdout_fd);
  saved_stdout_fd = -1;
  errno = e;
  return rc;
}

/* Fill a non-blocking pipe fd until it returns -1 with EAGAIN. */
static void fill_fd_nonblock_to_full(int fd)
{
  char buf[4096];
  memset(buf, 0xCD, sizeof buf);
  for (;;)
    {
      ssize_t wr = write(fd, buf, sizeof buf);
      if (wr < 0)
        {
          if (errno == EAGAIN)
            break;
          /* Unexpected error: just break to avoid infinite loop. */
          break;
        }
      else if (wr == 0)
        {
          /* Shouldn't happen on pipe; break defensively. */
          break;
        }
      /* else continue writing until full */
    }
}

/* Parse "X+Y records out" from an English stderr buffer; return 0 on success. */
static int parse_records_out(const char *buf, long long *full, long long *partial)
{
  const char *p = strstr(buf, "records out");
  if (!p)
    return -1;

  /* Find start of the line containing "records out" */
  const char *line_start = p;
  while (line_start > buf && line_start[-1] != '\n')
    line_start--;

  /* Try to scan two integers separated by '+' before "records out" */
  long long f = -1, pa = -1;
  /* We'll scan backwards to the '+'; simpler approach: scan from line_start. */
  /* Expected formats:
     "%jd+%jd records out"
  */
  if (sscanf(line_start, "%lld+%lld", &f, &pa) == 2)
    {
      if (full) *full = f;
      if (partial) *partial = pa;
      return 0;
    }
  return -1;
}

void setUp(void) {
  /* No-op by default. */
}

void tearDown(void) {
  /* Ensure stdout is restored in case a test forgot. */
  if (saved_stdout_fd != -1)
    restore_stdout();
}

/* Test: successful write updates counters and resets oc */
void test_write_output_success_basic(void)
{
  /* Prepare globals */
  output_blocksize = 1024;
  oc = 123; /* non-zero, should be reset */
  w_bytes = 0;
  w_full = 0;
  w_partial = 0;
  conversions_mask = 0;
  output_flags = 0;

  /* Prepare obuf */
  if (obuf)
    {
      /* Free any existing buffer to avoid memory waste between tests. */
      free(obuf);
      obuf = NULL;
    }
  obuf = (char*)malloc(output_blocksize);
  TEST_ASSERT_NOT_NULL(obuf);
  memset(obuf, 'A', output_blocksize);

  /* Redirect stdout to /dev/null to keep Unity output clean. */
  TEST_ASSERT_EQUAL_INT(0, redirect_stdout_to_devnull());

  /* Call target */
  write_output();

  /* Restore stdout before any assertions (Unity prints to stdout). */
  TEST_ASSERT_EQUAL_INT(0, restore_stdout());

  /* Validate */
  TEST_ASSERT_EQUAL_INT64((int64_t)output_blocksize, (int64_t)w_bytes);
  TEST_ASSERT_EQUAL_INT64(1, (int64_t)w_full);
  TEST_ASSERT_EQUAL_INT64(0, (int64_t)w_partial);
  TEST_ASSERT_EQUAL_UINT(0, oc);

  /* Cleanup */
  free(obuf);
  obuf = NULL;
}

/* Test: short write (0 bytes due to non-blocking full pipe) causes exit failure */
void test_write_output_short_write_zero_bytes_exits(void)
{
  fflush(stdout);
    fflush(stderr);
    pid_t pid = fork();
  TEST_ASSERT_TRUE(pid >= 0);

  if (pid == 0)
    {
      /* Child: set up environment to force write() to return EAGAIN immediately. */

      output_blocksize = 4096;
      oc = 7;
      w_bytes = 0;
      w_full = 0;
      w_partial = 0;
      r_full = 0;
      r_partial = 0;
      conversions_mask = 0;
      output_flags = 0;
      output_file = "child_stdout_zero";

      if (obuf)
        {
          free(obuf);
          obuf = NULL;
        }
      obuf = (char*)malloc(output_blocksize);
      if (!obuf) _exit(123);
      memset(obuf, 'B', output_blocksize);

      int p[2];
      if (pipe(p) != 0) _exit(124);

      /* Set write end non-blocking and redirect stdout there. */
      int flags = fcntl(p[1], F_GETFL);
      if (flags < 0) _exit(125);
      if (fcntl(p[1], F_SETFL, flags | O_NONBLOCK) != 0) _exit(126);

      if (dup2(p[1], STDOUT_FILENO) < 0) _exit(127);
      close(p[1]); /* stdout now points to pipe write end */

      /* Keep read end open so the pipe doesn't get closed */
      /* Fill the pipe completely; subsequent writes will EAGAIN immediately. */
      fill_fd_nonblock_to_full(STDOUT_FILENO);

      /* Now call target; this should diagnose and exit(EXIT_FAILURE). */
      write_output();

      /* Should not reach here */
      _exit(200);
    }
  else
    {
      int status = 0;
      waitpid(pid, &status, 0);
      TEST_ASSERT_TRUE(WIFEXITED(status));
      TEST_ASSERT_EQUAL_INT(EXIT_FAILURE, WEXITSTATUS(status));
    }
}

/* Test: short write (nonzero bytes) increments w_partial and exits failure; verify via stderr stats */
void test_write_output_short_write_nonzero_increments_partial(void)
{
  fflush(stdout);
    fflush(stderr);
    pid_t pid = fork();
  TEST_ASSERT_TRUE(pid >= 0);

  if (pid == 0)
    {
      /* Child setup */
      output_blocksize = 4096;
      oc = 5;
      w_bytes = 0;
      w_full = 0;
      w_partial = 0;
      r_full = 0;
      r_partial = 0;
      conversions_mask = 0;
      output_flags = 0;
      output_file = "child_stdout_partial";

      if (obuf)
        {
          free(obuf);
          obuf = NULL;
        }
      obuf = (char*)malloc(output_blocksize);
      if (!obuf) _exit(223);
      memset(obuf, 'C', output_blocksize);

      int pout[2];
      if (pipe(pout) != 0) _exit(224);
      /* Redirect stderr to pout[1] to capture stats */
      if (dup2(pout[1], STDERR_FILENO) < 0) _exit(225);
      close(pout[1]); /* stderr now pipe write end */
      int perr_read = pout[0];

      int p[2];
      if (pipe(p) != 0) _exit(226);

      /* Non-blocking write end as stdout */
      int flags = fcntl(p[1], F_GETFL);
      if (flags < 0) _exit(227);
      if (fcntl(p[1], F_SETFL, flags | O_NONBLOCK) != 0) _exit(228);
      if (dup2(p[1], STDOUT_FILENO) < 0) _exit(229);
      close(p[1]); /* stdout -> pipe write end */

      /* Fill pipe to full, then drain 1 byte to allow a partial write. */
      fill_fd_nonblock_to_full(STDOUT_FILENO);
      char tmp;
      ssize_t r = read(p[0], &tmp, 1); /* create 1-byte space */
      (void)r; /* ignore failures: best effort */
      /* Keep p[0] open to keep pipe valid. */

      /* Now call target; expected to increment w_partial (nwritten != 0) and then exit failure. */
      write_output();

      _exit(240);
    }
  else
    {
      /* Parent: capture child's stderr */
      int status = 0;
      waitpid(pid, &status, 0);
      TEST_ASSERT_TRUE(WIFEXITED(status));
      TEST_ASSERT_EQUAL_INT(EXIT_FAILURE, WEXITSTATUS(status));

      /* We cannot directly read child's stderr here since we didn't set up in parent;
         Instead, rely on child having redirected its stderr to a pipe and the stats printed
         before exit. To actually capture, we'd need the parent end of the pipe.
         Since file descriptors do not cross fork automatically unless inherited,
         we need to restructure: We'll accept verifying exit status only if
         pipe capture isn't feasible.
         However, to keep verification of w_partial increment, we redo a local reproduction
         but this time we will run the child with a dedicated stderr pipe that we can read.
       */
      /* Re-run with proper pipe capture */
      int cap[2];
      TEST_ASSERT_EQUAL_INT(0, pipe(cap));
      pid_t pid2 = fork();
      TEST_ASSERT_TRUE(pid2 >= 0);
      if (pid2 == 0)
        {
          /* Child 2: same setup as above but ensure parent can read stderr from cap[0] */
          output_blocksize = 4096;
          oc = 5;
          w_bytes = 0;
          w_full = 0;
          w_partial = 0;
          r_full = 0;
          r_partial = 0;
          conversions_mask = 0;
          output_flags = 0;
          output_file = "child2_stdout_partial";

          if (obuf)
            {
              free(obuf);
              obuf = NULL;
            }
          obuf = (char*)malloc(output_blocksize);
          if (!obuf) _exit(253);
          memset(obuf, 'D', output_blocksize);

          /* Connect stderr to cap[1] */
          if (dup2(cap[1], STDERR_FILENO) < 0) _exit(254);
          close(cap[1]);
          /* Set up non-blocking stdout pipe */
          int p[2];
          if (pipe(p) != 0) _exit(255);
          int flags = fcntl(p[1], F_GETFL);
          if (flags < 0) _exit(256);
          if (fcntl(p[1], F_SETFL, flags | O_NONBLOCK) != 0) _exit(257);
          if (dup2(p[1], STDOUT_FILENO) < 0) _exit(258);
          close(p[1]); /* stdout is pipe write end now */

          fill_fd_nonblock_to_full(STDOUT_FILENO);
          char tmp;
          (void)read(p[0], &tmp, 1); /* free 1 byte for a partial write */

          write_output();

          _exit(260);
        }
      else
        {
          close(cap[1]); /* parent reads from cap[0] */
          char buffer[8192];
          ssize_t total = 0;
          for (;;)
            {
              ssize_t nr = read(cap[0], buffer + total, sizeof(buffer) - total - 1);
              if (nr > 0)
                {
                  total += nr;
                  if ((size_t)total >= sizeof(buffer) - 1)
                    break;
                }
              else if (nr == 0)
                break;
              else if (errno == EINTR)
                continue;
              else
                break;
            }
          if (total < (ssize_t)sizeof(buffer))
            buffer[total] = '\0';
          else
            buffer[sizeof(buffer) - 1] = '\0';
          close(cap[0]);

          int st2 = 0;
          waitpid(pid2, &st2, 0);
          TEST_ASSERT_TRUE(WIFEXITED(st2));
          TEST_ASSERT_EQUAL_INT(EXIT_FAILURE, WEXITSTATUS(st2));

          long long full = -1, partial = -1;
          /* We expect "0+1 records out" */
          TEST_ASSERT_EQUAL_INT(0, parse_records_out(buffer, &full, &partial));
          TEST_ASSERT_EQUAL_INT64(0, full);
          TEST_ASSERT_EQUAL_INT64(1, partial);
        }
    }
}

int main(void)
{
  UNITY_BEGIN();
  RUN_TEST(test_write_output_success_basic);
  RUN_TEST(test_write_output_short_write_zero_bytes_exits);
  RUN_TEST(test_write_output_short_write_nonzero_increments_partial);
  return UNITY_END();
}