#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>

/* The following tests assume they are included into dd.c's
   translation unit and thus can access its internal static symbols. */

/* Forward declarations of internal symbols from dd.c used here (all are in the same TU). */
extern int dd_copy(void); /* static in dd.c, but visible here due to same TU inclusion */

/* Helpers: we'll reference internal globals from dd.c directly. */

static long get_page_size_for_tests(void)
{
#ifdef _SC_PAGESIZE
    long ps = sysconf(_SC_PAGESIZE);
    if (ps > 0) return ps;
#endif
#ifdef _SC_PAGE_SIZE
    long ps = sysconf(_SC_PAGE_SIZE);
    if (ps > 0) return ps;
#endif
    return 4096;
}

/* Prototypes for internal globals (available due to same TU). */
extern char const *input_file;
extern char const *output_file;
extern idx_t page_size;
extern idx_t input_blocksize;
extern idx_t output_blocksize;
extern idx_t conversion_blocksize;
extern intmax_t skip_records;
extern idx_t skip_bytes;
extern intmax_t seek_records;
extern intmax_t seek_bytes;
extern bool final_op_was_seek;
extern intmax_t max_records;
extern idx_t max_bytes;
extern int conversions_mask;
extern int input_flags;
extern int output_flags;
extern int status_level;
extern bool translation_needed;
extern intmax_t w_partial;
extern intmax_t w_full;
extern intmax_t r_partial;
extern intmax_t r_full;
extern intmax_t w_bytes;
extern intmax_t reported_w_bytes;
extern bool input_seekable;
extern int input_seek_errno;
extern off_t input_offset;
extern bool warn_partial_read;
extern intmax_t r_truncate;
extern char newline_character;
extern char space_character;
extern char *ibuf;
extern char *obuf;
extern idx_t oc;
extern idx_t col;
extern bool i_nocache, o_nocache;
extern bool i_nocache_eof, o_nocache_eof;
extern ssize_t (*iread_fnc) (int fd, char *buf, idx_t size);

/* Internal functions we will point to. */
extern ssize_t iread (int fd, char *buf, idx_t size);

/* Conversion flags from dd.c enum. */
#ifndef C_TWOBUFS
#define C_TWOBUFS 04000
#endif
#ifndef C_SWAB
#define C_SWAB 0200
#endif
#ifndef C_BLOCK
#define C_BLOCK 010
#endif
#ifndef C_UNBLOCK
#define C_UNBLOCK 020
#endif
#ifndef C_SYNC
#define C_SYNC 02000
#endif
#ifndef STATUS_DEFAULT
#define STATUS_DEFAULT 3
#endif

/* Reset internal dd state to a known baseline before each scenario. */
static void dd_reset_state(void)
{
    input_file = NULL;
    output_file = NULL;
    page_size = (idx_t)get_page_size_for_tests();

    input_blocksize = 0;
    output_blocksize = 0;
    conversion_blocksize = 0;

    skip_records = 0;
    skip_bytes = 0;
    seek_records = 0;
    seek_bytes = 0;

    final_op_was_seek = false;

    max_records = INTMAX_MAX;
    max_bytes = 0;

    conversions_mask = 0;
    input_flags = 0;
    output_flags = 0;

    status_level = STATUS_DEFAULT;
    translation_needed = false;

    w_partial = 0;
    w_full = 0;
    r_partial = 0;
    r_full = 0;
    w_bytes = 0;
    reported_w_bytes = -1;

    input_seekable = false;
    input_seek_errno = 0;
    input_offset = 0;

    warn_partial_read = false;

    r_truncate = 0;

    newline_character = '\n';
    space_character = ' ';

    /* Force fresh allocation on next use. */
    ibuf = NULL;
    obuf = NULL;

    oc = 0;
    col = 0;

    i_nocache = false;
    o_nocache = false;
    i_nocache_eof = false;
    o_nocache_eof = false;

    /* Default read function. */
    iread_fnc = iread;
}

