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

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

/* The test file is included into the du.c translation unit, so we can
   access its internal symbols directly (functions, statics, flags). */

/* Helper: create a temporary directory. Caller must free returned string. */
static char *make_temp_dir(void)
{
  char *tmpl = malloc(64);
  if (!tmpl) return NULL;
  snprintf(tmpl, 64, "/tmp/du_ut_%ld_XXXXXX", (long)getpid());
  if (!mkdtemp(tmpl))
    {
      free(tmpl);
      return NULL;
    }
  return tmpl; /* Keep allocated for later free */
}

/* Helper: join path components into buffer. buffer must be big enough. */
static void join_path(char *buf, size_t bufsz, const char *a, const char *b)
{
  size_t la = strlen(a);
  snprintf(buf, bufsz, "%s%s%s", a, (la && a[la-1] == '/') ? "" : "/", b);
}

/* Helper: create a file with specified size filled with 'X' bytes. */
static int create_file_with_size(const char *path, size_t size)
{
  int fd = open(path, O_CREAT | O_TRUNC | O_WRONLY, 0644);
  if (fd < 0) return -1;
  size_t remaining = size;
  char buf[4096];
  memset(buf, 'X', sizeof buf);
  while (remaining > 0)
    {
      size_t chunk = remaining < sizeof buf ? remaining : sizeof buf;
      ssize_t w = write(fd, buf, chunk);
      if (w < 0)
        {
          int e = errno;
          close(fd);
          errno = e;
          return -1;
        }
      remaining -= (size_t)w;
    }
  if (close(fd) < 0) return -1;
  return 0;
}

/* Helper: recursively remove directory tree. */
static int remove_tree(const char *path)
{
  struct stat st;
  if (lstat(path, &st) != 0)
    return -1;

  if (S_ISDIR(st.st_mode))
    {
      DIR *d = opendir(path);
      if (!d) return -1;
      struct dirent *de;
      while ((de = readdir(d)) != NULL)
        {
          if (strcmp(de->d_name, ".") == 0 || strcmp(de->d_name, "..") == 0)
            continue;
          char child[PATH_MAX];
          join_path(child, sizeof child, path, de->d_name);
          if (remove_tree(child) != 0)
            {
              /* Attempt to continue cleanup even on failure */
            }
        }
      closedir(d);
      if (rmdir(path) != 0)
        return -1;
      return 0;
    }
  else
    {
      return unlink(path);
    }
}

/* Reset du global state to deterministic defaults for testing. */
static void reset_du_state(void)
{
  /* flags declared static in du.c; accessible here */
  opt_all = false;
  apparent_size = true;      /* sizes count as st_size */
  opt_count_all = false;
  hash_all = false;
  opt_nul_terminate_output = false;
  print_grand_total = false;
  opt_separate_dirs = false;
  max_depth = IDX_MAX;
  opt_threshold = 0;
  human_output_opts = 0;
  opt_inodes = false;        /* print sizes, not inode counts */
  opt_time = false;
  time_format = NULL;
  time_style = NULL;
  output_block_size = 1;     /* print raw number of bytes */
  exclude = NULL;
  duinfo_init(&tot_dui);
  prev_level = 0;
}

/* Structure for captured output */
struct CaptureResult {
  char *data;
  size_t len;
  bool ok;
};

/* Capture stdout produced by du_files(files, bit_flags).
   While stdout is redirected, do not call Unity asserts. */
static struct CaptureResult run_du_and_capture(char **files, int bit_flags)
{
  struct CaptureResult res = { NULL, 0, false };

  fflush(stdout);
  int pipefd[2];
  if (pipe(pipefd) != 0)
    {
      return res;
    }

  int saved = dup(STDOUT_FILENO);
  if (saved < 0)
    {
      close(pipefd[0]);
      close(pipefd[1]);
      return res;
    }

  if (dup2(pipefd[1], STDOUT_FILENO) < 0)
    {
      close(saved);
      close(pipefd[0]);
      close(pipefd[1]);
      return res;
    }
  close(pipefd[1]); /* writer now is stdout */

  bool ok = du_files(files, bit_flags);

  fflush(stdout);

  /* Restore stdout before any assertions or further output */
  dup2(saved, STDOUT_FILENO);
  close(saved);

