#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 <fcntl.h>
#include <errno.h>
#include <stdbool.h>

/* Helper structure to hold captured results from a child process run of usage(). */
struct CaptureResult {
    int exited;          /* boolean: 1 if child exited normally */
    int exit_status;     /* exit code if exited == 1 */
    char *out;           /* captured stdout */
    size_t out_len;
    char *err;           /* captured stderr */
    size_t err_len;
};

/* Read all data from fd into a malloc'd buffer, return pointer and set *len. */
static char *read_all_from_fd(int fd, size_t *len_out) {
    const size_t chunk = 4096;
    size_t cap = 0;
    size_t len = 0;
    char *buf = NULL;

    for (;;) {
        if (len + chunk + 1 > cap) {
            size_t new_cap = cap ? cap * 2 : 8192;
            if (new_cap < len + chunk + 1) new_cap = len + chunk + 1;
            char *nb = (char *)realloc(buf, new_cap);
            if (!nb) {
                free(buf);
                *len_out = 0;
                return NULL;
            }
            buf = nb;
            cap = new_cap;
        }
        ssize_t r = read(fd, buf + len, chunk);
        if (r < 0) {
            if (errno == EINTR) continue;
            /* On read error, stop and return what we have. */
            break;
        }
        if (r == 0) break;
        len += (size_t)r;
    }
    if (!buf) {
        buf = (char *)malloc(1);
        if (!buf) {
            *len_out = 0;
            return NULL;
        }
    }
    buf[len] = '\0';
    *len_out = len;
    return buf;
}

/* Run usage(status_arg) in a child process, capturing stdout/stderr and exit status. */
static struct CaptureResult run_usage_and_capture(int status_arg) {
    int out_pipe[2];
    int err_pipe[2];
    struct CaptureResult res;
    memset(&res, 0, sizeof(res));

    if (pipe(out_pipe) != 0) {
        /* Fatal setup error; emulate failed child with errno-coded status. */
        res.exited = 1;
        res.exit_status = 127;
        return res;
    }
    if (pipe(err_pipe) != 0) {
        close(out_pipe[0]); close(out_pipe[1]);
        res.exited = 1;
        res.exit_status = 127;
        return res;
    }

    fflush(stdout);
    fflush(stderr);
    pid_t pid = fork();
    if (pid < 0) {
        /* Fork failed. */
        close(out_pipe[0]); close(out_pipe[1]);
        close(err_pipe[0]); close(err_pipe[1]);
        res.exited = 1;
        res.exit_status = 127;
        return res;
    }

    if (pid == 0) {
        /* Child: redirect stdout/stderr and call usage(). */
        /* Close read ends in child. */
        close(out_pipe[0]);
        close(err_pipe[0]);

        /* Redirect stdout and stderr to pipes. */
        if (dup2(out_pipe[1], STDOUT_FILENO) < 0) _exit(127);
        if (dup2(err_pipe[1], STDERR_FILENO) < 0) _exit(127);

        /* Close original write fds after dup. */
        close(out_pipe[1]);
        close(err_pipe[1]);

        /* Call the target function; it should exit. */
        usage(status_arg);

        /* If it returns unexpectedly, ensure child terminates. */
        _exit(255);
    }

    /* Parent: close write ends, read from read ends. */
    close(out_pipe[1]);
    close(err_pipe[1]);

    res.out = read_all_from_fd(out_pipe[0], &res.out_len);
    res.err = read_all_from_fd(err_pipe[0], &res.err_len);
    close(out_pipe[0]);
    close(err_pipe[0]);

    int status = 0;
    if (waitpid(pid, &status, 0) < 0) {
        res.exited = 0;
        res.exit_status = -1;
        return res;
    }

    if (WIFEXITED(status)) {
        res.exited = 1;
        res.exit_status = WEXITSTATUS(status);
    } else {
        res.exited = 0;
        res.exit_status = -1;
    }

    return res;
}

void setUp(void) {
    /* No global setup needed. */
}

void tearDown(void) {
    /* No global teardown needed. */
}

/* Test that usage(EXIT_SUCCESS) writes help to stdout only, not stderr,
   and exits with status 0. */
void test_usage_success_outputs_to_stdout_only_and_mentions_program_name(void) {
    struct CaptureResult r = run_usage_and_capture(EXIT_SUCCESS);

    TEST_ASSERT_TRUE_MESSAGE(r.exited, "Child did not exit normally");
    TEST_ASSERT_EQUAL_INT_MESSAGE(0, r.exit_status, "Exit status should be 0 for success");

    /* Ensure stdout has content and stderr is empty. */
    TEST_ASSERT_NOT_NULL(r.out);
    TEST_ASSERT_TRUE_MESSAGE(r.out_len > 0, "Expected non-empty stdout for success");
    TEST_ASSERT_TRUE_MESSAGE(r.err_len == 0, "Expected empty stderr for success");

    /* The usage text should mention the program name 'du'. */
    bool mentions_du = (r.out && strstr(r.out, "du") != NULL);
    TEST_ASSERT_TRUE_MESSAGE(mentions_du, "Expected usage output to mention 'du'");

    free(r.out);
    free(r.err);
}

/* Test that usage(nonzero) writes brief message to stderr only, not stdout,
   and exits with the provided non-zero status. */
void test_usage_nonzero_status_outputs_to_stderr_only_and_mentions_help(void) {
    int code = 2;
    struct CaptureResult r = run_usage_and_capture(code);

    TEST_ASSERT_TRUE_MESSAGE(r.exited, "Child did not exit normally");
    TEST_ASSERT_EQUAL_INT_MESSAGE(code, r.exit_status, "Exit status should match provided code");

    /* Ensure stderr has content and stdout is empty. */
    TEST_ASSERT_TRUE_MESSAGE(r.out_len == 0, "Expected empty stdout for error status");
    TEST_ASSERT_TRUE_MESSAGE(r.err_len > 0, "Expected non-empty stderr for error status");

    /* The try-help message should include 'du' and '--help' tokens. */
    bool mentions_du = (r.err && strstr(r.err, "du") != NULL);
    bool mentions_help = (r.err && strstr(r.err, "--help") != NULL);
    TEST_ASSERT_TRUE_MESSAGE(mentions_du, "Expected error hint to mention 'du'");
    TEST_ASSERT_TRUE_MESSAGE(mentions_help, "Expected error hint to mention '--help'");

    free(r.out);
    free(r.err);
}

/* Test that arbitrary non-zero statuses are propagated. */
void test_usage_propagates_exact_exit_status(void) {
    int code = 123;
    struct CaptureResult r = run_usage_and_capture(code);

    TEST_ASSERT_TRUE_MESSAGE(r.exited, "Child did not exit normally");
    TEST_ASSERT_EQUAL_INT_MESSAGE(code, r.exit_status, "Exit status should be exactly the provided code");

    /* For non-zero, expect stderr output and empty stdout. */
    TEST_ASSERT_TRUE_MESSAGE(r.out_len == 0, "Expected empty stdout for non-zero status");
    TEST_ASSERT_TRUE_MESSAGE(r.err_len > 0, "Expected non-empty stderr for non-zero status");

    free(r.out);
    free(r.err);
}

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_usage_success_outputs_to_stdout_only_and_mentions_program_name);
    RUN_TEST(test_usage_nonzero_status_outputs_to_stderr_only_and_mentions_help);
    RUN_TEST(test_usage_propagates_exact_exit_status);
    return UNITY_END();
}