/* Utility to write all bytes to fd. */
static int write_all(int fd, const void* buf, size_t len)
{
    const char* p = (const char*)buf;
    while (len) {
        ssize_t n = write(fd, p, len);
        if (n < 0) {
            if (errno == EINTR) continue;
            return -1;
        }
        if (n == 0) return -1;
        p += n;
        len -= (size_t)n;
    }
    return 0;
}

/* Helper that sets up stdin/stdout redirection, invokes dd_copy, restores stdout,
   and captures the output into a newly malloc'd buffer. On error returns a static
   string; on success returns NULL and sets *out_buf and *out_len. */
static const char* invoke_dd_and_capture(const unsigned char* in, size_t in_len,
                                         unsigned char** out_buf, size_t* out_len)
{
    int saved_stdin = -1, saved_stdout = -1;
    int in_fd = -1, out_fd = -1;
    char in_template[] = "/tmp/dd_test_in_XXXXXX";
    char out_template[] = "/tmp/dd_test_out_XXXXXX";

    *out_buf = NULL;
    *out_len = 0;

    in_fd = mkstemp(in_template);
    if (in_fd < 0) return "mkstemp input failed";
    if (fchmod(in_fd, S_IRUSR | S_IWUSR) != 0) { /* best effort */ }
    if (in_len) {
        if (write_all(in_fd, in, in_len) != 0) {
            close(in_fd);
            unlink(in_template);
            return "writing input failed";
        }
    }
    if (lseek(in_fd, 0, SEEK_SET) < 0) {
        close(in_fd);
        unlink(in_template);
        return "lseek input failed";
    }

    out_fd = mkstemp(out_template);
    if (out_fd < 0) {
        close(in_fd);
        unlink(in_template);
        return "mkstemp output failed";
    }
    if (fchmod(out_fd, S_IRUSR | S_IWUSR) != 0) { /* best effort */ }

    saved_stdin = dup(STDIN_FILENO);
    saved_stdout = dup(STDOUT_FILENO);
    if (saved_stdin < 0 || saved_stdout < 0) {
        if (saved_stdin >= 0) close(saved_stdin);
        if (saved_stdout >= 0) close(saved_stdout);
        close(in_fd);
        close(out_fd);
        unlink(in_template);
        unlink(out_template);
        return "dup save stdio failed";
    }

    if (dup2(in_fd, STDIN_FILENO) < 0) {
        close(saved_stdin);
        close(saved_stdout);
        close(in_fd);
        close(out_fd);
        unlink(in_template);
        unlink(out_template);
        return "dup2 stdin failed";
    }
    if (dup2(out_fd, STDOUT_FILENO) < 0) {
        /* Try restore stdin before returning */
        dup2(saved_stdin, STDIN_FILENO);
        close(saved_stdin);
        close(saved_stdout);
        close(in_fd);
        close(out_fd);
        unlink(in_template);
        unlink(out_template);
        return "dup2 stdout failed";
    }

    /* Close original temp fds; stdio now refer to them. */
    close(in_fd);
    close(out_fd);

    /* Call dd_copy while stdout is redirected; do not use Unity here. */
    int rc = dd_copy();

    /* Restore stdout and stdin before any assertions. */
    if (dup2(saved_stdout, STDOUT_FILENO) < 0) {
        close(saved_stdout);
        close(saved_stdin);
        /* We can't safely report via Unity here, but we can bail. */
        return "restore stdout failed";
    }
    if (dup2(saved_stdin, STDIN_FILENO) < 0) {
        close(saved_stdout);
        close(saved_stdin);
        return "restore stdin failed";
    }
    close(saved_stdout);
    close(saved_stdin);

    /* Re-open the output file for reading. */
    int rd_fd = open(out_template, O_RDONLY);
    if (rd_fd < 0) {
        unlink(in_template);
        unlink(out_template);
        return "open output for read failed";
    }

    /* Read all content */
    struct stat st;
    if (fstat(rd_fd, &st) != 0) {
        close(rd_fd);
        unlink(in_template);
        unlink(out_template);
        return "fstat output failed";
    }

    size_t sz = (size_t)st.st_size;
    unsigned char* buf = (unsigned char*)malloc(sz ? sz : 1);
    if (!buf && sz) {
        close(rd_fd);
        unlink(in_template);
        unlink(out_template);
        return "malloc output buffer failed";
    }

    size_t off = 0;
    while (off < sz) {
        ssize_t n = read(rd_fd, buf + off, sz - off);
        if (n < 0) {
            if (errno == EINTR) continue;
            free(buf);
            close(rd_fd);
            unlink(in_template);
            unlink(out_template);
            return "read output failed";
        }
        if (n == 0) break;
        off += (size_t)n;
    }
    close(rd_fd);

    /* Cleanup temp files. */
    unlink(in_template);
    unlink(out_template);

    if (rc != 0) {
        free(buf);
        return "dd_copy returned non-zero";
    }

    *out_buf = buf;
    *out_len = sz;
    return NULL;
}

