#include "../../unity/unity.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>

/* Helper: read all content from a FILE* into a newly-allocated buffer.
   Returns buffer (NUL-terminated) and sets *len_out. Caller must free().
   Returns NULL on error. */
static char *read_all_from_FILE(FILE *f, size_t *len_out)
{
    if (!f || !len_out)
        return NULL;

    if (fflush(f) != 0)
        return NULL;

    if (fseek(f, 0, SEEK_END) != 0)
        return NULL;
    long end = ftell(f);
    if (end < 0)
        return NULL;
    if (fseek(f, 0, SEEK_SET) != 0)
        return NULL;

    size_t len = (size_t)end;
    char *buf = (char *)malloc(len + 1);
    if (!buf)
        return NULL;

    size_t n = fread(buf, 1, len, f);
    if (n != len)
    {
        free(buf);
        return NULL;
    }
    buf[len] = '\0';
    *len_out = len;
    return buf;
}

/* Helper: run wrap_write while capturing stdout to a temporary file.
   - Returns a malloc'd buffer containing stdout content (NUL-terminated).
   - Sets *captured_len to the length captured.
   - Returns NULL on error.
   IMPORTANT: No Unity asserts while stdout is redirected. */
static char *capture_stdout_from_wrap(const char *buffer, idx_t len,
                                      idx_t wrap_column, idx_t *current_column,
                                      FILE *out_stream, size_t *captured_len)
{
    if (!captured_len)
        return NULL;

    FILE *cap = tmpfile();
    if (!cap)
        return NULL;

    int saved = dup(fileno(stdout));
    if (saved < 0)
    {
        fclose(cap);
        return NULL;
    }

    int cap_fd = fileno(cap);
    if (cap_fd < 0)
    {
        close(saved);
        fclose(cap);
        return NULL;
    }

    if (fflush(stdout) != 0)
    {
        close(saved);
        fclose(cap);
        return NULL;
    }

    if (dup2(cap_fd, fileno(stdout)) < 0)
    {
        close(saved);
        fclose(cap);
        return NULL;
    }

    /* Call the target function while stdout is redirected. */
    wrap_write(buffer, len, wrap_column, current_column, out_stream);

    /* Flush and restore stdout. Avoid Unity asserts before restoration. */
    fflush(stdout);
    if (dup2(saved, fileno(stdout)) < 0)
    {
        close(saved);
        fclose(cap);
        return NULL;
    }
    close(saved);

    /* Now safe to read the captured stdout and assert later. */
    size_t outlen = 0;
    char *captured = read_all_from_FILE(cap, &outlen);
    fclose(cap);

    if (!captured)
        return NULL;

    *captured_len = outlen;
    return captured;
}

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

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

/* Test: wrap_column == 0 => write all to stdout, do not change current_column, no newline to out. */
void test_wrap_write_no_wrap(void)
{
    const char *input = "hello";
    idx_t wrap = 0;
    idx_t cur = (idx_t)7; /* arbitrary, should remain unchanged */

    FILE *out = tmpfile();
    TEST_ASSERT_NOT_NULL(out);

    size_t cap_len = 0;
    char *cap = capture_stdout_from_wrap(input, (idx_t)strlen(input), wrap, &cur, out, &cap_len);
    TEST_ASSERT_NOT_NULL_MESSAGE(cap, "Failed to capture stdout");

    /* Check stdout content */
    TEST_ASSERT_EQUAL_UINT64((uint64_t)strlen(input), (uint64_t)cap_len);
    TEST_ASSERT_EQUAL_MEMORY(input, cap, cap_len);

    /* out stream should be empty (no newlines) */
    size_t out_len = 0;
    char *out_buf = read_all_from_FILE(out, &out_len);
    TEST_ASSERT_NOT_NULL(out_buf);
    TEST_ASSERT_EQUAL_UINT64(0, (uint64_t)out_len);

    /* current_column unchanged */
    TEST_ASSERT_EQUAL_INT(7, (int)cur);

    free(cap);
    free(out_buf);
    fclose(out);
}

