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

/* The test file is included into the same translation unit as dd.
   We can access and manipulate internal globals directly: obuf, oc,
   output_blocksize, w_full, w_partial, w_bytes, conversions_mask. */

/* Helper to set output blocksize and allocate obuf. */
static void setup_buffer_and_blocksize(idx_t bs)
{
    /* Free any existing buffer owned by tests (if any). */
    if (obuf)
    {
        free(obuf);
        obuf = NULL;
    }
    output_blocksize = bs;
    oc = 0;
    obuf = (char *)malloc((size_t)bs);
    TEST_ASSERT_NOT_NULL(obuf);
}

/* Capture stdout around the call under test. */
typedef struct {
    int saved_stdout;
    int pipe_rd;
} capture_t;

static capture_t begin_capture_stdout(void)
{
    capture_t c;
    c.saved_stdout = -1;
    c.pipe_rd = -1;

    int pipefd[2];
    if (pipe(pipefd) != 0)
    {
        TEST_FAIL_MESSAGE("pipe() failed");
        return c;
    }

    fflush(stdout);
    int saved = dup(STDOUT_FILENO);
    if (saved < 0)
    {
        close(pipefd[0]);
        close(pipefd[1]);
        TEST_FAIL_MESSAGE("dup(STDOUT_FILENO) failed");
        return c;
    }
    if (dup2(pipefd[1], STDOUT_FILENO) < 0)
    {
        close(saved);
        close(pipefd[0]);
        close(pipefd[1]);
        TEST_FAIL_MESSAGE("dup2() failed for STDOUT");
        return c;
    }
    close(pipefd[1]); /* only keep read end */
    c.saved_stdout = saved;
    c.pipe_rd = pipefd[0];
    return c;
}

static void end_capture_stdout(capture_t *c, unsigned char **out, size_t *out_len)
{
    /* Restore STDOUT so we can safely use Unity asserts. */
    if (dup2(c->saved_stdout, STDOUT_FILENO) < 0)
    {
        /* Try to proceed; assertions below will flag issues. */
    }
    close(c->saved_stdout);

    /* Read all captured bytes. */
    size_t cap = 256;
    size_t len = 0;
    unsigned char *buf = (unsigned char *)malloc(cap);
    if (!buf)
    {
        close(c->pipe_rd);
        TEST_FAIL_MESSAGE("malloc failed in end_capture_stdout");
        return;
    }

    for (;;)
    {
        unsigned char tmp[256];
        ssize_t n = read(c->pipe_rd, tmp, sizeof tmp);
        if (n < 0)
        {
            if (errno == EINTR) continue;
            free(buf);
            close(c->pipe_rd);
            TEST_FAIL_MESSAGE("read() failed on capture pipe");
            return;
        }
        if (n == 0)
            break;
        if (len + (size_t)n > cap)
        {
            size_t new_cap = (len + (size_t)n) * 2;
            unsigned char *nb = (unsigned char *)realloc(buf, new_cap);
            if (!nb)
            {
                free(buf);
                close(c->pipe_rd);
                TEST_FAIL_MESSAGE("realloc failed in end_capture_stdout");
                return;
            }
            buf = nb;
            cap = new_cap;
        }
        memcpy(buf + len, tmp, (size_t)n);
        len += (size_t)n;
    }
    close(c->pipe_rd);

    *out = buf;
    *out_len = len;
}

/* Unity-required hooks */
void setUp(void) {
    /* Reset counters and state likely to be affected. */
    oc = 0;
    w_bytes = 0;
    w_full = 0;
    w_partial = 0;
    conversions_mask = 0;
    /* Ensure no stale buffer persists across tests. */
    if (obuf) {
        free(obuf);
        obuf = NULL;
    }
}

void tearDown(void) {
    if (obuf) {
        free(obuf);
        obuf = NULL;
    }
}

/* Tests */

void test_copy_simple_partial_no_flush(void)
{
    setup_buffer_and_blocksize((idx_t)8);
    const char *input = "12345";
    idx_t nread = (idx_t)strlen(input);

    capture_t cap = begin_capture_stdout();
    /* Call function under test while stdout is redirected (no assertions here). */
    copy_simple(input, nread);
    unsigned char *captured = NULL;
    size_t captured_len = 0;
    end_capture_stdout(&cap, &captured, &captured_len);

    /* Now assert postconditions. */
    TEST_ASSERT_EQUAL_UINT32(5, oc);
    TEST_ASSERT_EQUAL_UINT32(0, captured_len);
    TEST_ASSERT_EQUAL_INT(0, w_full);
    TEST_ASSERT_EQUAL_INT(0, w_partial);
    TEST_ASSERT_EQUAL_INT(0, w_bytes);
    TEST_ASSERT_EQUAL_INT(0, memcmp(obuf, input, 5));

    free(captured);
}

void test_copy_simple_exact_flush(void)
{
    setup_buffer_and_blocksize((idx_t)4);
    const char *input = "ABCD";
    idx_t nread = (idx_t)strlen(input);

    capture_t cap = begin_capture_stdout();
    copy_simple(input, nread);
    unsigned char *captured = NULL;
    size_t captured_len = 0;
    end_capture_stdout(&cap, &captured, &captured_len);

    TEST_ASSERT_EQUAL_UINT32(0, oc);
    TEST_ASSERT_EQUAL_UINT32(4, captured_len);
    TEST_ASSERT_EQUAL_INT(0, memcmp(captured, input, 4));
    TEST_ASSERT_EQUAL_INT(1, w_full);
    TEST_ASSERT_EQUAL_INT(0, w_partial);
    TEST_ASSERT_EQUAL_INT(4, w_bytes);

    free(captured);
}

