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

/* Access the target symbols from the program under test. */
extern void usage (int status);
extern char const *program_name;

/* Helper to capture child process stdout/stderr when calling usage(status). */
typedef struct {
    char *out_buf;
    char *err_buf;
    int exit_code;
} usage_capture_t;

/* Read all data from fd into a dynamically allocated buffer. */
static char *read_all_from_fd(int fd) {
    size_t cap = 1024;
    size_t len = 0;
    char *buf = (char *)malloc(cap);
    if (!buf) return NULL;

    for (;;) {
        if (len + 512 > cap) {
            size_t new_cap = cap * 2;
            char *new_buf = (char *)realloc(buf, new_cap);
            if (!new_buf) {
                free(buf);
                return NULL;
            }
            buf = new_buf;
            cap = new_cap;
        }
        ssize_t r = read(fd, buf + len, cap - len);
        if (r < 0) {
            if (errno == EINTR) continue;
            free(buf);
            return NULL;
        }
        if (r == 0) break;
        len += (size_t)r;
    }
    /* Null-terminate */
    if (len + 1 > cap) {
        char *new_buf = (char *)realloc(buf, len + 1);
        if (!new_buf) {
            free(buf);
            return NULL;
        }
        buf = new_buf;
    }
    buf[len] = '\0';
    return buf;
}

/* Spawn a child, redirect stdout/stderr, call usage(status), collect outputs and exit code. */
static usage_capture_t run_usage_and_capture(int status_to_call) {
    usage_capture_t cap = {0};
    int pout[2] = {-1, -1};
    int perr[2] = {-1, -1};

    if (pipe(pout) != 0) {
        TEST_FAIL_MESSAGE("Failed to create pipe for stdout");
    }
    if (pipe(perr) != 0) {
        close(pout[0]); close(pout[1]);
        TEST_FAIL_MESSAGE("Failed to create pipe for stderr");
    }

    fflush(stdout);
    fflush(stderr);
    pid_t pid = fork();
    if (pid < 0) {
        close(pout[0]); close(pout[1]);
        close(perr[0]); close(perr[1]);
        TEST_FAIL_MESSAGE("fork() failed");
    }

    if (pid == 0) {
        /* Child: redirect stdout/stderr and call usage(). No Unity calls here. */
        /* Close read ends */
        close(pout[0]);
        close(perr[0]);

        /* Duplicate to stdout/stderr */
        if (dup2(pout[1], STDOUT_FILENO) < 0) _exit(127);
        if (dup2(perr[1], STDERR_FILENO) < 0) _exit(127);

        /* Close original write ends after dup */
        close(pout[1]);
        close(perr[1]);

        /* Call the function; it will exit() */
        usage(status_to_call);

        /* If it ever returns, exit explicitly to avoid running test harness in child */
        _exit(126);
    }

    /* Parent: close write ends, read outputs */
    close(pout[1]);
    close(perr[1]);

    char *out = read_all_from_fd(pout[0]);
    char *err = read_all_from_fd(perr[0]);
    close(pout[0]);
    close(perr[0]);

    if (!out || !err) {
        free(out);
        free(err);
        TEST_FAIL_MESSAGE("Failed to read from pipes");
    }

    int wstatus = 0;
    pid_t w = 0;
    do {
        w = waitpid(pid, &wstatus, 0);
    } while (w < 0 && errno == EINTR);

    if (w < 0) {
        free(out);
        free(err);
        TEST_FAIL_MESSAGE("waitpid() failed");
    }

    int exit_code = -1;
    if (WIFEXITED(wstatus)) {
        exit_code = WEXITSTATUS(wstatus);
    } else if (WIFSIGNALED(wstatus)) {
        /* Encode signal-based termination in a recognizable way. */
        exit_code = 128 + WTERMSIG(wstatus);
    } else {
        exit_code = -1;
    }

    cap.out_buf = out;
    cap.err_buf = err;
    cap.exit_code = exit_code;
    return cap;
}

void setUp(void) {
    /* Ensure predictable, untranslated output. */
    setenv("LC_ALL", "C", 1);
    /* Set default program_name for tests. This is a global const char*. */
    program_name = "mkdir";
}

void tearDown(void) {
    /* Nothing persistent to clean up. */
}