  /* Read all data */
  char *buf = NULL;
  size_t cap = 0;
  size_t len = 0;
  char tmp[4096];
  ssize_t r;
  while ((r = read(pipefd[0], tmp, sizeof tmp)) > 0)
    {
      if (len + (size_t)r > cap)
        {
          size_t newcap = cap ? cap * 2 : 8192;
          while (newcap < len + (size_t)r) newcap *= 2;
          char *nb = realloc(buf, newcap);
          if (!nb)
            {
              free(buf);
              close(pipefd[0]);
              res.data = NULL;
              res.len = 0;
              res.ok = ok;
              return res;
            }
          buf = nb;
          cap = newcap;
        }
      memcpy(buf + len, tmp, (size_t)r);
      len += (size_t)r;
    }
  close(pipefd[0]);

  res.data = buf;
  res.len = len;
  res.ok = ok;
  return res;
}

/* Parse the captured buffer and find the size printed for a specific path.
   Lines are delimited by term ('\n' or '\0'). Returns UINTMAX_MAX+1 on failure.
*/
static unsigned long long find_size_for_path(const char *buf, size_t len,
                                             char term, const char *path)
{
  size_t i = 0;
  size_t pathlen = strlen(path);
  while (i < len)
    {
      size_t line_start = i;
      /* find end */
      while (i < len && buf[i] != term) i++;
      size_t line_end = i; /* exclusive */
      if (line_end > line_start)
        {
          /* Tab-separated: "<number>\t<path>" */
          const char *tab = memchr(buf + line_start, '\t', line_end - line_start);
          if (tab)
            {
              size_t num_len = (size_t)(tab - (buf + line_start));
              size_t p_len = line_end - (size_t)(tab - buf) - 1;
              const char *p = tab + 1;
              if (p_len == pathlen && memcmp(p, path, pathlen) == 0)
                {
                  char numstr[64];
                  size_t cpy = num_len < sizeof numstr - 1 ? num_len : sizeof numstr - 1;
                  memcpy(numstr, buf + line_start, cpy);
                  numstr[cpy] = '\0';
                  char *endp = NULL;
                  unsigned long long v = strtoull(numstr, &endp, 10);
                  if (endp && *endp == '\0')
                    return v;
                }
            }
        }
      /* advance past term */
      if (i < len) i++;
    }
  return (unsigned long long)(~0ULL); /* sentinel: not found */
}

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

/* Tests */

static void test_du_files_empty_input_no_output(void)
{
  reset_du_state();

  char *files[] = { NULL };
  struct CaptureResult cr = run_du_and_capture(files, FTS_PHYSICAL | FTS_NOCHDIR);

  TEST_ASSERT_TRUE(cr.ok);
  TEST_ASSERT_EQUAL_UINT64(0, cr.len);
  free(cr.data);
}

static void test_du_files_single_file_apparent_size(void)
{
  reset_du_state();

  char *td = make_temp_dir();
  TEST_ASSERT_NOT_NULL(td);

  char fpath[PATH_MAX];
  join_path(fpath, sizeof fpath, td, "one.txt");
  TEST_ASSERT_EQUAL_INT(0, create_file_with_size(fpath, 7));

  char *files[] = { fpath, NULL };
  struct CaptureResult cr = run_du_and_capture(files, FTS_PHYSICAL | FTS_NOCHDIR);

  TEST_ASSERT_TRUE(cr.ok);

  /* For newline-terminated output, verify exact size/path line exists */
  unsigned long long sz = find_size_for_path(cr.data, cr.len, '\n', fpath);
  TEST_ASSERT_NOT_EQUAL_UINT64((unsigned long long)(~0ULL), sz);
  TEST_ASSERT_EQUAL_UINT64(7ULL, sz);

  free(cr.data);
  remove_tree(td);
  free(td);
}