void test_copy_simple_more_than_block(void)
{
    setup_buffer_and_blocksize((idx_t)4);
    const char *input = "ABCDEFG"; /* 7 bytes -> 1 full block + 3 remainder */
    idx_t nread = (idx_t)strlen(input);

    capture_t cap = begin_capture_stdout();
    copy_simple(input, nread);
    unsigned char *captured = NULL;
    size_t captured_len = 0;
    end_capture_stdout(&cap, &captured, &captured_len);

    TEST_ASSERT_EQUAL_UINT32(3, oc);
    TEST_ASSERT_EQUAL_UINT32(4, captured_len);
    TEST_ASSERT_EQUAL_INT(0, memcmp(captured, "ABCD", 4));
    TEST_ASSERT_EQUAL_INT(0, memcmp(obuf, "EFG", 3));
    TEST_ASSERT_EQUAL_INT(1, w_full);
    TEST_ASSERT_EQUAL_INT(0, w_partial);
    TEST_ASSERT_EQUAL_INT(4, w_bytes);

    free(captured);
}

void test_copy_simple_with_existing_oc_and_partial_flush(void)
{
    setup_buffer_and_blocksize((idx_t)5);

    /* Preload obuf with 3 bytes to simulate partially filled buffer. */
    oc = 3;
    obuf[0] = 'A';
    obuf[1] = 'B';
    obuf[2] = 'C';

    const char *input = "DEFGH"; /* First 2 complete block, then 3 leftover */
    idx_t nread = (idx_t)strlen(input);

    capture_t cap = begin_capture_stdout();
    copy_simple(input, nread);
    unsigned char *captured = NULL;
    size_t captured_len = 0;
    end_capture_stdout(&cap, &captured, &captured_len);

    TEST_ASSERT_EQUAL_UINT32(3, oc);
    TEST_ASSERT_EQUAL_UINT32(5, captured_len);
    TEST_ASSERT_EQUAL_INT(0, memcmp(captured, "ABCDE", 5));
    TEST_ASSERT_EQUAL_INT(0, memcmp(obuf, "FGH", 3));
    TEST_ASSERT_EQUAL_INT(1, w_full);
    TEST_ASSERT_EQUAL_INT(0, w_partial);
    TEST_ASSERT_EQUAL_INT(5, w_bytes);

    free(captured);
}

void test_copy_simple_zero_length(void)
{
    setup_buffer_and_blocksize((idx_t)5);

    /* Preload oc and content; should remain unchanged. */
    oc = 2;
    obuf[0] = 'Z';
    obuf[1] = 'Y';

    capture_t cap = begin_capture_stdout();
    copy_simple("ignored", (idx_t)0);
    unsigned char *captured = NULL;
    size_t captured_len = 0;
    end_capture_stdout(&cap, &captured, &captured_len);

    TEST_ASSERT_EQUAL_UINT32(2, oc);
    TEST_ASSERT_EQUAL_UINT32(0, captured_len);
    TEST_ASSERT_EQUAL_INT(0, w_full);
    TEST_ASSERT_EQUAL_INT(0, w_partial);
    TEST_ASSERT_EQUAL_INT(0, w_bytes);
    TEST_ASSERT_EQUAL_INT('Z', obuf[0]);
    TEST_ASSERT_EQUAL_INT('Y', obuf[1]);

    free(captured);
}

void test_copy_simple_multiple_full_blocks_and_remainder(void)
{
    setup_buffer_and_blocksize((idx_t)3);
    const char *input = "ABCDEFGHIJK"; /* 11 bytes: 3 blocks (9 bytes) + 2 remainder */
    idx_t nread = (idx_t)strlen(input);

    capture_t cap = begin_capture_stdout();
    copy_simple(input, nread);
    unsigned char *captured = NULL;
    size_t captured_len = 0;
    end_capture_stdout(&cap, &captured, &captured_len);

    TEST_ASSERT_EQUAL_UINT32(2, oc);
    TEST_ASSERT_EQUAL_UINT32(9, captured_len);
    TEST_ASSERT_EQUAL_INT(0, memcmp(captured, "ABCDEFGHI", 9));
    TEST_ASSERT_EQUAL_INT(0, memcmp(obuf, "JK", 2));
    TEST_ASSERT_EQUAL_INT(3, w_full);
    TEST_ASSERT_EQUAL_INT(0, w_partial);
    TEST_ASSERT_EQUAL_INT(9, w_bytes);

    free(captured);
}

int main(void)
{
    UNITY_BEGIN();
    RUN_TEST(test_copy_simple_partial_no_flush);
    RUN_TEST(test_copy_simple_exact_flush);
    RUN_TEST(test_copy_simple_more_than_block);
    RUN_TEST(test_copy_simple_with_existing_oc_and_partial_flush);
    RUN_TEST(test_copy_simple_zero_length);
    RUN_TEST(test_copy_simple_multiple_full_blocks_and_remainder);
    return UNITY_END();
}