#include "../../unity/unity.h"

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <stdbool.h>
#include <unistd.h>
#include <sys/wait.h>
#include <fcntl.h>
#include <errno.h>

/* usage(int) is defined in the program under test. We will call it in a child. */

static char *read_all_from_fd(int fd, size_t *out_len) {
    if (out_len) *out_len = 0;
    size_t cap = 0;
    size_t len = 0;
    char *buf = NULL;

    for (;;) {
        char tmp[4096];
        ssize_t n = read(fd, tmp, sizeof tmp);
        if (n > 0) {
            if (len + (size_t)n + 1 > cap) {
                size_t newcap = cap ? cap : 64;
                while (len + (size_t)n + 1 > newcap) {
                    newcap *= 2;
                }
                char *nb = (char *)realloc(buf, newcap);
                if (!nb) {
                    free(buf);
                    return NULL;
                }
                buf = nb;
                cap = newcap;
            }
            memcpy(buf + len, tmp, (size_t)n);
            len += (size_t)n;
        } else if (n == 0) {
            break;
        } else {
            if (errno == EINTR) continue;
            free(buf);
            return NULL;
        }
    }

    if (!buf) {
        buf = (char *)malloc(1);
        if (!buf) return NULL;
        buf[0] = '\0';
    } else {
        buf[len] = '\0';
    }
    if (out_len) *out_len = len;
    return buf;
}

typedef struct {
    char *out_data;
    size_t out_len;
    char *err_data;
    size_t err_len;
    int exit_code;
} usage_result;

/* Run usage(status) in a child, capturing stdout/stderr. Returns 0 on success, -1 on error. */
static int run_usage_capture(int status, usage_result *res) {
    if (!res) return -1;
    res->out_data = NULL;
    res->err_data = NULL;
    res->out_len = 0;
    res->err_len = 0;
    res->exit_code = -1;

    int out_pipe[2] = {-1, -1};
    int err_pipe[2] = {-1, -1};

    if (pipe(out_pipe) < 0) return -1;
    if (pipe(err_pipe) < 0) {
        close(out_pipe[0]); close(out_pipe[1]);
        return -1;
    }

    fflush(stdout);
    fflush(stderr);
    pid_t pid = fork();
    if (pid < 0) {
        close(out_pipe[0]); close(out_pipe[1]);
        close(err_pipe[0]); close(err_pipe[1]);
        return -1;
    }

    if (pid == 0) {
        /* Child: redirect stdout/stderr to pipes, set locale, then call usage(status). */
        /* Close read ends in child. */
        close(out_pipe[0]);
        close(err_pipe[0]);

        /* Redirect */
        if (dup2(out_pipe[1], STDOUT_FILENO) < 0) _exit(127);
        if (dup2(err_pipe[1], STDERR_FILENO) < 0) _exit(127);

        /* Close original write ends after dup2 */
        close(out_pipe[1]);
        close(err_pipe[1]);

        /* Ensure predictable locale for messages */
        setenv("LC_ALL", "C", 1);
        setenv("LANGUAGE", "C", 1);

        /* Call the function under test; it will exit(status). */
        usage(status);

        /* Should not reach here, but ensure exit if it does. */
        _exit(255);
    }

    /* Parent: close write ends, read outputs, wait for child. */
    close(out_pipe[1]);
    close(err_pipe[1]);

    size_t out_len = 0, err_len = 0;
    char *out_buf = read_all_from_fd(out_pipe[0], &out_len);
    char *err_buf = read_all_from_fd(err_pipe[0], &err_len);

    /* Close read ends after reading */
    close(out_pipe[0]);
    close(err_pipe[0]);

    int wstatus = 0;
    if (waitpid(pid, &wstatus, 0) < 0) {
        free(out_buf);
        free(err_buf);
        return -1;
    }

    int exit_code = -1;
    if (WIFEXITED(wstatus)) {
        exit_code = WEXITSTATUS(wstatus);
    } else if (WIFSIGNALED(wstatus)) {
        /* encode signal as 128+signal like shells often do */
        exit_code = 128 + WTERMSIG(wstatus);
    }

    res->out_data = out_buf ? out_buf : strdup("");
    res->err_data = err_buf ? err_buf : strdup("");
    res->out_len = out_len;
    res->err_len = err_len;
    res->exit_code = exit_code;

    if (!res->out_data || !res->err_data) {
        free(out_buf);
        free(err_buf);
        return -1;
    }
    return 0;
}

void setUp(void) {
    /* Ensure C locale in parent as well, though we also set it in child. */
    setenv("LC_ALL", "C", 1);
    setenv("LANGUAGE", "C", 1);
}

void tearDown(void) {
    /* No global cleanup needed */
}

/* Test that usage(EXIT_SUCCESS) writes a non-empty usage/help to stdout, nothing to stderr, and exits 0. */
void test_usage_success_prints_to_stdout_and_exits_zero(void) {
    usage_result r;
    int rc = run_usage_capture(EXIT_SUCCESS, &r);
    TEST_ASSERT_EQUAL_INT(0, rc);

    /* Expect exit code 0 */
    TEST_ASSERT_EQUAL_INT(0, r.exit_code);

    /* stdout should contain something, and typically includes "Usage:" in C locale */
    TEST_ASSERT_TRUE(r.out_len > 0);
    TEST_ASSERT_NOT_NULL(r.out_data);
    TEST_ASSERT_NOT_NULL(r.err_data);

    /* stderr should be empty for success case */
    TEST_ASSERT_EQUAL_size_t(0, r.err_len);

    /* In C locale, verify minimal expected token appears */
    TEST_ASSERT_NOT_NULL(strstr(r.out_data, "Usage:"));

    free(r.out_data);
    free(r.err_data);
}

/* Test that nonzero status prints brief try-help to stderr, nothing to stdout, and exits with given status. */
void test_usage_failure_prints_try_help_to_stderr_and_exits_status(void) {
    int status = 2;
    usage_result r;
    int rc = run_usage_capture(status, &r);
    TEST_ASSERT_EQUAL_INT(0, rc);

    TEST_ASSERT_EQUAL_INT(status, r.exit_code);

    /* stderr should have content; stdout should be empty */
    TEST_ASSERT_TRUE(r.err_len > 0);
    TEST_ASSERT_EQUAL_size_t(0, r.out_len);

    /* In C locale, expect a hint about --help */
    TEST_ASSERT_NOT_NULL(strstr(r.err_data, "--help"));

    free(r.out_data);
    free(r.err_data);
}

/* Test EXIT_FAILURE path specifically. */
void test_usage_exit_failure_status_code(void) {
    usage_result r;
    int rc = run_usage_capture(EXIT_FAILURE, &r);
    TEST_ASSERT_EQUAL_INT(0, rc);

    TEST_ASSERT_EQUAL_INT(EXIT_FAILURE, r.exit_code);
    TEST_ASSERT_TRUE(r.err_len > 0);
    TEST_ASSERT_EQUAL_size_t(0, r.out_len);

    free(r.out_data);
    free(r.err_data);
}

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_usage_success_prints_to_stdout_and_exits_zero);
    RUN_TEST(test_usage_failure_prints_try_help_to_stderr_and_exits_status);
    RUN_TEST(test_usage_exit_failure_status_code);
    return UNITY_END();
}