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

/* The function under test and relevant globals are defined in the dd source
   file that includes this test. We reference them directly here. */

/* Snapshots of global state to restore between tests */
static int saved_conversions_mask;
static int saved_output_flags;
static idx_t saved_output_blocksize;
static bool saved_o_nocache;
static bool saved_o_nocache_eof;
static bool saved_final_op_was_seek;
static int saved_status_level;

void setUp(void) {
    saved_conversions_mask = conversions_mask;
    saved_output_flags = output_flags;
    saved_output_blocksize = output_blocksize;
    saved_o_nocache = o_nocache;
    saved_o_nocache_eof = o_nocache_eof;
    saved_final_op_was_seek = final_op_was_seek;
    saved_status_level = status_level;
}

void tearDown(void) {
    conversions_mask = saved_conversions_mask;
    output_flags = saved_output_flags;
    output_blocksize = saved_output_blocksize;
    o_nocache = saved_o_nocache;
    o_nocache_eof = saved_o_nocache_eof;
    final_op_was_seek = saved_final_op_was_seek;
    status_level = saved_status_level;
}

/* Helper to create a temporary file; returns fd, and writes the path into buf. */
static int helper_mkstemp(char *buf, size_t buflen) {
    /* Use a template under /tmp */
    const char *prefix = "/tmp/dd_iwrite_test_";
    size_t plen = strlen(prefix);
    TEST_ASSERT_TRUE_MESSAGE(buflen > plen + 7, "Buffer too small for mkstemp template");
    memcpy(buf, prefix, plen);
    memcpy(buf + plen, "XXXXXX", 7); /* includes NUL */
    int fd = mkstemp(buf);
    TEST_ASSERT_TRUE_MESSAGE(fd >= 0, "mkstemp failed");
    return fd;
}

/* Basic write to a regular file should write all data and not set final_op_was_seek. */
void test_iwrite_basic_regular_file(void) {
    char path[128];
    int fd = helper_mkstemp(path, sizeof(path));

    conversions_mask = 0;
    output_flags = 0;
    final_op_was_seek = false;

    const char *data = "abc123XYZ";
    idx_t len = (idx_t)strlen(data);

    idx_t ret = iwrite(fd, data, len);
    TEST_ASSERT_EQUAL_INT((int)len, (int)ret);
    TEST_ASSERT_FALSE(final_op_was_seek);

    /* Verify file contents */
    ssize_t r;
    char buf[64] = {0};
    lseek(fd, 0, SEEK_SET);
    r = read(fd, buf, sizeof(buf));
    TEST_ASSERT_EQUAL_INT((int)len, (int)r);
    TEST_ASSERT_EQUAL_INT(0, memcmp(buf, data, (size_t)len));

    close(fd);
    unlink(path);
}

/* With conv=sparse on a seekable fd, iwrite should seek instead of writing
   for all-NUL buffers, set final_op_was_seek, and not extend file size. */
void test_iwrite_sparse_regular_file_seek(void) {
    char path[128];
    int fd = helper_mkstemp(path, sizeof(path));

    conversions_mask = C_SPARSE; /* enable sparse */
    output_flags = 0;
    final_op_was_seek = false;

    const idx_t size = 4096;
    char *zeros = (char *)calloc(1, (size_t)size);
    TEST_ASSERT_NOT_NULL(zeros);

    idx_t ret = iwrite(fd, zeros, size);
    TEST_ASSERT_EQUAL_INT((int)size, (int)ret);
    TEST_ASSERT_TRUE(final_op_was_seek);

    /* File size should remain 0, offset should advance by size */
    struct stat st;
    TEST_ASSERT_EQUAL_INT(0, fstat(fd, &st));
    TEST_ASSERT_EQUAL_INT(0, (int)st.st_size);

    off_t off = lseek(fd, 0, SEEK_CUR);
    TEST_ASSERT_EQUAL_INT((int)size, (int)off);

    free(zeros);
    close(fd);
    unlink(path);
}

/* With conv=sparse on a non-seekable fd (pipe), iwrite should fall back
   to writing actual bytes, and clear C_SPARSE in conversions_mask. */