/* Test that usage(EXIT_SUCCESS) prints full help to stdout, nothing to stderr, exit code 0. */
void test_usage_success_prints_full_help_to_stdout_and_exits_zero(void) {
    usage_capture_t cap = run_usage_and_capture(EXIT_SUCCESS);

    TEST_ASSERT_EQUAL_INT_MESSAGE(0, cap.exit_code, "usage(EXIT_SUCCESS) should exit with code 0");
    TEST_ASSERT_TRUE_MESSAGE(cap.out_buf != NULL, "stdout buffer should not be NULL");
    TEST_ASSERT_TRUE_MESSAGE(strlen(cap.out_buf) > 0, "stdout should have content for success usage");
    TEST_ASSERT_TRUE_MESSAGE(strstr(cap.out_buf, "Usage:") != NULL, "stdout should contain 'Usage:'");
    TEST_ASSERT_TRUE_MESSAGE(strstr(cap.out_buf, "mkdir") != NULL, "stdout should contain program name");
    TEST_ASSERT_TRUE_MESSAGE(cap.err_buf != NULL, "stderr buffer should not be NULL");
    TEST_ASSERT_EQUAL_UINT_MESSAGE(0u, (unsigned)strlen(cap.err_buf), "stderr should be empty for success usage");

    free(cap.out_buf);
    free(cap.err_buf);
}

/* Test that usage(nonzero) prints short help to stderr, nothing to stdout, and preserves exit code. */
void test_usage_nonzero_prints_try_help_to_stderr_and_exits_with_status(void) {
    int code = 2;
    usage_capture_t cap = run_usage_and_capture(code);

    TEST_ASSERT_EQUAL_INT_MESSAGE(code, cap.exit_code, "usage(nonzero) should exit with provided status");
    TEST_ASSERT_TRUE_MESSAGE(cap.out_buf != NULL, "stdout buffer should not be NULL");
    TEST_ASSERT_EQUAL_UINT_MESSAGE(0u, (unsigned)strlen(cap.out_buf), "stdout should be empty for nonzero status");
    TEST_ASSERT_TRUE_MESSAGE(cap.err_buf != NULL, "stderr buffer should not be NULL");

    /* Expect the canonical short help message, containing program_name. */
    TEST_ASSERT_TRUE_MESSAGE(strstr(cap.err_buf, "Try 'mkdir --help' for more information.") != NULL,
                             "stderr should contain short 'Try ... --help' message");

    free(cap.out_buf);
    free(cap.err_buf);
}

/* Test that program_name substitution works in the success branch. */
void test_usage_success_reflects_program_name_variable(void) {
    program_name = "mytool";
    usage_capture_t cap = run_usage_and_capture(EXIT_SUCCESS);

    TEST_ASSERT_EQUAL_INT(0, cap.exit_code);
    TEST_ASSERT_TRUE_MESSAGE(strstr(cap.out_buf, "Usage:") != NULL, "stdout should contain 'Usage:'");
    TEST_ASSERT_TRUE_MESSAGE(strstr(cap.out_buf, "mytool") != NULL, "stdout should contain the current program_name");
    TEST_ASSERT_EQUAL_UINT(0u, (unsigned)strlen(cap.err_buf));

    free(cap.out_buf);
    free(cap.err_buf);
}

/* Test that program_name substitution works in the failure branch. */
void test_usage_failure_reflects_program_name_variable(void) {
    program_name = "another-prog";
    usage_capture_t cap = run_usage_and_capture(3);

    TEST_ASSERT_EQUAL_INT(3, cap.exit_code);
    TEST_ASSERT_EQUAL_UINT(0u, (unsigned)strlen(cap.out_buf));
    TEST_ASSERT_TRUE_MESSAGE(strstr(cap.err_buf, "Try 'another-prog --help' for more information.") != NULL,
                             "stderr short help should reference the current program_name");

    free(cap.out_buf);
    free(cap.err_buf);
}

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_usage_success_prints_full_help_to_stdout_and_exits_zero);
    RUN_TEST(test_usage_nonzero_prints_try_help_to_stderr_and_exits_with_status);
    RUN_TEST(test_usage_success_reflects_program_name_variable);
    RUN_TEST(test_usage_failure_reflects_program_name_variable);
    return UNITY_END();
}