static void test_du_files_directory_sums_no_all(void)
{
  reset_du_state();

  char *td = make_temp_dir();
  TEST_ASSERT_NOT_NULL(td);

  /* Create structure: td/a (3B), td/b (4B), td/s/c (5B) */
  char apath[PATH_MAX], bpath[PATH_MAX], sdir[PATH_MAX], cpath[PATH_MAX];
  join_path(apath, sizeof apath, td, "a");
  join_path(bpath, sizeof bpath, td, "b");
  join_path(sdir, sizeof sdir, td, "s");
  TEST_ASSERT_EQUAL_INT(0, mkdir(sdir, 0755));
  join_path(cpath, sizeof cpath, sdir, "c");

  TEST_ASSERT_EQUAL_INT(0, create_file_with_size(apath, 3));
  TEST_ASSERT_EQUAL_INT(0, create_file_with_size(bpath, 4));
  TEST_ASSERT_EQUAL_INT(0, create_file_with_size(cpath, 5));

  char *files[] = { td, NULL };
  struct CaptureResult cr = run_du_and_capture(files, FTS_PHYSICAL | FTS_NOCHDIR);

  TEST_ASSERT_TRUE(cr.ok);

  /* Expect two lines: for s (5) and for td (12). Files not printed without --all. */
  unsigned long long sz_s = find_size_for_path(cr.data, cr.len, '\n', sdir);
  unsigned long long sz_td = find_size_for_path(cr.data, cr.len, '\n', td);

  TEST_ASSERT_NOT_EQUAL_UINT64((unsigned long long)(~0ULL), sz_s);
  TEST_ASSERT_NOT_EQUAL_UINT64((unsigned long long)(~0ULL), sz_td);
  TEST_ASSERT_EQUAL_UINT64(5ULL, sz_s);
  TEST_ASSERT_EQUAL_UINT64(12ULL, sz_td);

  free(cr.data);
  remove_tree(td);
  free(td);
}

static void test_du_files_opt_all_includes_files(void)
{
  reset_du_state();
  opt_all = true;

  char *td = make_temp_dir();
  TEST_ASSERT_NOT_NULL(td);

  char apath[PATH_MAX], bpath[PATH_MAX], sdir[PATH_MAX], cpath[PATH_MAX];
  join_path(apath, sizeof apath, td, "a");
  join_path(bpath, sizeof bpath, td, "b");
  join_path(sdir, sizeof sdir, td, "s");
  TEST_ASSERT_EQUAL_INT(0, mkdir(sdir, 0755));
  join_path(cpath, sizeof cpath, sdir, "c");

  TEST_ASSERT_EQUAL_INT(0, create_file_with_size(apath, 3));
  TEST_ASSERT_EQUAL_INT(0, create_file_with_size(bpath, 4));
  TEST_ASSERT_EQUAL_INT(0, create_file_with_size(cpath, 5));

  char *files[] = { td, NULL };
  struct CaptureResult cr = run_du_and_capture(files, FTS_PHYSICAL | FTS_NOCHDIR);

  TEST_ASSERT_TRUE(cr.ok);

  /* Expect printed sizes for files and directories */
  unsigned long long sz_a = find_size_for_path(cr.data, cr.len, '\n', apath);
  unsigned long long sz_b = find_size_for_path(cr.data, cr.len, '\n', bpath);
  unsigned long long sz_c = find_size_for_path(cr.data, cr.len, '\n', cpath);
  unsigned long long sz_s = find_size_for_path(cr.data, cr.len, '\n', sdir);
  unsigned long long sz_td = find_size_for_path(cr.data, cr.len, '\n', td);

  TEST_ASSERT_EQUAL_UINT64(3ULL, sz_a);
  TEST_ASSERT_EQUAL_UINT64(4ULL, sz_b);
  TEST_ASSERT_EQUAL_UINT64(5ULL, sz_c);
  TEST_ASSERT_EQUAL_UINT64(5ULL, sz_s);
  TEST_ASSERT_EQUAL_UINT64(12ULL, sz_td);

  free(cr.data);
  remove_tree(td);
  free(td);
}

static void test_du_files_max_depth_zero(void)
{
  reset_du_state();
  max_depth = 0;

  char *td = make_temp_dir();
  TEST_ASSERT_NOT_NULL(td);

  char sdir[PATH_MAX], cpath[PATH_MAX];
  join_path(sdir, sizeof sdir, td, "s");
  TEST_ASSERT_EQUAL_INT(0, mkdir(sdir, 0755));
  join_path(cpath, sizeof cpath, sdir, "c");
  TEST_ASSERT_EQUAL_INT(0, create_file_with_size(cpath, 5));

  char *files[] = { td, NULL };
  struct CaptureResult cr = run_du_and_capture(files, FTS_PHYSICAL | FTS_NOCHDIR);

  TEST_ASSERT_TRUE(cr.ok);

  /* Only root directory (td) should be printed */
  unsigned long long sz_td = find_size_for_path(cr.data, cr.len, '\n', td);
  unsigned long long sz_s = find_size_for_path(cr.data, cr.len, '\n', sdir);
  unsigned long long sz_c = find_size_for_path(cr.data, cr.len, '\n', cpath);

  TEST_ASSERT_NOT_EQUAL_UINT64((unsigned long long)(~0ULL), sz_td);
  TEST_ASSERT_EQUAL_UINT64((unsigned long long)(~0ULL), sz_s);
  TEST_ASSERT_EQUAL_UINT64((unsigned long long)(~0ULL), sz_c);
  TEST_ASSERT_EQUAL_UINT64(5ULL, sz_td);

  free(cr.data);
  remove_tree(td);
  free(td);
}