/* Test: simple write below wrap boundary => data to stdout, no newline, current_column updated. */
void test_wrap_write_simple_partial_line(void)
{
    const char *input = "abcde";
    idx_t wrap = (idx_t)10;
    idx_t cur = (idx_t)0;

    FILE *out = tmpfile();
    TEST_ASSERT_NOT_NULL(out);

    size_t cap_len = 0;
    char *cap = capture_stdout_from_wrap(input, (idx_t)strlen(input), wrap, &cur, out, &cap_len);
    TEST_ASSERT_NOT_NULL(cap);

    TEST_ASSERT_EQUAL_UINT64((uint64_t)strlen(input), (uint64_t)cap_len);
    TEST_ASSERT_EQUAL_MEMORY(input, cap, cap_len);

    size_t out_len = 0;
    char *out_buf = read_all_from_FILE(out, &out_len);
    TEST_ASSERT_NOT_NULL(out_buf);
    TEST_ASSERT_EQUAL_UINT64(0, (uint64_t)out_len);

    TEST_ASSERT_EQUAL_INT(5, (int)cur);

    free(cap);
    free(out_buf);
    fclose(out);
}

/* Test: crossing a boundary emits a newline to out and continues writing; current_column updated */
void test_wrap_write_crosses_boundary(void)
{
    const char *input = "WXYZ"; /* length 4 */
    idx_t wrap = (idx_t)5;
    idx_t cur = (idx_t)3; /* 2 chars to fill line, then newline, then 2 more */

    FILE *out = tmpfile();
    TEST_ASSERT_NOT_NULL(out);

    size_t cap_len = 0;
    char *cap = capture_stdout_from_wrap(input, (idx_t)strlen(input), wrap, &cur, out, &cap_len);
    TEST_ASSERT_NOT_NULL(cap);

    /* All data should be on stdout, no newlines in stdout */
    TEST_ASSERT_EQUAL_UINT64((uint64_t)strlen(input), (uint64_t)cap_len);
    TEST_ASSERT_EQUAL_MEMORY(input, cap, cap_len);

    /* out should contain exactly one newline */
    size_t out_len = 0;
    char *out_buf = read_all_from_FILE(out, &out_len);
    TEST_ASSERT_NOT_NULL(out_buf);
    TEST_ASSERT_EQUAL_UINT64(1, (uint64_t)out_len);
    TEST_ASSERT_EQUAL_CHAR('\n', out_buf[0]);

    /* current_column is remaining after last write: 2 */
    TEST_ASSERT_EQUAL_INT(2, (int)cur);

    free(cap);
    free(out_buf);
    fclose(out);
}

/* Test: multiple wraps across one call: ensure multiple newlines to out, stdout has contiguous data */
void test_wrap_write_multiple_wraps(void)
{
    const char *input = "abcdefghij"; /* 10 bytes */
    idx_t wrap = (idx_t)4;
    idx_t cur = (idx_t)0;

    FILE *out = tmpfile();
    TEST_ASSERT_NOT_NULL(out);

    size_t cap_len = 0;
    char *cap = capture_stdout_from_wrap(input, (idx_t)strlen(input), wrap, &cur, out, &cap_len);
    TEST_ASSERT_NOT_NULL(cap);

    /* stdout should contain exactly the input (no newlines added there) */
    TEST_ASSERT_EQUAL_UINT64((uint64_t)strlen(input), (uint64_t)cap_len);
    TEST_ASSERT_EQUAL_MEMORY(input, cap, cap_len);

    /* out should contain two newlines (after 4 and after 8) */
    size_t out_len = 0;
    char *out_buf = read_all_from_FILE(out, &out_len);
    TEST_ASSERT_NOT_NULL(out_buf);
    TEST_ASSERT_EQUAL_UINT64(2, (uint64_t)out_len);
    TEST_ASSERT_EQUAL_CHAR('\n', out_buf[0]);
    TEST_ASSERT_EQUAL_CHAR('\n', out_buf[1]);

    /* 10 bytes => after two wraps (8 written), remaining current_column = 2 */
    TEST_ASSERT_EQUAL_INT(2, (int)cur);

    free(cap);
    free(out_buf);
    fclose(out);
}

