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

/* The function under test is static in the included program source and is
   available here because this test file is included into that translation unit:
   static bool head_bytes (char const *filename, int fd, uintmax_t bytes_to_write);
*/

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

/* Helper: create a temporary file with given content.
   Returns an open fd positioned at start, or -1 on failure. */
static int create_temp_file_with_content(const void *data, size_t len)
{
  char tmpl[] = "/tmp/head_bytes_test_XXXXXX";
  int fd = mkstemp(tmpl);
  if (fd < 0) {
    return -1;
  }

  /* Unlink the file so it is removed automatically. */
  unlink(tmpl);

  /* Write content if any. */
  size_t off = 0;
  while (off < len) {
    ssize_t w = write(fd, (const char*)data + off, len - off);
    if (w < 0) {
      close(fd);
      return -1;
    }
    off += (size_t)w;
  }

  /* Rewind to start for reading. */
  if (lseek(fd, 0, SEEK_SET) < 0) {
    close(fd);
    return -1;
  }

  return fd;
}

/* Helper: capture stdout output of head_bytes while calling it.
   IMPORTANT: No Unity asserts in this function while stdout is redirected. */
typedef struct {
  char *data;
  size_t len;
  bool func_ret;
} CaptureResult;

static CaptureResult capture_head_bytes_output(int fd, uintmax_t n_bytes, const char *filename)
{
  CaptureResult out = { NULL, 0, false };
  int saved_stdout = dup(STDOUT_FILENO);
  if (saved_stdout < 0) {
    return out;
  }

  int pipefd[2];
  if (pipe(pipefd) != 0) {
    close(saved_stdout);
    return out;
  }

  fflush(stdout);

  /* Redirect stdout to pipe write end */
  if (dup2(pipefd[1], STDOUT_FILENO) < 0) {
    close(saved_stdout);
    close(pipefd[0]);
    close(pipefd[1]);
    return out;
  }
  /* Close our copy of the write end; stdout now refers to it. */
  close(pipefd[1]);

  /* Call the function under test; do not use Unity asserts here. */
  bool r = head_bytes(filename, fd, n_bytes);

  /* Ensure all data is flushed to the pipe. */
  fflush(stdout);

  /* Restore stdout; this closes the old fd 1 (the pipe write end). */
  dup2(saved_stdout, STDOUT_FILENO);
  close(saved_stdout);

  /* Read everything from the pipe read end. */
  size_t cap = 1024;
  char *buf = (char *)malloc(cap);
  if (!buf) {
    close(pipefd[0]);
    out.func_ret = r;
    return out;
  }
  size_t total = 0;
  while (1) {
    if (total == cap) {
      size_t new_cap = cap * 2;
      char *nb = (char *)realloc(buf, new_cap);
      if (!nb) {
        free(buf);
        close(pipefd[0]);
        out.func_ret = r;
        return out;
      }
      buf = nb;
      cap = new_cap;
    }
    ssize_t nr = read(pipefd[0], buf + total, cap - total);
    if (nr < 0) {
      free(buf);
      close(pipefd[0]);
      out.func_ret = r;
      return out;
    }
    if (nr == 0) {
      break;
    }
    total += (size_t)nr;
  }
  close(pipefd[0]);

  out.data = buf;
  out.len = total;
  out.func_ret = r;
  return out;
}

/* Tests */

static void free_capture(CaptureResult *c) {
  if (c && c->data) {
    free(c->data);
    c->data = NULL;
  }
}

/* 1) Request 0 bytes: no output, success */
void test_head_bytes_zero_bytes(void)
{
  const char *content = "abcdef";
  int fd = create_temp_file_with_content(content, strlen(content));
  TEST_ASSERT_MESSAGE(fd >= 0, "Failed to create temp file");

  CaptureResult cap = capture_head_bytes_output(fd, 0, "temp-zero");
  close(fd);

  TEST_ASSERT_TRUE(cap.func_ret);
  TEST_ASSERT_EQUAL_UINT64((uint64_t)0, (uint64_t)cap.len);
  free_capture(&cap);
}

/* 2) Request less than file size: partial output, success */
void test_head_bytes_less_than_file_size(void)
{
  const char *content = "Hello, Unity!";
  size_t content_len = strlen(content);
  int fd = create_temp_file_with_content(content, content_len);
  TEST_ASSERT_MESSAGE(fd >= 0, "Failed to create temp file");

  uintmax_t req = 5;
  CaptureResult cap = capture_head_bytes_output(fd, req, "temp-partial");
  close(fd);

  TEST_ASSERT_TRUE(cap.func_ret);
  TEST_ASSERT_EQUAL_UINT64((uint64_t)req, (uint64_t)cap.len);
  TEST_ASSERT_EQUAL_INT(0, memcmp(content, cap.data, (size_t)req));
  free_capture(&cap);
}