void setUp(void) {
    /* leave empty; tests call dd_reset_state() explicitly */
}

void tearDown(void) {
    /* leave empty */
}

/* Test 1: Simple copy without C_TWOBUFS: ibuf==obuf path with direct writes. */
void test_dd_copy_simple_no_twobuffers(void)
{
    dd_reset_state();
    input_blocksize = 4;
    output_blocksize = 4;
    conversions_mask = 0; /* ensure ibuf==obuf */

    const unsigned char in[] = "Hello, world!"; /* 13 bytes */

    unsigned char* out = NULL; size_t out_len = 0;
    const char* err = invoke_dd_and_capture(in, sizeof(in)-1, &out, &out_len);
    if (err) {
        TEST_FAIL_MESSAGE(err);
    }

    TEST_ASSERT_EQUAL_UINT32(sizeof(in)-1, out_len);
    TEST_ASSERT_EQUAL_UINT8_ARRAY(in, out, out_len);
    /* Verify counters */
    TEST_ASSERT_EQUAL_INT64(13, w_bytes);
    TEST_ASSERT_EQUAL_INT64(2, w_full);   /* 4+4 full writes */
    TEST_ASSERT_EQUAL_INT64(1, w_partial);/* final 5 bytes? actually 5? Wait 13 -> 4,4,5 => 1 partial */
    /* And reads mirror writes */
    TEST_ASSERT_EQUAL_INT64(2, r_full);
    TEST_ASSERT_EQUAL_INT64(1, r_partial);

    free(out);
}

/* Test 2: C_TWOBUFS buffered copying crossing output block boundary. */
void test_dd_copy_twobuffers_writes(void)
{
    dd_reset_state();
    input_blocksize = 4;
    output_blocksize = 8;
    conversions_mask = C_TWOBUFS; /* use separate obuf to exercise write_output */

    unsigned char in[20];
    for (size_t i = 0; i < sizeof(in); i++) in[i] = (unsigned char)('A' + (int)(i % 26));

    unsigned char* out = NULL; size_t out_len = 0;
    const char* err = invoke_dd_and_capture(in, sizeof(in), &out, &out_len);
    if (err) {
        TEST_FAIL_MESSAGE(err);
    }

    TEST_ASSERT_EQUAL_UINT32(sizeof(in), out_len);
    TEST_ASSERT_EQUAL_UINT8_ARRAY(in, out, out_len);

    TEST_ASSERT_EQUAL_INT64(20, w_bytes);
    TEST_ASSERT_EQUAL_INT64(2, w_full);    /* two full 8-byte buffered writes */
    TEST_ASSERT_EQUAL_INT64(1, w_partial); /* final 4-byte flush */

    free(out);
}

