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

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

/* Accessing internals since this file is included into the same translation unit */
extern bool print_headers;       /* static in the TU, visible due to inclusion */
extern bool presume_input_pipe;  /* static in the TU, visible due to inclusion */
extern bool have_read_stdin;     /* static in the TU, visible due to inclusion */
extern char line_end;            /* static in the TU, visible due to inclusion */

/* Target function prototype (already defined above in the TU) */
static bool head_file (char const *filename, uintmax_t n_units, bool count_lines,
                       bool elide_from_end);

/* Helper: create temp file with content; return malloc'd path string. */
static char* create_temp_file_with_content(const void* data, size_t len)
{
  char tmpl[] = "/tmp/head_file_test_XXXXXX";
  int fd = mkstemp(tmpl);
  if (fd < 0)
  {
    /* We can use TEST_FAIL_MESSAGE here since stdout is not redirected. */
    TEST_FAIL_MESSAGE("mkstemp failed creating temporary file");
    return NULL;
  }

  ssize_t written = 0;
  if (len)
  {
    written = write(fd, data, len);
    if (written < 0 || (size_t)written != len)
    {
      close(fd);
      unlink(tmpl);
      TEST_FAIL_MESSAGE("Failed to write expected content to temporary file");
      return NULL;
    }
  }

  if (close(fd) != 0)
  {
    unlink(tmpl);
    TEST_FAIL_MESSAGE("Failed to close temporary file descriptor");
    return NULL;
  }

  char* path = (char*)malloc(strlen(tmpl) + 1);
  TEST_ASSERT_NOT_NULL(path);
  strcpy(path, tmpl);
  return path;
}

/* Helper: capture stdout while invoking head_file; returns malloc'd buffer and sets out_len and out_ok.
   IMPORTANT: No Unity assertions while stdout is redirected. */
static char* capture_head_file_output(const char* filename,
                                      uintmax_t n_units,
                                      bool count_lines,
                                      bool elide_from_end,
                                      size_t* out_len,
                                      bool* out_ok)
{
  if (!out_len || !out_ok)
    return NULL;

  *out_len = 0;
  *out_ok = false;

  fflush(stdout);
  int saved_stdout = dup(fileno(stdout));
  if (saved_stdout < 0)
  {
    TEST_FAIL_MESSAGE("dup(stdout) failed");
    return NULL;
  }

  char outtmpl[] = "/tmp/head_file_out_XXXXXX";
  int outfd = mkstemp(outtmpl);
  if (outfd < 0)
  {
    close(saved_stdout);
    TEST_FAIL_MESSAGE("mkstemp failed for capture file");
    return NULL;
  }

  if (dup2(outfd, fileno(stdout)) < 0)
  {
    close(outfd);
    unlink(outtmpl);
    close(saved_stdout);
    TEST_FAIL_MESSAGE("dup2 to redirect stdout failed");
    return NULL;
  }

  bool ok = head_file(filename, n_units, count_lines, elide_from_end);

  fflush(stdout);
  off_t endpos = lseek(outfd, 0, SEEK_END);
  if (endpos < 0)
  {
    /* Restore stdout before reporting failure */
    dup2(saved_stdout, fileno(stdout));
    close(saved_stdout);
    close(outfd);
    unlink(outtmpl);
    TEST_FAIL_MESSAGE("lseek on capture fd failed");
    return NULL;
  }
  if (lseek(outfd, 0, SEEK_SET) < 0)
  {
    dup2(saved_stdout, fileno(stdout));
    close(saved_stdout);
    close(outfd);
    unlink(outtmpl);
    TEST_FAIL_MESSAGE("lseek rewind on capture fd failed");
    return NULL;
  }

  size_t len = (size_t)endpos;
  char* buf = (char*)malloc(len + 1);
  if (!buf)
  {
    dup2(saved_stdout, fileno(stdout));
    close(saved_stdout);
    close(outfd);
    unlink(outtmpl);
    TEST_FAIL_MESSAGE("malloc failed for capture buffer");
    return NULL;
  }

  size_t total = 0;
  while (total < len)
  {
    ssize_t r = read(outfd, buf + total, len - total);
    if (r < 0)
    {
      free(buf);
      dup2(saved_stdout, fileno(stdout));
      close(saved_stdout);
      close(outfd);
      unlink(outtmpl);
      TEST_FAIL_MESSAGE("read on capture fd failed");
      return NULL;
    }
    if (r == 0)
      break;
    total += (size_t)r;
  }
  buf[len] = '\0';

  /* Restore stdout */
  dup2(saved_stdout, fileno(stdout));
  close(saved_stdout);
  close(outfd);
  unlink(outtmpl);

  *out_len = len;
  *out_ok = ok;
  return buf;
}

