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

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <errno.h>

/* The program's source (this test file is included into it) has already
   included the necessary headers like base64.h when BASE_TYPE==42. */

/* Helpers to capture child process output */
static char *read_all_fd(int fd, size_t *len_out)
{
    size_t cap = 1024;
    size_t len = 0;
    char *buf = (char *)malloc(cap);
    if (!buf) return NULL;

    while (1) {
        ssize_t n = read(fd, buf + len, cap - len);
        if (n > 0) {
            len += (size_t)n;
            if (cap - len == 0) {
                size_t new_cap = cap * 2;
                char *nbuf = (char *)realloc(buf, new_cap);
                if (!nbuf) { free(buf); return NULL; }
                buf = nbuf;
                cap = new_cap;
            }
        } else if (n == 0) {
            break; /* EOF */
        } else {
            if (errno == EINTR) continue;
            free(buf);
            return NULL;
        }
    }

    char *out = (char *)malloc(len + 1);
    if (!out) { free(buf); return NULL; }
    memcpy(out, buf, len);
    out[len] = '\0';
    free(buf);
    if (len_out) *len_out = len;
    return out;
}

/* Run do_encode in a child process, capturing stdout, and optionally
   directing the 'out' stream to stderr to observe where newlines go. */
static int run_do_encode_capture(const char *input, size_t inlen,
                                 int wrap_column, int out_to_stderr,
                                 char **captured_out, size_t *captured_out_len,
                                 char **captured_err, size_t *captured_err_len)
{
    int out_pipe[2];
    int err_pipe[2];

    if (pipe(out_pipe) != 0)
        return -1;
    int need_err = out_to_stderr ? 1 : 0;
    if (need_err && pipe(err_pipe) != 0) {
        close(out_pipe[0]); close(out_pipe[1]);
        return -1;
    }

    fflush(stdout);
    fflush(stderr);
    pid_t pid = fork();
    if (pid < 0) {
        close(out_pipe[0]); close(out_pipe[1]);
        if (need_err) { close(err_pipe[0]); close(err_pipe[1]); }
        return -1;
    }

    if (pid == 0) {
        /* Child: redirect stdout (and stderr if needed). */
        if (dup2(out_pipe[1], STDOUT_FILENO) < 0) _exit(127);
        close(out_pipe[0]); close(out_pipe[1]);

        if (need_err) {
            if (dup2(err_pipe[1], STDERR_FILENO) < 0) _exit(127);
            close(err_pipe[0]); close(err_pipe[1]);
        }

        /* Prepare input FILE* via tmpfile. */
        FILE *inf = tmpfile();
        if (!inf) _exit(127);
        if (inlen) {
            if (fwrite(input, 1, inlen, inf) != inlen) _exit(127);
        }
        rewind(inf);

        /* Call the target function. It will exit() the process. */
        if (out_to_stderr)
            do_encode(inf, "-", stderr, (idx_t)wrap_column);
        else
            do_encode(inf, "-", stdout, (idx_t)wrap_column);

        /* Not reached */
        _exit(127);
    }

    /* Parent: close write ends and read. */
    close(out_pipe[1]);
    if (need_err) close(err_pipe[1]);

    char *outbuf = read_all_fd(out_pipe[0], captured_out_len);
    close(out_pipe[0]);

    char *errbuf = NULL; size_t errlen = 0;
    if (need_err) {
        errbuf = read_all_fd(err_pipe[0], &errlen);
        close(err_pipe[0]);
    }

    int status = 0;
    if (waitpid(pid, &status, 0) < 0) status = -1;

    if (captured_out) *captured_out = outbuf; else free(outbuf);
    if (need_err) {
        if (captured_err) { *captured_err = errbuf; if (captured_err_len) *captured_err_len = errlen; }
        else free(errbuf);
    } else {
        if (captured_err) *captured_err = NULL; if (captured_err_len) *captured_err_len = 0;
    }

    return status;
}

void setUp(void) {
    /* For basenc (BASE_TYPE==42), initialize encoding function pointers to base64. */
#if BASE_TYPE == 42
    /* Ensure we select the base64 variants for encoding. */
    base_length = base64_length_wrapper;      /* BASE64_LENGTH wrapper */
    required_padding = base64_required_padding; /* not used in encode path but harmless */
    isubase = isubase64;                      /* not used in encode path */
    base_encode = base64_encode;              /* core encoder used by do_encode */

    /* Ensure context-based encoding is disabled by default (so do_encode uses base_encode). */
    base_encode_ctx_init = NULL;
    base_encode_ctx = NULL;
    base_encode_ctx_finalize = NULL;
#endif
}

