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

/* Globals from the program (file-scope static in the same translation unit). */
extern char line_end;         /* declared above in the program, static, visible here */
extern bool print_headers;    /* visible here */
extern bool presume_input_pipe; /* visible here */

/* Helper: create a temp file with given content, return read-only fd and path for cleanup. */
static int make_temp_file_with_content(const char* tag, const void* data, size_t len,
                                       char* out_path, size_t out_path_size)
{
  char tmpl[PATH_MAX];
  snprintf(tmpl, sizeof(tmpl), "/tmp/head_test_%s_XXXXXX", tag ? tag : "tmp");
  if (out_path && out_path_size > 0)
    out_path[0] = '\0';

  int wfd = mkstemp(tmpl);
  if (wfd < 0)
    return -1;

  ssize_t wr = 0;
  const char* p = (const char*)data;
  size_t left = len;
  while (left > 0) {
    ssize_t n = write(wfd, p + wr, left);
    if (n < 0) {
      close(wfd);
      unlink(tmpl);
      return -1;
    }
    wr += n;
    left -= n;
  }
  /* Rewind and close writer, reopen read-only to emulate a normal input. */
  if (lseek(wfd, 0, SEEK_SET) < 0) { /* Not strictly needed. */ }
  close(wfd);

  if (out_path && out_path_size > 0)
    snprintf(out_path, out_path_size, "%s", tmpl);

  int rfd = open(tmpl, O_RDONLY);
  if (rfd < 0) {
    unlink(tmpl);
    return -1;
  }
  return rfd;
}

/* Stdout capture helpers. Do NOT assert while stdout is redirected. */
typedef struct {
  int saved_stdout_fd;
  int cap_fd;
  char path[PATH_MAX];
  bool active;
} capture_ctx_t;

static bool capture_stdout_begin(capture_ctx_t* ctx, const char* tag)
{
  if (!ctx) return false;
  memset(ctx, 0, sizeof(*ctx));
  ctx->saved_stdout_fd = dup(STDOUT_FILENO);
  if (ctx->saved_stdout_fd < 0) return false;

  char tmpl[PATH_MAX];
  snprintf(tmpl, sizeof(tmpl), "/tmp/head_cap_%s_XXXXXX", tag ? tag : "out");
  int fd = mkstemp(tmpl);
  if (fd < 0) {
    close(ctx->saved_stdout_fd);
    return false;
  }
  snprintf(ctx->path, sizeof(ctx->path), "%s", tmpl);
  /* Redirect stdout to fd */
  fflush(stdout);
  if (dup2(fd, STDOUT_FILENO) < 0) {
    close(fd);
    unlink(ctx->path);
    close(ctx->saved_stdout_fd);
    return false;
  }
  ctx->cap_fd = fd;
  ctx->active = true;
  return true;
}

static bool capture_stdout_end(capture_ctx_t* ctx, char** out_buf, size_t* out_len)
{
  if (!ctx || !ctx->active) return false;

  /* Ensure all data is written */
  fflush(stdout);
  fsync(STDOUT_FILENO);

  /* Determine size, read into buffer */
  off_t end = lseek(ctx->cap_fd, 0, SEEK_END);
  if (end < 0) end = 0;
  if (lseek(ctx->cap_fd, 0, SEEK_SET) < 0) { /* ignore */ }

  size_t len = (size_t)end;
  char* buf = (char*)malloc(len + 1);
  if (!buf) {
    /* Still restore stdout even on failure */
    dup2(ctx->saved_stdout_fd, STDOUT_FILENO);
    close(ctx->saved_stdout_fd);
    close(ctx->cap_fd);
    unlink(ctx->path);
    ctx->active = false;
    return false;
  }

  size_t rd_total = 0;
  while (rd_total < len) {
    ssize_t n = read(ctx->cap_fd, buf + rd_total, len - rd_total);
    if (n < 0) {
      free(buf);
      dup2(ctx->saved_stdout_fd, STDOUT_FILENO);
      close(ctx->saved_stdout_fd);
      close(ctx->cap_fd);
      unlink(ctx->path);
      ctx->active = false;
      return false;
    }
    if (n == 0) break;
    rd_total += n;
  }
  buf[rd_total] = '\0'; /* NUL-terminate for convenience in string compares */

  /* Restore stdout */
  dup2(ctx->saved_stdout_fd, STDOUT_FILENO);
  close(ctx->saved_stdout_fd);
  close(ctx->cap_fd);
  unlink(ctx->path);
  ctx->active = false;

  if (out_buf) *out_buf = buf; else free(buf);
  if (out_len) *out_len = rd_total;
  return true;
}

/* setUp/tearDown for Unity */
void setUp(void) {
  /* Default to newline-separated lines. */
  line_end = '\n';
  print_headers = false;
  presume_input_pipe = false;
}