void setUp(void) {
  /* Reset relevant globals for predictable behavior */
  print_headers = false;
  presume_input_pipe = false;
  have_read_stdin = false;
  line_end = '\n';
}

void tearDown(void) {
  /* No persistent resources */
}

/* Tests */

static void test_head_file_bytes_first_N(void)
{
  const char* content = "abcdef\n12345\n";
  char* path = create_temp_file_with_content(content, strlen(content));

  size_t out_len = 0;
  bool ok = false;
  char* out = capture_head_file_output(path, 5, false, false, &out_len, &ok);

  TEST_ASSERT_TRUE(ok);
  TEST_ASSERT_EQUAL_UINT64(5, (uint64_t)out_len);
  TEST_ASSERT_EQUAL_MEMORY("abcde", out, 5);

  free(out);
  unlink(path);
  free(path);
}

static void test_head_file_lines_first_N(void)
{
  const char* content = "line1\nline2\nline3\n";
  char* path = create_temp_file_with_content(content, strlen(content));

  size_t out_len = 0;
  bool ok = false;
  char* out = capture_head_file_output(path, 2, true, false, &out_len, &ok);

  const char* expected = "line1\nline2\n";
  TEST_ASSERT_TRUE(ok);
  TEST_ASSERT_EQUAL_UINT64((uint64_t)strlen(expected), (uint64_t)out_len);
  TEST_ASSERT_EQUAL_MEMORY(expected, out, strlen(expected));

  free(out);
  unlink(path);
  free(path);
}

static void test_head_file_bytes_zero_prints_nothing(void)
{
  const char* content = "abcdef";
  char* path = create_temp_file_with_content(content, strlen(content));

  size_t out_len = 0;
  bool ok = false;
  char* out = capture_head_file_output(path, 0, false, false, &out_len, &ok);

  TEST_ASSERT_TRUE(ok);
  TEST_ASSERT_EQUAL_UINT64(0, (uint64_t)out_len);

  free(out);
  unlink(path);
  free(path);
}

static void test_head_file_lines_zero_prints_nothing(void)
{
  const char* content = "a\nb\nc\n";
  char* path = create_temp_file_with_content(content, strlen(content));

  size_t out_len = 0;
  bool ok = false;
  char* out = capture_head_file_output(path, 0, true, false, &out_len, &ok);

  TEST_ASSERT_TRUE(ok);
  TEST_ASSERT_EQUAL_UINT64(0, (uint64_t)out_len);

  free(out);
  unlink(path);
  free(path);
}

static void test_head_file_elide_bytes_seekable(void)
{
  const char* content = "abcdef";
  char* path = create_temp_file_with_content(content, strlen(content));

  size_t out_len = 0;
  bool ok = false;
  /* elide last 2 bytes -> "abcd" */
  char* out = capture_head_file_output(path, 2, false, true, &out_len, &ok);

  const char* expected = "abcd";
  TEST_ASSERT_TRUE(ok);
  TEST_ASSERT_EQUAL_UINT64((uint64_t)strlen(expected), (uint64_t)out_len);
  TEST_ASSERT_EQUAL_MEMORY(expected, out, strlen(expected));

  free(out);
  unlink(path);
  free(path);
}

static void test_head_file_elide_lines_seekable(void)
{
  const char* content = "a\nb\nc\n";
  char* path = create_temp_file_with_content(content, strlen(content));

  size_t out_len = 0;
  bool ok = false;
  /* elide last 1 line -> "a\nb\n" */
  char* out = capture_head_file_output(path, 1, true, true, &out_len, &ok);

  const char* expected = "a\nb\n";
  TEST_ASSERT_TRUE(ok);
  TEST_ASSERT_EQUAL_UINT64((uint64_t)strlen(expected), (uint64_t)out_len);
  TEST_ASSERT_EQUAL_MEMORY(expected, out, strlen(expected));

  free(out);
  unlink(path);
  free(path);
}