void tearDown(void) {
    /* Nothing to clean up. */
}

/* Tests */

static void assert_child_ok(int status)
{
    TEST_ASSERT_TRUE_MESSAGE(WIFEXITED(status), "Child did not exit normally");
    TEST_ASSERT_EQUAL_INT_MESSAGE(0, WEXITSTATUS(status), "Child exit status not 0");
}

void test_do_encode_empty_input_no_wrap(void)
{
    const char *input = "";
    char *out = NULL; size_t outlen = 0;
    int status = run_do_encode_capture(input, strlen(input), 0, 0, &out, &outlen, NULL, NULL);
    TEST_ASSERT_NOT_EQUAL(-1, status);
    assert_child_ok(status);
    TEST_ASSERT_EQUAL_size_t(0, outlen);
    if (out) free(out);
}

void test_do_encode_base64_no_wrap_foo(void)
{
    const char *input = "foo";     /* base64: Zm9v */
    char *out = NULL; size_t outlen = 0;
    int status = run_do_encode_capture(input, strlen(input), 0, 0, &out, &outlen, NULL, NULL);
    TEST_ASSERT_NOT_EQUAL(-1, status);
    assert_child_ok(status);
    TEST_ASSERT_EQUAL_STRING_LEN("Zm9v", out, 4);
    TEST_ASSERT_EQUAL_size_t(4, outlen);
    free(out);
}

void test_do_encode_wrap4_exact_boundary_newline_once(void)
{
    const char *input = "foo";     /* base64: Zm9v (length 4 exactly one chunk) */
    char *out = NULL; size_t outlen = 0;
    int status = run_do_encode_capture(input, strlen(input), 4, 0, &out, &outlen, NULL, NULL);
    TEST_ASSERT_NOT_EQUAL(-1, status);
    assert_child_ok(status);
    /* Expect a single trailing newline when wrap>0 and output length is exact multiple. */
    TEST_ASSERT_EQUAL_STRING_LEN("Zm9v\n", out, 5);
    TEST_ASSERT_EQUAL_size_t(5, outlen);
    free(out);
}

void test_do_encode_wrap4_inserts_mid_and_final_newlines(void)
{
    const char *input = "foobar";  /* base64: Zm9vYmFy (8 chars) */
    char *out = NULL; size_t outlen = 0;
    int status = run_do_encode_capture(input, strlen(input), 4, 0, &out, &outlen, NULL, NULL);
    TEST_ASSERT_NOT_EQUAL(-1, status);
    assert_child_ok(status);
    /* Expect split after 4 chars, and final newline: "Zm9v\nYmFy\n" */
    const char *expected = "Zm9v\nYmFy\n";
    TEST_ASSERT_EQUAL_size_t(strlen(expected), outlen);
    TEST_ASSERT_EQUAL_STRING_LEN(expected, out, outlen);
    free(out);
}

void test_do_encode_newlines_go_to_out_stream(void)
{
    const char *input = "foobar";  /* base64: Zm9vYmFy */
    char *capt_out = NULL, *capt_err = NULL; size_t outlen = 0, errlen = 0;
    /* Pass 'out' as stderr so that wrap_write writes newlines to stderr. */
    int status = run_do_encode_capture(input, strlen(input), 4, 1,
                                       &capt_out, &outlen,
                                       &capt_err, &errlen);
    TEST_ASSERT_NOT_EQUAL(-1, status);
    assert_child_ok(status);

    /* stdout should contain only the payload without any newlines: */
    TEST_ASSERT_EQUAL_STRING_LEN("Zm9vYmFy", capt_out, 8);
    TEST_ASSERT_EQUAL_size_t(8, outlen);

    /* stderr should contain the mid and final newlines */
    TEST_ASSERT_NOT_NULL(capt_err);
    TEST_ASSERT_EQUAL_size_t(2, errlen);
    TEST_ASSERT_EQUAL_UINT8('\n', (unsigned char)capt_err[0]);
    TEST_ASSERT_EQUAL_UINT8('\n', (unsigned char)capt_err[1]);

    free(capt_out);
    free(capt_err);
}

int main(void)
{
    UNITY_BEGIN();
    RUN_TEST(test_do_encode_empty_input_no_wrap);
    RUN_TEST(test_do_encode_base64_no_wrap_foo);
    RUN_TEST(test_do_encode_wrap4_exact_boundary_newline_once);
    RUN_TEST(test_do_encode_wrap4_inserts_mid_and_final_newlines);
    RUN_TEST(test_do_encode_newlines_go_to_out_stream);
    return UNITY_END();
}