#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/types.h>
#include <sys/stat.h>
#include <errno.h>

/* We rely on being included in the same translation unit as head.c,
   so we can access the static function and static globals: head_lines and line_end. */

extern void dummy_reference_to_quiet_linker(void); /* avoid empty TU warning if needed */

/* Prototypes from the program (static in same TU):
   static bool head_lines (char const *filename, int fd, uintmax_t lines_to_write);
   static char line_end;  (we will assign to it) */

/* Helper: create a pipe preloaded with input data for reading. */
static int make_input_pipe(const char *data, size_t len) {
    int p[2];
    if (pipe(p) != 0) return -1;
    ssize_t off = 0;
    while ((size_t)off < len) {
        ssize_t w = write(p[1], data + off, len - off);
        if (w < 0) { close(p[0]); close(p[1]); return -1; }
        off += w;
    }
    close(p[1]);
    return p[0]; /* return read end */
}

/* Helper: start capturing stdout by redirecting it to a pipe. */
static int begin_capture_stdout(int *saved_stdout_fd, int *capture_read_fd) {
    int p[2];
    if (pipe(p) != 0) return -1;
    if (fflush(stdout) != 0) {
        close(p[0]); close(p[1]); return -1;
    }
    int saved = dup(STDOUT_FILENO);
    if (saved < 0) { close(p[0]); close(p[1]); return -1; }
    if (dup2(p[1], STDOUT_FILENO) < 0) {
        close(saved); close(p[0]); close(p[1]); return -1;
    }
    close(p[1]);
    *saved_stdout_fd = saved;
    *capture_read_fd = p[0];
    return 0;
}

/* Helper: finish capturing stdout and return buffer+length. The returned buffer is malloc'd. */
static char *end_capture_stdout(int saved_stdout_fd, int capture_read_fd, size_t *out_len) {
    /* Restore stdout first */
    fflush(stdout);
    dup2(saved_stdout_fd, STDOUT_FILENO);
    close(saved_stdout_fd);

    /* Read all data from capture_read_fd */
    size_t cap = 1024;
    size_t len = 0;
    char *buf = (char *)malloc(cap);
    if (!buf) { close(capture_read_fd); return NULL; }

    while (1) {
        char tmp[4096];
        ssize_t r = read(capture_read_fd, tmp, sizeof tmp);
        if (r < 0) { free(buf); close(capture_read_fd); return NULL; }
        if (r == 0) break;
        if (len + (size_t)r > cap) {
            size_t new_cap = (cap * 2 > len + (size_t)r) ? cap * 2 : len + (size_t)r;
            char *nb = (char *)realloc(buf, new_cap);
            if (!nb) { free(buf); close(capture_read_fd); return NULL; }
            buf = nb; cap = new_cap;
        }
        memcpy(buf + len, tmp, r);
        len += (size_t)r;
    }
    close(capture_read_fd);
    *out_len = len;
    return buf;
}

/* Run head_lines on a given input fd and capture stdout output. */
static int run_head_lines_and_capture(int input_fd, const char *fname, uintmax_t n_lines,
                                      bool *ret_ok, char **out_buf, size_t *out_len) {
    int saved_out = -1, cap_r = -1;
    if (begin_capture_stdout(&saved_out, &cap_r) != 0) {
        return -1;
    }

    /* Call target function while stdout is redirected. No Unity asserts here. */
    bool ok = head_lines(fname, input_fd, n_lines);

    /* Now restore stdout and read captured data */
    size_t got_len = 0;
    char *data = end_capture_stdout(saved_out, cap_r, &got_len);
    if (!data) {
        *ret_ok = ok;
        return -2;
    }

    *ret_ok = ok;
    *out_buf = data;
    *out_len = got_len;
    return 0;
}

/* Helper: create a temporary regular file with given content, return fd open for read-only. */
static int make_temp_regular_file(const char *data, size_t len, char *out_path, size_t out_path_sz) {
    const char *dir = "/tmp";
    char templ[256];
    snprintf(templ, sizeof templ, "%s/head_lines_test_XXXXXX", dir);
    int tfd = mkstemp(templ);
    if (tfd < 0) return -1;
    /* Write data, then reopen read-only to simulate typical usage. */
    ssize_t off = 0;
    while ((size_t)off < len) {
        ssize_t w = write(tfd, data + off, len - off);
        if (w < 0) { close(tfd); unlink(templ); return -1; }
        off += w;
    }
    /* Rewind */
    if (lseek(tfd, 0, SEEK_SET) < 0) { close(tfd); unlink(templ); return -1; }

    /* Duplicate to a read-only fd, then close original to allow unlink. */
    int ro_fd = open(templ, O_RDONLY);
    if (ro_fd < 0) { close(tfd); unlink(templ); return -1; }

    /* Keep path if requested */
    if (out_path && out_path_sz > 0) {
        snprintf(out_path, out_path_sz, "%s", templ);
    }

    /* Unlink the file so it is removed after close; ro_fd remains valid */
    unlink(templ);
    close(tfd);
    return ro_fd;
}

