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

/* We rely on being included in the same translation unit as paste.c,
   so we can use:
   - static bool paste_parallel(size_t, char **)
   - static int collapse_escapes(const char *)
   - globals: have_read_stdin, delims, delim_end, line_delim
*/

/* Helper: create a temporary file with the given bytes.
   Returns a malloc'd path string that the caller must free (and unlink). */
static char *create_temp_file_with_bytes(const void *data, size_t len)
{
    char tmpl[] = "/tmp/paste_test_in_XXXXXX";
    int fd = mkstemp(tmpl);
    if (fd < 0) {
        return NULL;
    }
    ssize_t w = write(fd, data, len);
    if (w < 0 || (size_t)w != len) {
        close(fd);
        unlink(tmpl);
        return NULL;
    }
    if (close(fd) != 0) {
        unlink(tmpl);
        return NULL;
    }
    char *path = (char *)malloc(strlen(tmpl) + 1);
    if (!path) {
        unlink(tmpl);
        return NULL;
    }
    strcpy(path, tmpl);
    return path;
}

/* Helper: run paste_parallel with given file list and delimiter spec,
   capturing stdout into a buffer. Does not use Unity asserts while
   stdout is redirected. Returns buffer (malloc'd) and length via out_len.
   Sets *out_ok to the function return value. Returns NULL on severe error. */
static unsigned char *run_and_capture(char **files, size_t nfiles,
                                      const char *delimspec,
                                      int use_zero_terminated,
                                      size_t *out_len, int *out_ok)
{
    if (!delimspec) delimspec = "\t";

    /* Reset relevant globals */
    have_read_stdin = false;
    line_delim = use_zero_terminated ? '\0' : '\n';
    /* Set delimiters via collapse_escapes */
    collapse_escapes(delimspec);

    /* Prepare temp output */
    char outtmpl[] = "/tmp/paste_test_out_XXXXXX";
    int outfd = mkstemp(outtmpl);
    if (outfd < 0) {
        return NULL;
    }

    /* Redirect stdout */
    fflush(stdout);
    int saved_stdout = dup(STDOUT_FILENO);
    if (saved_stdout < 0) {
        close(outfd);
        unlink(outtmpl);
        return NULL;
    }
    if (dup2(outfd, STDOUT_FILENO) < 0) {
        close(saved_stdout);
        close(outfd);
        unlink(outtmpl);
        return NULL;
    }

    /* Call the function under test */
    bool ok = paste_parallel(nfiles, files);

    /* Flush and restore stdout */
    fflush(stdout);
    fsync(STDOUT_FILENO);
    if (dup2(saved_stdout, STDOUT_FILENO) < 0) {
        /* Attempt to restore failed; clean what we can */
        close(saved_stdout);
        close(outfd);
        unlink(outtmpl);
        return NULL;
    }
    close(saved_stdout);

    /* Read captured output */
    struct stat st;
    if (fstat(outfd, &st) != 0) {
        close(outfd);
        unlink(outtmpl);
        return NULL;
    }
    size_t len = (size_t)st.st_size;
    if (lseek(outfd, 0, SEEK_SET) < 0) {
        close(outfd);
        unlink(outtmpl);
        return NULL;
    }
    unsigned char *buf = (unsigned char *)malloc(len ? len : 1);
    if (!buf) {
        close(outfd);
        unlink(outtmpl);
        return NULL;
    }
    size_t total = 0;
    while (total < len) {
        ssize_t r = read(outfd, buf + total, len - total);
        if (r < 0) {
            free(buf);
            close(outfd);
            unlink(outtmpl);
            return NULL;
        }
        if (r == 0) break;
        total += (size_t)r;
    }
    close(outfd);
    unlink(outtmpl);

    *out_len = total;
    *out_ok = ok ? 1 : 0;
    return buf;
}

void setUp(void) {
    /* Nothing for now */
}

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

static void free_and_unlink_paths(char **paths, size_t n)
{
    if (!paths) return;
    for (size_t i = 0; i < n; i++) {
        if (paths[i]) {
            unlink(paths[i]);
            free(paths[i]);
        }
    }
}

/* Test 1: Basic two files with equal number of lines and default tab delimiter. */
void test_paste_parallel_basic_two_files(void)
{
    const char *a = "a1\na2\n";
    const char *b = "b1\nb2\n";

    char *paths[2];
    paths[0] = create_temp_file_with_bytes(a, strlen(a));
    paths[1] = create_temp_file_with_bytes(b, strlen(b));
    TEST_ASSERT_NOT_NULL_MESSAGE(paths[0], "failed to create temp file A");
    TEST_ASSERT_NOT_NULL_MESSAGE(paths[1], "failed to create temp file B");

    char *files[3] = { paths[0], paths[1], NULL };

    size_t out_len = 0;
    int out_ok = 0;
    unsigned char *out = run_and_capture(files, 2, "\t", 0, &out_len, &out_ok);

    TEST_ASSERT_NOT_NULL(out);
    TEST_ASSERT_EQUAL_INT(1, out_ok);

    const char *expected = "a1\tb1\na2\tb2\n";
    TEST_ASSERT_EQUAL_size_t(strlen(expected), out_len);
    TEST_ASSERT_EQUAL_INT(0, memcmp(expected, out, out_len));

    free(out);
    free_and_unlink_paths(paths, 2);
}