void tearDown(void) {
  /* Nothing */
}

/* Prototype for the target under test (already defined above in the same TU). */
static bool head (char const *filename, int fd, uintmax_t n_units, bool count_lines,
                  bool elide_from_end);

/* Test 1: Header printing on first call and line output */
void test_head_headers_first_call_prints_banner_and_content(void)
{
  const char* content = "12345\n67890\n";
  char path[PATH_MAX];
  int fd = make_temp_file_with_content("hdr", content, strlen(content), path, sizeof(path));
  TEST_ASSERT_MESSAGE(fd >= 0, "Failed to create temp file");
  /* Ensure cleanup */
  unlink(path);

  print_headers = true; /* enable header output */
  line_end = '\n';

  capture_ctx_t cap;
  TEST_ASSERT_TRUE_MESSAGE(capture_stdout_begin(&cap, "hdr"), "Failed to start capture");
  bool ok = head("file1", fd, 1, true, false); /* first 1 line */
  char* out = NULL; size_t out_len = 0;
  TEST_ASSERT_TRUE_MESSAGE(capture_stdout_end(&cap, &out, &out_len), "Failed to finish capture");

  close(fd);

  TEST_ASSERT_TRUE(ok);
  const char* expected = "==> file1 <==\n12345\n";
  TEST_ASSERT_EQUAL_size_t(strlen(expected), out_len);
  TEST_ASSERT_EQUAL_INT(0, memcmp(expected, out, out_len));

  free(out);
}

/* Test 2: Basic bytes from start */
void test_head_bytes_basic(void)
{
  const char* content = "abcdef";
  char path[PATH_MAX];
  int fd = make_temp_file_with_content("bytes", content, strlen(content), path, sizeof(path));
  TEST_ASSERT_MESSAGE(fd >= 0, "Failed to create temp file");
  unlink(path);

  capture_ctx_t cap;
  TEST_ASSERT_TRUE(capture_stdout_begin(&cap, "bytes"));
  bool ok = head("bytes_file", fd, 3, false, false); /* first 3 bytes */
  char* out = NULL; size_t out_len = 0;
  TEST_ASSERT_TRUE(capture_stdout_end(&cap, &out, &out_len));

  close(fd);

  TEST_ASSERT_TRUE(ok);
  const char* expected = "abc";
  TEST_ASSERT_EQUAL_size_t(3, out_len);
  TEST_ASSERT_EQUAL_INT(0, memcmp(expected, out, out_len));

  free(out);
}

/* Test 3: Basic lines from start */
void test_head_lines_basic_newline(void)
{
  const char* content = "a\nb\nc\n";
  char path[PATH_MAX];
  int fd = make_temp_file_with_content("lines", content, strlen(content), path, sizeof(path));
  TEST_ASSERT_MESSAGE(fd >= 0, "Failed to create temp file");
  unlink(path);

  line_end = '\n';

  capture_ctx_t cap;
  TEST_ASSERT_TRUE(capture_stdout_begin(&cap, "lines"));
  bool ok = head("lines_file", fd, 2, true, false); /* first 2 lines */
  char* out = NULL; size_t out_len = 0;
  TEST_ASSERT_TRUE(capture_stdout_end(&cap, &out, &out_len));

  close(fd);

  TEST_ASSERT_TRUE(ok);
  const char* expected = "a\nb\n";
  TEST_ASSERT_EQUAL_size_t(strlen(expected), out_len);
  TEST_ASSERT_EQUAL_INT(0, memcmp(expected, out, out_len));

  free(out);
}

/* Test 4: Elide bytes from the end (seekable path) */
void test_head_bytes_elide_end_seekable(void)
{
  const char* content = "abcdef";
  char path[PATH_MAX];
  int fd = make_temp_file_with_content("elideb", content, strlen(content), path, sizeof(path));
  TEST_ASSERT_MESSAGE(fd >= 0, "Failed to create temp file");
  unlink(path);

  /* elide last 2 bytes => output "abcd" */
  capture_ctx_t cap;
  TEST_ASSERT_TRUE(capture_stdout_begin(&cap, "elideb"));
  bool ok = head("eb_file", fd, 2, false, true);
  char* out = NULL; size_t out_len = 0;
  TEST_ASSERT_TRUE(capture_stdout_end(&cap, &out, &out_len));
  TEST_ASSERT_TRUE(ok);
  TEST_ASSERT_EQUAL_size_t(4, out_len);
  TEST_ASSERT_EQUAL_INT(0, memcmp("abcd", out, 4));
  free(out);

  /* elide more than file size => produce nothing, still success */
  lseek(fd, 0, SEEK_SET);
  capture_ctx_t cap2;
  TEST_ASSERT_TRUE(capture_stdout_begin(&cap2, "elideb2"));
  ok = head("eb_file", fd, 99, false, true);
  out = NULL; out_len = 0;
  TEST_ASSERT_TRUE(capture_stdout_end(&cap2, &out, &out_len));
  TEST_ASSERT_TRUE(ok);
  TEST_ASSERT_EQUAL_size_t(0, out_len);
  free(out);

  close(fd);
}