void setUp(void) {
    /* Ensure default delimiter is newline for most tests. */
    extern char line_end; /* available via inclusion in same TU */
    line_end = '\n';
}

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

static void assert_mem_eq(const char *exp, size_t exp_len, const char *got, size_t got_len) {
    TEST_ASSERT_EQUAL_UINT64_MESSAGE((uint64_t)exp_len, (uint64_t)got_len, "Length mismatch");
    if (exp_len == got_len && exp_len > 0) {
        TEST_ASSERT_EQUAL_INT_MESSAGE(0, memcmp(exp, got, exp_len), "Content mismatch");
    }
}

void test_head_lines_zero_lines(void) {
    const char *in = "x\ny\nz\n";
    int fd = make_input_pipe(in, strlen(in));
    TEST_ASSERT_TRUE_MESSAGE(fd >= 0, "Failed to create input pipe");

    bool ok = false; char *out = NULL; size_t out_len = 0;
    int rc = run_head_lines_and_capture(fd, "pipe", 0, &ok, &out, &out_len);
    close(fd);

    TEST_ASSERT_EQUAL_INT_MESSAGE(0, rc, "Capture failed");
    TEST_ASSERT_TRUE_MESSAGE(ok, "head_lines returned false for zero lines");
    TEST_ASSERT_EQUAL_UINT64(0, (uint64_t)out_len);
    free(out);
}

void test_head_lines_basic_two_lines(void) {
    const char *in = "a\nb\nc\n";
    const char *exp = "a\nb\n";
    int fd = make_input_pipe(in, strlen(in));
    TEST_ASSERT_TRUE(fd >= 0);

    bool ok = false; char *out = NULL; size_t out_len = 0;
    int rc = run_head_lines_and_capture(fd, "pipe", 2, &ok, &out, &out_len);
    close(fd);

    TEST_ASSERT_EQUAL_INT(0, rc);
    TEST_ASSERT_TRUE(ok);
    assert_mem_eq(exp, strlen(exp), out, out_len);
    free(out);
}

void test_head_lines_more_than_available(void) {
    const char *in = "x\ny\n";
    int fd = make_input_pipe(in, strlen(in));
    TEST_ASSERT_TRUE(fd >= 0);

    bool ok = false; char *out = NULL; size_t out_len = 0;
    int rc = run_head_lines_and_capture(fd, "pipe", 10, &ok, &out, &out_len);
    close(fd);

    TEST_ASSERT_EQUAL_INT(0, rc);
    TEST_ASSERT_TRUE(ok);
    assert_mem_eq(in, strlen(in), out, out_len);
    free(out);
}

void test_head_lines_no_trailing_newline(void) {
    const char *in = "one\ntwo"; /* only one newline total */
    int fd = make_input_pipe(in, strlen(in));
    TEST_ASSERT_TRUE(fd >= 0);

    bool ok = false; char *out = NULL; size_t out_len = 0;
    int rc = run_head_lines_and_capture(fd, "pipe", 2, &ok, &out, &out_len);
    close(fd);

    TEST_ASSERT_EQUAL_INT(0, rc);
    TEST_ASSERT_TRUE(ok);
    assert_mem_eq(in, strlen(in), out, out_len);
    free(out);
}

