#include <stdio.h>
#include <string.h>
#include "arboricx.h"

static uint32_t make_pure(arb_ctx_t* ctx, uint32_t val) {
    uint32_t zero = arb_of_number(ctx, 0);
    return arb_fork(ctx, zero, val);
}

static uint32_t make_io_sentinel(arb_ctx_t* ctx, uint32_t action) {
    uint32_t sentinel = arb_of_string(ctx, "tricuIO");
    uint32_t version = arb_of_number(ctx, 1);
    uint32_t version_action = arb_fork(ctx, version, action);
    return arb_fork(ctx, sentinel, version_action);
}

int main(void) {
    arb_ctx_t* ctx = arboricx_init();
    if (!ctx) {
        fprintf(stderr, "Failed to initialize Arboricx context\n");
        return 1;
    }

    arb_io_perms_t perms = { 0, 0 };

    /* Test 1: pure "hello" wrapped in IO sentinel */
    {
        uint32_t hello = arb_of_string(ctx, "hello");
        uint32_t pure_hello = make_pure(ctx, hello);
        uint32_t program = make_io_sentinel(ctx, pure_hello);

        uint32_t result = arb_run_io(ctx, program, &perms);
        if (result == 0) {
            fprintf(stderr, "FAIL: pure hello returned 0\n");
            arboricx_free(ctx);
            return 1;
        }

        uint8_t* decoded;
        size_t decoded_len;
        if (!arb_to_string(ctx, result, &decoded, &decoded_len) ||
            decoded_len != 5 || memcmp(decoded, "hello", 5) != 0) {
            fprintf(stderr, "FAIL: pure hello result mismatch\n");
            arboricx_free(ctx);
            return 1;
        }
        arboricx_free_buf(ctx, decoded, decoded_len);
        printf("PASS: pure hello\n");
    }

    /* Test 2: bind (pure "a") (\_ : pure "done") */
    {
        uint32_t a = arb_of_string(ctx, "a");
        uint32_t done = arb_of_string(ctx, "done");
        uint32_t pure_a = make_pure(ctx, a);
        uint32_t pure_done = make_pure(ctx, done);

        /* K pure_done = Fork Leaf pure_done */
        uint32_t k = arb_fork(ctx, arb_leaf(ctx), pure_done);
        uint32_t bind_pair = arb_fork(ctx, pure_a, k);
        uint32_t one = arb_of_number(ctx, 1);
        uint32_t bind_action = arb_fork(ctx, one, bind_pair);
        uint32_t program = make_io_sentinel(ctx, bind_action);

        uint32_t result = arb_run_io(ctx, program, &perms);
        if (result == 0) {
            fprintf(stderr, "FAIL: bind returned 0\n");
            arboricx_free(ctx);
            return 1;
        }

        uint8_t* decoded;
        size_t decoded_len;
        if (!arb_to_string(ctx, result, &decoded, &decoded_len) ||
            decoded_len != 4 || memcmp(decoded, "done", 4) != 0) {
            fprintf(stderr, "FAIL: bind result mismatch\n");
            arboricx_free(ctx);
            return 1;
        }
        arboricx_free_buf(ctx, decoded, decoded_len);
        printf("PASS: bind pure\n");
    }

    /* Test 3: putStr "test" (no permissions needed) */
    {
        uint32_t test = arb_of_string(ctx, "test");
        uint32_t ten = arb_of_number(ctx, 10);
        uint32_t putStr_action = arb_fork(ctx, ten, test);
        uint32_t program = make_io_sentinel(ctx, putStr_action);

        printf("EXPECT: test\n");
        uint32_t result = arb_run_io(ctx, program, &perms);
        if (result == 0) {
            fprintf(stderr, "FAIL: putStr returned 0\n");
            arboricx_free(ctx);
            return 1;
        }
        if (!arb_is_leaf(ctx, result)) {
            fprintf(stderr, "FAIL: putStr should return Leaf\n");
            arboricx_free(ctx);
            return 1;
        }
        printf("PASS: putStr\n");
    }

    /* Test 4: readFile without permission returns err */
    {
        uint32_t path = arb_of_string(ctx, "/etc/passwd");
        uint32_t twenty = arb_of_number(ctx, 20);
        uint32_t readFile_action = arb_fork(ctx, twenty, path);
        uint32_t program = make_io_sentinel(ctx, readFile_action);

        uint32_t result = arb_run_io(ctx, program, &perms);
        if (result == 0) {
            fprintf(stderr, "FAIL: readFile denied returned 0\n");
            arboricx_free(ctx);
            return 1;
        }

        /* Should be an err result: Fork Leaf (Fork code Leaf) */
        uint32_t left, right;
        if (!arb_get_fork_children(ctx, result, &left, &right) ||
            !arb_is_leaf(ctx, left)) {
            fprintf(stderr, "FAIL: readFile denied should be err result\n");
            arboricx_free(ctx);
            return 1;
        }

        uint32_t code, rest;
        if (!arb_get_fork_children(ctx, right, &code, &rest) ||
            !arb_is_leaf(ctx, rest)) {
            fprintf(stderr, "FAIL: readFile denied err shape mismatch\n");
            arboricx_free(ctx);
            return 1;
        }

        uint64_t code_num;
        if (!arb_to_number(ctx, code, &code_num) || code_num != 20) {
            fprintf(stderr, "FAIL: readFile denied code should be 20, got %llu\n",
                    (unsigned long long)code_num);
            arboricx_free(ctx);
            return 1;
        }
        printf("PASS: readFile denied\n");
    }

    /* Test 5: readFile with permission succeeds */
    {
        /* Create a temp file first */
        const char* tmp = "/tmp/tricu_io_test.txt";
        FILE* f = fopen(tmp, "w");
        if (!f) {
            fprintf(stderr, "FAIL: could not create temp file\n");
            arboricx_free(ctx);
            return 1;
        }
        fprintf(f, "hi");
        fclose(f);

        arb_io_perms_t unsafe_perms = { 1, 0 };
        uint32_t path = arb_of_string(ctx, tmp);
        uint32_t twenty = arb_of_number(ctx, 20);
        uint32_t readFile_action = arb_fork(ctx, twenty, path);
        uint32_t program = make_io_sentinel(ctx, readFile_action);

        uint32_t result = arb_run_io(ctx, program, &unsafe_perms);
        if (result == 0) {
            fprintf(stderr, "FAIL: readFile allowed returned 0\n");
            arboricx_free(ctx);
            return 1;
        }

        /* Should be ok result: Fork (Stem Leaf) (Fork val Leaf) */
        uint32_t ok_tag, ok_rest;
        if (!arb_get_fork_children(ctx, result, &ok_tag, &ok_rest) ||
            !arb_is_stem(ctx, ok_tag)) {
            fprintf(stderr, "FAIL: readFile allowed should be ok result\n");
            arboricx_free(ctx);
            return 1;
        }

        uint32_t val, leaf;
        if (!arb_get_fork_children(ctx, ok_rest, &val, &leaf) ||
            !arb_is_leaf(ctx, leaf)) {
            fprintf(stderr, "FAIL: readFile allowed ok shape mismatch\n");
            arboricx_free(ctx);
            return 1;
        }

        uint8_t* decoded;
        size_t decoded_len;
        if (!arb_to_string(ctx, val, &decoded, &decoded_len) ||
            decoded_len != 2 || memcmp(decoded, "hi", 2) != 0) {
            fprintf(stderr, "FAIL: readFile allowed contents mismatch\n");
            arboricx_free(ctx);
            return 1;
        }
        arboricx_free_buf(ctx, decoded, decoded_len);
        printf("PASS: readFile allowed\n");
    }

    /* Test 6: invalid sentinel returns 0 */
    {
        uint32_t bad = arb_fork(ctx, arb_leaf(ctx), arb_leaf(ctx));
        uint32_t result = arb_run_io(ctx, bad, &perms);
        if (result != 0) {
            fprintf(stderr, "FAIL: invalid sentinel should return 0\n");
            arboricx_free(ctx);
            return 1;
        }
        printf("PASS: invalid sentinel\n");
    }

    arboricx_free(ctx);
    printf("\nAll IO run tests passed.\n");
    return 0;
}