/* Test: exact fill does not emit newline in the same call; next call emits the newline first */
void test_wrap_write_exact_fill_then_next_call_emits_newline(void)
{
    idx_t wrap = (idx_t)4;
    idx_t cur = (idx_t)2;

    FILE *out = tmpfile();
    TEST_ASSERT_NOT_NULL(out);

    /* First call fills to wrap exactly (no newline). */
    const char *first = "XY"; /* len 2: 2+2=4 */
    size_t cap1_len = 0;
    char *cap1 = capture_stdout_from_wrap(first, (idx_t)strlen(first), wrap, &cur, out, &cap1_len);
    TEST_ASSERT_NOT_NULL(cap1);

    TEST_ASSERT_EQUAL_UINT64((uint64_t)strlen(first), (uint64_t)cap1_len);
    TEST_ASSERT_EQUAL_MEMORY(first, cap1, cap1_len);

    size_t out_len1 = 0;
    char *out_buf1 = read_all_from_FILE(out, &out_len1);
    TEST_ASSERT_NOT_NULL(out_buf1);
    TEST_ASSERT_EQUAL_UINT64(0, (uint64_t)out_len1);
    TEST_ASSERT_EQUAL_INT(4, (int)cur);

    free(cap1);
    free(out_buf1);

    /* Second call with 1 byte should emit newline first, then write the byte. */
    const char *second = "Z";
    size_t cap2_len = 0;
    char *cap2 = capture_stdout_from_wrap(second, (idx_t)strlen(second), wrap, &cur, out, &cap2_len);
    TEST_ASSERT_NOT_NULL(cap2);

    TEST_ASSERT_EQUAL_UINT64((uint64_t)strlen(second), (uint64_t)cap2_len);
    TEST_ASSERT_EQUAL_MEMORY(second, cap2, cap2_len);

    size_t out_len2 = 0;
    char *out_buf2 = read_all_from_FILE(out, &out_len2);
    TEST_ASSERT_NOT_NULL(out_buf2);
    TEST_ASSERT_EQUAL_UINT64(1, (uint64_t)out_len2);
    TEST_ASSERT_EQUAL_CHAR('\n', out_buf2[0]);

    TEST_ASSERT_EQUAL_INT(1, (int)cur);

    free(cap2);
    free(out_buf2);
    fclose(out);
}

/* Test: zero-length input produces no output and no newline; current_column unchanged */
void test_wrap_write_zero_length_no_effect(void)
{
    idx_t wrap = (idx_t)6;
    idx_t cur = (idx_t)3;

    FILE *out = tmpfile();
    TEST_ASSERT_NOT_NULL(out);

    size_t cap_len = 0;
    char *cap = capture_stdout_from_wrap("", (idx_t)0, wrap, &cur, out, &cap_len);
    TEST_ASSERT_NOT_NULL(cap);

    TEST_ASSERT_EQUAL_UINT64(0, (uint64_t)cap_len);

    size_t out_len = 0;
    char *out_buf = read_all_from_FILE(out, &out_len);
    TEST_ASSERT_NOT_NULL(out_buf);
    TEST_ASSERT_EQUAL_UINT64(0, (uint64_t)out_len);

    TEST_ASSERT_EQUAL_INT(3, (int)cur);

    free(cap);
    free(out_buf);
    fclose(out);
}

/* Test: if current_column == wrap_column at call start, a newline is emitted before writing */
void test_wrap_write_starts_full_line_emits_newline_first(void)
{
    const char *input = "123";
    idx_t wrap = (idx_t)4;
    idx_t cur = (idx_t)4; /* already full */

    FILE *out = tmpfile();
    TEST_ASSERT_NOT_NULL(out);

    size_t cap_len = 0;
    char *cap = capture_stdout_from_wrap(input, (idx_t)strlen(input), wrap, &cur, out, &cap_len);
    TEST_ASSERT_NOT_NULL(cap);

    TEST_ASSERT_EQUAL_UINT64((uint64_t)strlen(input), (uint64_t)cap_len);
    TEST_ASSERT_EQUAL_MEMORY(input, cap, cap_len);

    size_t out_len = 0;
    char *out_buf = read_all_from_FILE(out, &out_len);
    TEST_ASSERT_NOT_NULL(out_buf);
    TEST_ASSERT_EQUAL_UINT64(1, (uint64_t)out_len);
    TEST_ASSERT_EQUAL_CHAR('\n', out_buf[0]);

    TEST_ASSERT_EQUAL_INT(3, (int)cur);

    free(cap);
    free(out_buf);
    fclose(out);
}

int main(void)
{
    UNITY_BEGIN();

    RUN_TEST(test_wrap_write_no_wrap);
    RUN_TEST(test_wrap_write_simple_partial_line);
    RUN_TEST(test_wrap_write_crosses_boundary);
    RUN_TEST(test_wrap_write_multiple_wraps);
    RUN_TEST(test_wrap_write_exact_fill_then_next_call_emits_newline);
    RUN_TEST(test_wrap_write_zero_length_no_effect);
    RUN_TEST(test_wrap_write_starts_full_line_emits_newline_first);

    return UNITY_END();
}