void test_head_lines_multiple_reads_and_limit(void) {
    /* Build input with many short lines to force multiple reads via pipe */
    const int total_lines = 3000;
    const int want_lines = 1500;
    size_t in_len = (size_t)total_lines * 2; /* "x\n" per line */
    char *in = (char *)malloc(in_len);
    TEST_ASSERT_NOT_NULL(in);
    for (int i = 0; i < total_lines; i++) { in[2*i] = 'x'; in[2*i + 1] = '\n'; }

    int fd = make_input_pipe(in, in_len);
    TEST_ASSERT_TRUE(fd >= 0);

    bool ok = false; char *out = NULL; size_t out_len = 0;
    int rc = run_head_lines_and_capture(fd, "pipe", (uintmax_t)want_lines, &ok, &out, &out_len);
    close(fd);

    TEST_ASSERT_EQUAL_INT(0, rc);
    TEST_ASSERT_TRUE(ok);

    size_t exp_len = (size_t)want_lines * 2;
    TEST_ASSERT_EQUAL_UINT64((uint64_t)exp_len, (uint64_t)out_len);
    for (size_t i = 0; i < out_len; i += 2) {
        TEST_ASSERT_EQUAL_CHAR('x', out[i]);
        TEST_ASSERT_EQUAL_CHAR('\n', out[i+1]);
    }

    free(out);
    free(in);
}

void test_head_lines_regular_file_seeks_back(void) {
    const char *in = "A\nB"; /* one newline, extra byte after */
    char pathbuf[256];
    int fd = make_temp_regular_file(in, strlen(in), pathbuf, sizeof pathbuf);
    TEST_ASSERT_TRUE_MESSAGE(fd >= 0, "Failed to create temp regular file");

    bool ok = false; char *out = NULL; size_t out_len = 0;
    int rc = run_head_lines_and_capture(fd, pathbuf, 1, &ok, &out, &out_len);

    TEST_ASSERT_EQUAL_INT(0, rc);
    TEST_ASSERT_TRUE(ok);
    const char *exp = "A\n";
    assert_mem_eq(exp, strlen(exp), out, out_len);

    /* The file position should now be right after the newline (offset 2) */
    off_t pos = lseek(fd, 0, SEEK_CUR);
    TEST_ASSERT_EQUAL_INT(2, (int)pos);

    /* Optional: the next byte should be 'B' */
    char c;
    ssize_t r = read(fd, &c, 1);
    TEST_ASSERT_EQUAL_INT(1, (int)r);
    TEST_ASSERT_EQUAL_CHAR('B', c);

    close(fd);
    free(out);
}

void test_head_lines_zero_terminated_delimiter(void) {
    extern char line_end;
    line_end = '\0';

    /* Build input: "one\0two\0tail" */
    const char raw_in[] = { 'o','n','e','\0','t','w','o','\0','t','a','i','l' };
    const size_t raw_in_len = sizeof raw_in;
    const char raw_exp[] = { 'o','n','e','\0','t','w','o','\0' };
    const size_t raw_exp_len = sizeof raw_exp;

    int fd = make_input_pipe(raw_in, raw_in_len);
    TEST_ASSERT_TRUE(fd >= 0);

    bool ok = false; char *out = NULL; size_t out_len = 0;
    int rc = run_head_lines_and_capture(fd, "pipe-z", 2, &ok, &out, &out_len);
    close(fd);

    TEST_ASSERT_EQUAL_INT(0, rc);
    TEST_ASSERT_TRUE(ok);
    assert_mem_eq(raw_exp, raw_exp_len, out, out_len);

    free(out);

    /* Restore default for subsequent tests */
    line_end = '\n';
}

void test_head_lines_read_error_invalid_fd(void) {
    /* Call with invalid FD to simulate read error; expect false return */
    bool ok = true; char *out = NULL; size_t out_len = 0;

    /* We do not need to capture stdout here, but use the same helper for consistency. */
    int saved_out = -1, cap_r = -1;
    TEST_ASSERT_EQUAL_INT(0, begin_capture_stdout(&saved_out, &cap_r));
    bool ret = head_lines("invalid-fd", -1, 1);
    char *captured = end_capture_stdout(saved_out, cap_r, &out_len);
    /* We don't assert on captured stderr (not captured) or stdout content. */
    free(captured);

    TEST_ASSERT_FALSE_MESSAGE(ret, "head_lines should return false on read error");
}

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_head_lines_zero_lines);
    RUN_TEST(test_head_lines_basic_two_lines);
    RUN_TEST(test_head_lines_more_than_available);
    RUN_TEST(test_head_lines_no_trailing_newline);
    RUN_TEST(test_head_lines_multiple_reads_and_limit);
    RUN_TEST(test_head_lines_regular_file_seeks_back);
    RUN_TEST(test_head_lines_zero_terminated_delimiter);
    RUN_TEST(test_head_lines_read_error_invalid_fd);
    return UNITY_END();
}