/* Test 2: Second file longer than first; ensure leading delimiter for missing field. */
void test_paste_parallel_second_file_longer(void)
{
    const char *a = "a1\na2\n";
    const char *b = "b1\nb2\nb3\n";

    char *paths[2];
    paths[0] = create_temp_file_with_bytes(a, strlen(a));
    paths[1] = create_temp_file_with_bytes(b, strlen(b));
    TEST_ASSERT_NOT_NULL(paths[0]);
    TEST_ASSERT_NOT_NULL(paths[1]);

    char *files[3] = { paths[0], paths[1], NULL };

    size_t out_len = 0;
    int out_ok = 0;
    unsigned char *out = run_and_capture(files, 2, "\t", 0, &out_len, &out_ok);

    TEST_ASSERT_NOT_NULL(out);
    TEST_ASSERT_EQUAL_INT(1, out_ok);

    const char *expected = "a1\tb1\na2\tb2\n\tb3\n";
    TEST_ASSERT_EQUAL_size_t(strlen(expected), out_len);
    TEST_ASSERT_EQUAL_INT(0, memcmp(expected, out, out_len));

    free(out);
    free_and_unlink_paths(paths, 2);
}

/* Test 3: Cycling delimiters across multiple files. */
void test_paste_parallel_cycling_delimiters(void)
{
    const char *a = "a\n";
    const char *b = "b\n";
    const char *c = "c\n";

    char *paths[3];
    paths[0] = create_temp_file_with_bytes(a, strlen(a));
    paths[1] = create_temp_file_with_bytes(b, strlen(b));
    paths[2] = create_temp_file_with_bytes(c, strlen(c));
    TEST_ASSERT_NOT_NULL(paths[0]);
    TEST_ASSERT_NOT_NULL(paths[1]);
    TEST_ASSERT_NOT_NULL(paths[2]);

    char *files[4] = { paths[0], paths[1], paths[2], NULL };

    size_t out_len = 0;
    int out_ok = 0;
    unsigned char *out = run_and_capture(files, 3, ",;", 0, &out_len, &out_ok);

    TEST_ASSERT_NOT_NULL(out);
    TEST_ASSERT_EQUAL_INT(1, out_ok);

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

    free(out);
    free_and_unlink_paths(paths, 3);
}

/* Test 4: EMPTY_DELIM behavior using "\0" (no delimiter between columns). */
void test_paste_parallel_empty_delimiter(void)
{
    const char *a = "a\n";
    const char *b = "b\n";

    char *paths[2];
    paths[0] = create_temp_file_with_bytes(a, strlen(a));
    paths[1] = create_temp_file_with_bytes(b, strlen(b));
    TEST_ASSERT_NOT_NULL(paths[0]);
    TEST_ASSERT_NOT_NULL(paths[1]);

    char *files[3] = { paths[0], paths[1], NULL };

    size_t out_len = 0;
    int out_ok = 0;
    /* Note: collapse_escapes expects backslash escapes; pass "\\0" to mean EMPTY_DELIM */
    unsigned char *out = run_and_capture(files, 2, "\\0", 0, &out_len, &out_ok);

    TEST_ASSERT_NOT_NULL(out);
    TEST_ASSERT_EQUAL_INT(1, out_ok);

    const char *expected = "ab\n";
    TEST_ASSERT_EQUAL_size_t(strlen(expected), out_len);
    TEST_ASSERT_EQUAL_INT(0, memcmp(expected, out, out_len));

    free(out);
    free_and_unlink_paths(paths, 2);
}

/* Test 5: Last line of last file lacks newline; ensure newline is printed. */
void test_paste_parallel_missing_newline_on_last_file(void)
{
    const char *a = "1\n2\n";
    const char *b = "a\nb"; /* No trailing newline */

    char *paths[2];
    paths[0] = create_temp_file_with_bytes(a, strlen(a));
    paths[1] = create_temp_file_with_bytes(b, strlen(b));
    TEST_ASSERT_NOT_NULL(paths[0]);
    TEST_ASSERT_NOT_NULL(paths[1]);

    char *files[3] = { paths[0], paths[1], NULL };

    size_t out_len = 0;
    int out_ok = 0;
    unsigned char *out = run_and_capture(files, 2, "\t", 0, &out_len, &out_ok);

    TEST_ASSERT_NOT_NULL(out);
    TEST_ASSERT_EQUAL_INT(1, out_ok);

    const char *expected = "1\ta\n2\tb\n";
    TEST_ASSERT_EQUAL_size_t(strlen(expected), out_len);
    TEST_ASSERT_EQUAL_INT(0, memcmp(expected, out, out_len));
    /* Ensure the last character is a newline */
    TEST_ASSERT_TRUE(out_len > 0);
    TEST_ASSERT_EQUAL_CHAR('\n', out[out_len - 1]);

    free(out);
    free_and_unlink_paths(paths, 2);
}