void test_iwrite_sparse_nonseekable_pipe_falls_back_and_clears_flag(void) {
    int fds[2];
    TEST_ASSERT_EQUAL_INT(0, pipe(fds));

    conversions_mask = C_SPARSE;
    output_flags = 0;
    final_op_was_seek = false;

    const idx_t size = 1024;
    char *zeros = (char *)calloc(1, (size_t)size);
    TEST_ASSERT_NOT_NULL(zeros);

    idx_t ret = iwrite(fds[1], zeros, size);
    TEST_ASSERT_EQUAL_INT((int)size, (int)ret);
    TEST_ASSERT_FALSE(final_op_was_seek);

    /* C_SPARSE should be cleared after fallback */
    TEST_ASSERT_EQUAL_INT(0, (conversions_mask & C_SPARSE));

    /* Read back and verify zeros were actually written */
    idx_t total = 0;
    char buf[512];
    while (total < size) {
        ssize_t n = read(fds[0], buf, sizeof(buf));
        TEST_ASSERT_TRUE_MESSAGE(n >= 0, "read from pipe failed");
        if (n == 0) break; /* shouldn't happen before size bytes read */
        /* Verify the chunk is all zeros */
        for (ssize_t i = 0; i < n; i++) {
            TEST_ASSERT_EQUAL_UINT8(0, (unsigned char)buf[i]);
        }
        total += (idx_t)n;
    }
    TEST_ASSERT_EQUAL_INT((int)size, (int)total);

    free(zeros);
    close(fds[0]);
    close(fds[1]);
}

/* Zero-size write should be a no-op. */
void test_iwrite_zero_size_noop(void) {
    char path[128];
    int fd = helper_mkstemp(path, sizeof(path));

    conversions_mask = 0;
    output_flags = 0;
    final_op_was_seek = false;

    const char *data = "ignored";
    idx_t ret = iwrite(fd, data, 0);
    TEST_ASSERT_EQUAL_INT(0, (int)ret);
    TEST_ASSERT_FALSE(final_op_was_seek);

    off_t off = lseek(fd, 0, SEEK_CUR);
    TEST_ASSERT_EQUAL_INT(0, (int)off);

    close(fd);
    unlink(path);
}

/* Attempting to write to a read-only fd should return 0 written (error). */
void test_iwrite_error_on_readonly_fd_returns_0(void) {
    char path[128];
    int wfd = helper_mkstemp(path, sizeof(path));
    close(wfd);

    int rfd = open(path, O_RDONLY);
    TEST_ASSERT_TRUE_MESSAGE(rfd >= 0, "open O_RDONLY failed");

    conversions_mask = 0;
    output_flags = 0;
    final_op_was_seek = false;

    const char *data = "cannot write";
    errno = 0;
    idx_t ret = iwrite(rfd, data, (idx_t)strlen(data));
    int err = errno;

    TEST_ASSERT_EQUAL_INT(0, (int)ret);
    TEST_ASSERT_FALSE(final_op_was_seek);
    TEST_ASSERT_TRUE_MESSAGE(err == EBADF || err == EINVAL || err == EPERM,
                             "Expected EBADF/EINVAL/EPERM for write to read-only fd");

    close(rfd);
    unlink(path);
}

/* Enabling o_nocache should not affect the number of bytes written. */
void test_iwrite_with_o_nocache(void) {
    char path[128];
    int fd = helper_mkstemp(path, sizeof(path));

    conversions_mask = 0;
    output_flags = 0;
    o_nocache = true; /* enable cache invalidation path */
    final_op_was_seek = false;

    const char payload[] = "cache test payload";
    idx_t sz = (idx_t)strlen(payload);

    idx_t ret = iwrite(fd, payload, sz);
    TEST_ASSERT_EQUAL_INT((int)sz, (int)ret);
    TEST_ASSERT_FALSE(final_op_was_seek);

    /* Verify contents */
    char buf[64] = {0};
    lseek(fd, 0, SEEK_SET);
    ssize_t r = read(fd, buf, sizeof(buf));
    TEST_ASSERT_EQUAL_INT((int)sz, (int)r);
    TEST_ASSERT_EQUAL_INT(0, memcmp(buf, payload, (size_t)sz));

    close(fd);
    unlink(path);
}

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_iwrite_basic_regular_file);
    RUN_TEST(test_iwrite_sparse_regular_file_seek);
    RUN_TEST(test_iwrite_sparse_nonseekable_pipe_falls_back_and_clears_flag);
    RUN_TEST(test_iwrite_zero_size_noop);
    RUN_TEST(test_iwrite_error_on_readonly_fd_returns_0);
    RUN_TEST(test_iwrite_with_o_nocache);
    return UNITY_END();
}