/* Test 5: Elide lines from the end using pipe path (force via presume_input_pipe) */
void test_head_lines_elide_end_pipe(void)
{
  const char* content = "a\nb\nc\nd\n";
  char path[PATH_MAX];
  int fd = make_temp_file_with_content("elidel", content, strlen(content), path, sizeof(path));
  TEST_ASSERT_MESSAGE(fd >= 0, "Failed to create temp file");
  unlink(path);

  presume_input_pipe = true; /* force pipe algorithm */
  line_end = '\n';

  capture_ctx_t cap;
  TEST_ASSERT_TRUE(capture_stdout_begin(&cap, "elidel"));
  bool ok = head("elidel_file", fd, 2, true, true); /* all but last 2 lines => a\nb\n */
  char* out = NULL; size_t out_len = 0;
  TEST_ASSERT_TRUE(capture_stdout_end(&cap, &out, &out_len));

  close(fd);

  TEST_ASSERT_TRUE(ok);
  const char* expected = "a\nb\n";
  TEST_ASSERT_EQUAL_size_t(strlen(expected), out_len);
  TEST_ASSERT_EQUAL_INT(0, memcmp(expected, out, out_len));

  free(out);
}

/* Test 6: Zero-terminated lines (embedded NULs) */
void test_head_lines_zero_terminated(void)
{
  const unsigned char data[] = { 'a','a','\0','b','b','b','b','\0','c','c','c','\0' };
  char path[PATH_MAX];
  int fd = make_temp_file_with_content("nul", data, sizeof(data), path, sizeof(path));
  TEST_ASSERT_MESSAGE(fd >= 0, "Failed to create temp file");
  unlink(path);

  line_end = '\0'; /* NUL-terminated records */

  capture_ctx_t cap;
  TEST_ASSERT_TRUE(capture_stdout_begin(&cap, "nul"));
  bool ok = head("nul_file", fd, 2, true, false); /* first 2 NUL-terminated records */
  char* out = NULL; size_t out_len = 0;
  TEST_ASSERT_TRUE(capture_stdout_end(&cap, &out, &out_len));

  close(fd);

  TEST_ASSERT_TRUE(ok);
  const unsigned char expected[] = { 'a','a','\0','b','b','b','b','\0' };
  TEST_ASSERT_EQUAL_size_t(sizeof(expected), out_len);
  TEST_ASSERT_EQUAL_INT(0, memcmp(expected, out, out_len));

  free(out);
}

/* Test 7: Infinite elide (UINTMAX_MAX) produces no output and returns true */
void test_head_infinite_elide_returns_true_and_no_output(void)
{
  const char* content = "data\n";
  char path[PATH_MAX];
  int fd = make_temp_file_with_content("inf", content, strlen(content), path, sizeof(path));
  TEST_ASSERT_MESSAGE(fd >= 0, "Failed to create temp file");
  unlink(path);

  capture_ctx_t cap;
  TEST_ASSERT_TRUE(capture_stdout_begin(&cap, "inf"));
  bool ok = head("inf_file", fd, UINTMAX_MAX, true, true);
  char* out = NULL; size_t out_len = 0;
  TEST_ASSERT_TRUE(capture_stdout_end(&cap, &out, &out_len));

  close(fd);

  TEST_ASSERT_TRUE(ok);
  TEST_ASSERT_EQUAL_size_t(0, out_len);

  free(out);
}

/* Test 8: Read error returns false (use closed fd) */
void test_head_read_error_returns_false(void)
{
  const char* content = "error";
  char path[PATH_MAX];
  int fd = make_temp_file_with_content("err", content, strlen(content), path, sizeof(path));
  TEST_ASSERT_MESSAGE(fd >= 0, "Failed to create temp file");
  unlink(path);

  close(fd); /* provoke read error */

  bool ok = head("err_file", fd, 3, false, false);
  TEST_ASSERT_FALSE(ok);
}

int main(void)
{
  UNITY_BEGIN();
  /* Ensure header test runs first (write_header keeps internal first-file state) */
  RUN_TEST(test_head_headers_first_call_prints_banner_and_content);
  RUN_TEST(test_head_bytes_basic);
  RUN_TEST(test_head_lines_basic_newline);
  RUN_TEST(test_head_bytes_elide_end_seekable);
  RUN_TEST(test_head_lines_elide_end_pipe);
  RUN_TEST(test_head_lines_zero_terminated);
  RUN_TEST(test_head_infinite_elide_returns_true_and_no_output);
  RUN_TEST(test_head_read_error_returns_false);
  return UNITY_END();
}