/* Test 6: Zero-terminated input (line_delim = '\0'). */
void test_paste_parallel_zero_terminated(void)
{
    const unsigned char a[] = { 'a','1','\0','a','2','\0' };
    const unsigned char b[] = { 'b','1','\0','b','2','\0' };

    char *paths[2];
    paths[0] = create_temp_file_with_bytes(a, sizeof(a));
    paths[1] = create_temp_file_with_bytes(b, sizeof(b));
    TEST_ASSERT_NOT_NULL(paths[0]);
    TEST_ASSERT_NOT_NULL(paths[1]);

    char *files[3] = { paths[0], paths[1], NULL };

    size_t out_len = 0;
    int out_ok = 0;
    unsigned char *out = run_and_capture(files, 2, "\t", 1, &out_len, &out_ok);

    TEST_ASSERT_NOT_NULL(out);
    TEST_ASSERT_EQUAL_INT(1, out_ok);

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

    free(out);
    free_and_unlink_paths(paths, 2);
}

/* Test 7: Three files with early closure and delimiter buffering across columns. */
void test_paste_parallel_three_files_early_close_delbuf(void)
{
    const char *a = "a1\n";
    const char *b = "b1\nb2\n";
    const char *c = "c1\nc2\nc3\n";

    char *paths[3];
    paths[0] = create_temp_file_with_bytes(a, strlen(a));
    paths[1] = create_temp_file_with_bytes(b, strlen(b));
    paths[2] = create_temp_file_with_bytes(c, strlen(c));
    TEST_ASSERT_NOT_NULL(paths[0]);
    TEST_ASSERT_NOT_NULL(paths[1]);
    TEST_ASSERT_NOT_NULL(paths[2]);

    char *files[4] = { paths[0], paths[1], paths[2], NULL };

    size_t out_len = 0;
    int out_ok = 0;
    unsigned char *out = run_and_capture(files, 3, "\t", 0, &out_len, &out_ok);

    TEST_ASSERT_NOT_NULL(out);
    TEST_ASSERT_EQUAL_INT(1, out_ok);

    const char *expected = "a1\tb1\tc1\n\tb2\tc2\n\t\tc3\n";
    TEST_ASSERT_EQUAL_size_t(strlen(expected), out_len);
    TEST_ASSERT_EQUAL_INT(0, memcmp(expected, out, out_len));

    free(out);
    free_and_unlink_paths(paths, 3);
}

/* Test 8: EMPTY_DELIM not saved in delbuf; use "\0," with three files,
   first file has fewer lines. */
void test_paste_parallel_empty_delim_not_saved_in_delbuf(void)
{
    const char *a = "a1\n";
    const char *b = "b1\nb2\n";
    const char *c = "c1\nc2\n";

    char *paths[3];
    paths[0] = create_temp_file_with_bytes(a, strlen(a));
    paths[1] = create_temp_file_with_bytes(b, strlen(b));
    paths[2] = create_temp_file_with_bytes(c, strlen(c));
    TEST_ASSERT_NOT_NULL(paths[0]);
    TEST_ASSERT_NOT_NULL(paths[1]);
    TEST_ASSERT_NOT_NULL(paths[2]);

    char *files[4] = { paths[0], paths[1], paths[2], NULL };

    size_t out_len = 0;
    int out_ok = 0;
    /* "\\0," => first delimiter position is EMPTY, second is ',' */
    unsigned char *out = run_and_capture(files, 3, "\\0,", 0, &out_len, &out_ok);

    TEST_ASSERT_NOT_NULL(out);
    TEST_ASSERT_EQUAL_INT(1, out_ok);

    const char *expected = "a1b1,c1\nb2,c2\n";
    TEST_ASSERT_EQUAL_size_t(strlen(expected), out_len);
    TEST_ASSERT_EQUAL_INT(0, memcmp(expected, out, out_len));

    free(out);
    free_and_unlink_paths(paths, 3);
}

int main(void)
{
    UNITY_BEGIN();
    RUN_TEST(test_paste_parallel_basic_two_files);
    RUN_TEST(test_paste_parallel_second_file_longer);
    RUN_TEST(test_paste_parallel_cycling_delimiters);
    RUN_TEST(test_paste_parallel_empty_delimiter);
    RUN_TEST(test_paste_parallel_missing_newline_on_last_file);
    RUN_TEST(test_paste_parallel_zero_terminated);
    RUN_TEST(test_paste_parallel_three_files_early_close_delbuf);
    RUN_TEST(test_paste_parallel_empty_delim_not_saved_in_delbuf);
    return UNITY_END();
}