/* Test 3: conv=swab with odd-length input exercises saved byte across calls. */
void test_dd_copy_swab_odd_length(void)
{
    dd_reset_state();
    input_blocksize = 5;   /* odd to allow testing saved byte */
    output_blocksize = 8;
    conversions_mask = C_TWOBUFS | C_SWAB;

    const unsigned char in[] = { 'a','b','c','d','e' }; /* 5 bytes */
    const unsigned char expected[] = { 'b','a','d','c','e' };

    unsigned char* out = NULL; size_t out_len = 0;
    const char* err = invoke_dd_and_capture(in, sizeof(in), &out, &out_len);
    if (err) {
        TEST_FAIL_MESSAGE(err);
    }

    TEST_ASSERT_EQUAL_UINT32(sizeof(expected), out_len);
    TEST_ASSERT_EQUAL_UINT8_ARRAY(expected, out, out_len);

    free(out);
}

/* Test 4: conv=block with cbs=4: pad and truncate as specified. */
void test_dd_copy_block_truncate(void)
{
    dd_reset_state();
    input_blocksize = 16;
    output_blocksize = 16;
    conversion_blocksize = 4;
    conversions_mask = C_TWOBUFS | C_BLOCK;

    const unsigned char in[] = "ab\nabcdef\nxy"; /* lines: "ab\n", "abcdef\n", "xy" */
    const unsigned char expected[] = "ab  " "abcd" "xy  ";

    unsigned char* out = NULL; size_t out_len = 0;
    const char* err = invoke_dd_and_capture(in, sizeof(in)-1, &out, &out_len);
    if (err) {
        TEST_FAIL_MESSAGE(err);
    }

    TEST_ASSERT_EQUAL_UINT32(sizeof(expected)-1, out_len);
    TEST_ASSERT_EQUAL_UINT8_ARRAY(expected, out, out_len);
    /* Two truncated characters ("ef") in the second line. */
    TEST_ASSERT_EQUAL_INT64(2, r_truncate);

    free(out);
}

/* Test 5: conv=unblock with cbs=4: trailing spaces removed, newline added per record. */
void test_dd_copy_unblock(void)
{
    dd_reset_state();
    input_blocksize = 12;
    output_blocksize = 16;
    conversion_blocksize = 4;
    conversions_mask = C_TWOBUFS | C_UNBLOCK;

    const unsigned char in[] = "ab  cd  ef  "; /* three 4-char records with trailing spaces */
    const unsigned char expected[] = "ab\ncd\nef\n";

    unsigned char* out = NULL; size_t out_len = 0;
    const char* err = invoke_dd_and_capture(in, sizeof(in)-1, &out, &out_len);
    if (err) {
        TEST_FAIL_MESSAGE(err);
    }

    TEST_ASSERT_EQUAL_UINT32(sizeof(expected)-1, out_len);
    TEST_ASSERT_EQUAL_UINT8_ARRAY(expected, out, out_len);

    free(out);
}

/* Test 6: conv=sync: short read should be padded with NULs to ibs. */
void test_dd_copy_sync_padding(void)
{
    dd_reset_state();
    input_blocksize = 5;
    output_blocksize = 16;
    conversions_mask = C_TWOBUFS | C_SYNC;

    const unsigned char in[] = { 'x','y','z' };
    const unsigned char expected[] = { 'x','y','z','\0','\0' };

    unsigned char* out = NULL; size_t out_len = 0;
    const char* err = invoke_dd_and_capture(in, sizeof(in), &out, &out_len);
    if (err) {
        TEST_FAIL_MESSAGE(err);
    }

    TEST_ASSERT_EQUAL_UINT32(sizeof(expected), out_len);
    TEST_ASSERT_EQUAL_UINT8_ARRAY(expected, out, out_len);

    free(out);
}

int main(void)
{
    UNITY_BEGIN();
    RUN_TEST(test_dd_copy_simple_no_twobuffers);
    RUN_TEST(test_dd_copy_twobuffers_writes);
    RUN_TEST(test_dd_copy_swab_odd_length);
    RUN_TEST(test_dd_copy_block_truncate);
    RUN_TEST(test_dd_copy_unblock);
    RUN_TEST(test_dd_copy_sync_padding);
    return UNITY_END();
}