/* 3) Request exactly file size: full output, success */
void test_head_bytes_exact_file_size(void)
{
  const char *content = "exactsize";
  size_t content_len = strlen(content);
  int fd = create_temp_file_with_content(content, content_len);
  TEST_ASSERT_MESSAGE(fd >= 0, "Failed to create temp file");

  CaptureResult cap = capture_head_bytes_output(fd, content_len, "temp-exact");
  close(fd);

  TEST_ASSERT_TRUE(cap.func_ret);
  TEST_ASSERT_EQUAL_UINT64((uint64_t)content_len, (uint64_t)cap.len);
  TEST_ASSERT_EQUAL_INT(0, memcmp(content, cap.data, content_len));
  free_capture(&cap);
}

/* 4) Request more than file size: output full file, success (unexpected EOF tolerated) */
void test_head_bytes_more_than_file_size(void)
{
  const char *content = "short";
  size_t content_len = strlen(content);
  int fd = create_temp_file_with_content(content, content_len);
  TEST_ASSERT_MESSAGE(fd >= 0, "Failed to create temp file");

  uintmax_t req = content_len + 10;
  CaptureResult cap = capture_head_bytes_output(fd, req, "temp-more");
  close(fd);

  TEST_ASSERT_TRUE(cap.func_ret);
  TEST_ASSERT_EQUAL_UINT64((uint64_t)content_len, (uint64_t)cap.len);
  TEST_ASSERT_EQUAL_INT(0, memcmp(content, cap.data, content_len));
  free_capture(&cap);
}

/* 5) Empty file with nonzero request: no output, success */
void test_head_bytes_empty_file_request_nonzero(void)
{
  int fd = create_temp_file_with_content("", 0);
  TEST_ASSERT_MESSAGE(fd >= 0, "Failed to create temp file");

  CaptureResult cap = capture_head_bytes_output(fd, 100, "temp-empty");
  close(fd);

  TEST_ASSERT_TRUE(cap.func_ret);
  TEST_ASSERT_EQUAL_UINT64((uint64_t)0, (uint64_t)cap.len);
  free_capture(&cap);
}

/* 6) Multi-buffer read: request > BUFSIZ ensuring multiple reads/writes */
void test_head_bytes_large_multi_buffer(void)
{
  size_t biglen = (size_t)(3 * BUFSIZ + 123);
  char *bigbuf = (char *)malloc(biglen);
  TEST_ASSERT_NOT_NULL_MESSAGE(bigbuf, "Allocation failed for big buffer");
  /* Fill with a pattern */
  for (size_t i = 0; i < biglen; i++) {
    bigbuf[i] = (char)('A' + (int)(i % 26));
  }

  int fd = create_temp_file_with_content(bigbuf, biglen);
  TEST_ASSERT_MESSAGE(fd >= 0, "Failed to create temp file");

  uintmax_t req = (uintmax_t)(2 * BUFSIZ + 10);
  CaptureResult cap = capture_head_bytes_output(fd, req, "temp-large");
  close(fd);

  TEST_ASSERT_TRUE(cap.func_ret);
  TEST_ASSERT_EQUAL_UINT64((uint64_t)req, (uint64_t)cap.len);
  TEST_ASSERT_EQUAL_INT(0, memcmp(bigbuf, cap.data, (size_t)req));

  free_capture(&cap);
  free(bigbuf);
}

/* 7) Read error: invalid fd should return false and produce no stdout output */
void test_head_bytes_read_error_returns_false_and_no_output(void)
{
  int invalid_fd = -1;
  CaptureResult cap = capture_head_bytes_output(invalid_fd, 100, "invalid-fd");
  TEST_ASSERT_FALSE(cap.func_ret);
  TEST_ASSERT_EQUAL_UINT64((uint64_t)0, (uint64_t)cap.len);
  free_capture(&cap);
}

int main(void)
{
  UNITY_BEGIN();
  RUN_TEST(test_head_bytes_zero_bytes);
  RUN_TEST(test_head_bytes_less_than_file_size);
  RUN_TEST(test_head_bytes_exact_file_size);
  RUN_TEST(test_head_bytes_more_than_file_size);
  RUN_TEST(test_head_bytes_empty_file_request_nonzero);
  RUN_TEST(test_head_bytes_large_multi_buffer);
  RUN_TEST(test_head_bytes_read_error_returns_false_and_no_output);
  return UNITY_END();
}