static void test_head_file_elide_bytes_pipe(void)
{
  const char* content = "abcdef";
  char* path = create_temp_file_with_content(content, strlen(content));

  bool prev_presume = presume_input_pipe;
  presume_input_pipe = true; /* Force pipe-like behavior even for regular file */

  size_t out_len = 0;
  bool ok = false;
  /* elide last 3 bytes -> "abc" */
  char* out = capture_head_file_output(path, 3, false, true, &out_len, &ok);

  presume_input_pipe = prev_presume;

  const char* expected = "abc";
  TEST_ASSERT_TRUE(ok);
  TEST_ASSERT_EQUAL_UINT64((uint64_t)strlen(expected), (uint64_t)out_len);
  TEST_ASSERT_EQUAL_MEMORY(expected, out, strlen(expected));

  free(out);
  unlink(path);
  free(path);
}

static void test_head_file_elide_lines_pipe(void)
{
  const char* content = "a\nb\nc\n";
  char* path = create_temp_file_with_content(content, strlen(content));

  bool prev_presume = presume_input_pipe;
  presume_input_pipe = true; /* Force pipe-like code path */

  size_t out_len = 0;
  bool ok = false;
  /* elide last 2 lines -> "a\n" */
  char* out = capture_head_file_output(path, 2, true, true, &out_len, &ok);

  presume_input_pipe = prev_presume;

  const char* expected = "a\n";
  TEST_ASSERT_TRUE(ok);
  TEST_ASSERT_EQUAL_UINT64((uint64_t)strlen(expected), (uint64_t)out_len);
  TEST_ASSERT_EQUAL_MEMORY(expected, out, strlen(expected));

  free(out);
  unlink(path);
  free(path);
}

static void test_head_file_stdin_bytes(void)
{
  const char* content = "hello\nworld\n";
  int p[2];
  TEST_ASSERT_EQUAL_INT(0, pipe(p));

  /* Write content and close write end so reader sees EOF */
  ssize_t w = write(p[1], content, strlen(content));
  TEST_ASSERT_EQUAL_INT((int)strlen(content), (int)w);
  close(p[1]);

  int saved_stdin = dup(STDIN_FILENO);
  TEST_ASSERT_TRUE(saved_stdin >= 0);
  TEST_ASSERT_TRUE(dup2(p[0], STDIN_FILENO) >= 0);
  close(p[0]);

  size_t out_len = 0;
  bool ok = false;
  char* out = capture_head_file_output("-", 5, false, false, &out_len, &ok);

  /* Restore stdin */
  TEST_ASSERT_TRUE(dup2(saved_stdin, STDIN_FILENO) >= 0);
  close(saved_stdin);

  const char* expected = "hello";
  TEST_ASSERT_TRUE(ok);
  TEST_ASSERT_EQUAL_UINT64((uint64_t)strlen(expected), (uint64_t)out_len);
  TEST_ASSERT_EQUAL_MEMORY(expected, out, strlen(expected));

  free(out);
}

static void test_head_file_nonexistent_error(void)
{
  /* No stdout redirection; just check return value. */
  const char* path = "/tmp/this_file_should_not_exist_head_test";
  /* Ensure it doesn't exist */
  unlink(path);

  bool ok = head_file(path, 1, false, false);
  TEST_ASSERT_FALSE(ok);
}

static void test_head_file_elide_infinite(void)
{
  const char* content = "abc\ndef\n";
  char* path = create_temp_file_with_content(content, strlen(content));

  size_t out_len = 0;
  bool ok = false;
  char* out = capture_head_file_output(path, UINTMAX_MAX, false, true, &out_len, &ok);

  TEST_ASSERT_TRUE(ok);
  TEST_ASSERT_EQUAL_UINT64(0, (uint64_t)out_len);

  free(out);
  unlink(path);
  free(path);
}

int main(void)
{
  UNITY_BEGIN();

  RUN_TEST(test_head_file_bytes_first_N);
  RUN_TEST(test_head_file_lines_first_N);
  RUN_TEST(test_head_file_bytes_zero_prints_nothing);
  RUN_TEST(test_head_file_lines_zero_prints_nothing);
  RUN_TEST(test_head_file_elide_bytes_seekable);
  RUN_TEST(test_head_file_elide_lines_seekable);
  RUN_TEST(test_head_file_elide_bytes_pipe);
  RUN_TEST(test_head_file_elide_lines_pipe);
  RUN_TEST(test_head_file_stdin_bytes);
  RUN_TEST(test_head_file_nonexistent_error);
  RUN_TEST(test_head_file_elide_infinite);

  return UNITY_END();
}