static void test_du_files_threshold_filtering(void)
{
  reset_du_state();
  opt_all = true;
  opt_threshold = 4; /* Only entries with size >= 4 */

  char *td = make_temp_dir();
  TEST_ASSERT_NOT_NULL(td);

  char a[PATH_MAX], b[PATH_MAX], cdir[PATH_MAX], c[PATH_MAX];
  join_path(a, sizeof a, td, "a");
  join_path(b, sizeof b, td, "b");
  join_path(cdir, sizeof cdir, td, "cdir");
  TEST_ASSERT_EQUAL_INT(0, mkdir(cdir, 0755));
  join_path(c, sizeof c, cdir, "c");

  TEST_ASSERT_EQUAL_INT(0, create_file_with_size(a, 3));
  TEST_ASSERT_EQUAL_INT(0, create_file_with_size(b, 4));
  TEST_ASSERT_EQUAL_INT(0, create_file_with_size(c, 5));

  char *files[] = { td, NULL };
  struct CaptureResult cr = run_du_and_capture(files, FTS_PHYSICAL | FTS_NOCHDIR);

  TEST_ASSERT_TRUE(cr.ok);

  /* 'a' (3) excluded; 'b'(4), 'c'(5), 'cdir'(5), and td(9+? actually 9) included */
  unsigned long long sz_a = find_size_for_path(cr.data, cr.len, '\n', a);
  unsigned long long sz_b = find_size_for_path(cr.data, cr.len, '\n', b);
  unsigned long long sz_c = find_size_for_path(cr.data, cr.len, '\n', c);
  unsigned long long sz_cdir = find_size_for_path(cr.data, cr.len, '\n', cdir);
  unsigned long long sz_td = find_size_for_path(cr.data, cr.len, '\n', td);

  TEST_ASSERT_EQUAL_UINT64((unsigned long long)(~0ULL), sz_a);
  TEST_ASSERT_EQUAL_UINT64(4ULL, sz_b);
  TEST_ASSERT_EQUAL_UINT64(5ULL, sz_c);
  TEST_ASSERT_EQUAL_UINT64(5ULL, sz_cdir);
  TEST_ASSERT_EQUAL_UINT64(9ULL, sz_td);

  free(cr.data);
  remove_tree(td);
  free(td);
}

static void test_du_files_nul_terminated_output(void)
{
  reset_du_state();
  opt_nul_terminate_output = true;

  char *td = make_temp_dir();
  TEST_ASSERT_NOT_NULL(td);

  char f[PATH_MAX];
  join_path(f, sizeof f, td, "nulf");
  TEST_ASSERT_EQUAL_INT(0, create_file_with_size(f, 10));

  char *files[] = { f, NULL };
  struct CaptureResult cr = run_du_and_capture(files, FTS_PHYSICAL | FTS_NOCHDIR);

  TEST_ASSERT_TRUE(cr.ok);
  TEST_ASSERT_TRUE(cr.len >= 2);
  /* Last character should be NUL, not newline */
  TEST_ASSERT_EQUAL_UINT8(0, (unsigned char)cr.data[cr.len - 1]);

  /* Verify size/path pair exists when parsing with term '\0' */
  unsigned long long sz = find_size_for_path(cr.data, cr.len, '\0', f);
  TEST_ASSERT_EQUAL_UINT64(10ULL, sz);

  free(cr.data);
  remove_tree(td);
  free(td);
}

int main(void)
{
  UNITY_BEGIN();

  RUN_TEST(test_du_files_empty_input_no_output);
  RUN_TEST(test_du_files_single_file_apparent_size);
  RUN_TEST(test_du_files_directory_sums_no_all);
  RUN_TEST(test_du_files_opt_all_includes_files);
  RUN_TEST(test_du_files_max_depth_zero);
  RUN_TEST(test_du_files_threshold_filtering);
  RUN_TEST(test_du_files_nul_terminated_output);

  return UNITY_END();
}