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

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

    /* Test: construct and verify pure action = Fork 0 Leaf */
    uint32_t leaf = arb_leaf(ctx);
    uint32_t zero = arb_of_number(ctx, 0);
    uint32_t pure_action = arb_fork(ctx, zero, leaf);

    if (!arb_is_fork(ctx, pure_action)) {
        fprintf(stderr, "FAIL: pure action should be fork\n");
        arboricx_free(ctx);
        return 1;
    }

    uint32_t tag, payload;
    if (!arb_get_fork_children(ctx, pure_action, &tag, &payload) ||
        tag != zero || payload != leaf) {
        fprintf(stderr, "FAIL: pure action children mismatch\n");
        arboricx_free(ctx);
        return 1;
    }

    uint64_t tag_num;
    if (!arb_to_number(ctx, tag, &tag_num) || tag_num != 0) {
        fprintf(stderr, "FAIL: pure action tag should be 0\n");
        arboricx_free(ctx);
        return 1;
    }
    printf("PASS: pure action shape\n");

    /* Test: construct and verify bind action = Fork 1 (Fork left k) */
    uint32_t one = arb_of_number(ctx, 1);
    uint32_t left = arb_fork(ctx, zero, leaf);  /* pure Leaf */
    uint32_t k = arb_fork(ctx, leaf, leaf);     /* identity as Fork Leaf Leaf */
    uint32_t bind_pair = arb_fork(ctx, left, k);
    uint32_t bind_action = arb_fork(ctx, one, bind_pair);

    if (!arb_get_fork_children(ctx, bind_action, &tag, &payload) ||
        !arb_to_number(ctx, tag, &tag_num) || tag_num != 1) {
        fprintf(stderr, "FAIL: bind action tag should be 1\n");
        arboricx_free(ctx);
        return 1;
    }

    uint32_t bind_left, bind_k;
    if (!arb_get_fork_children(ctx, payload, &bind_left, &bind_k) ||
        bind_left != left || bind_k != k) {
        fprintf(stderr, "FAIL: bind payload should be Fork left k\n");
        arboricx_free(ctx);
        return 1;
    }
    printf("PASS: bind action shape\n");

    /* Test: construct and verify IO sentinel = Fork "tricuIO" (Fork 1 action) */
    uint32_t sentinel_str = arb_of_string(ctx, "tricuIO");
    uint32_t version = arb_of_number(ctx, 1);
    uint32_t version_action_pair = arb_fork(ctx, version, pure_action);
    uint32_t io_sentinel = arb_fork(ctx, sentinel_str, version_action_pair);

    if (!arb_is_fork(ctx, io_sentinel)) {
        fprintf(stderr, "FAIL: IO sentinel should be fork\n");
        arboricx_free(ctx);
        return 1;
    }

    uint32_t sent_left, sent_right;
    if (!arb_get_fork_children(ctx, io_sentinel, &sent_left, &sent_right)) {
        fprintf(stderr, "FAIL: get_fork_children on IO sentinel\n");
        arboricx_free(ctx);
        return 1;
    }

    /* Verify sentinel string */
    uint8_t* decoded_sentinel;
    size_t decoded_len;
    if (!arb_to_string(ctx, sent_left, &decoded_sentinel, &decoded_len) ||
        decoded_len != 7 || memcmp(decoded_sentinel, "tricuIO", 7) != 0) {
        fprintf(stderr, "FAIL: IO sentinel string mismatch\n");
        arboricx_free(ctx);
        return 1;
    }
    arboricx_free_buf(ctx, decoded_sentinel, decoded_len);

    /* Verify version = 1 and action = pure */
    uint32_t ver, act;
    if (!arb_get_fork_children(ctx, sent_right, &ver, &act) ||
        !arb_to_number(ctx, ver, &tag_num) || tag_num != 1 ||
        act != pure_action) {
        fprintf(stderr, "FAIL: IO sentinel version/action mismatch\n");
        arboricx_free(ctx);
        return 1;
    }
    printf("PASS: IO sentinel shape\n");

    /* Test: putStr action = Fork 10 string */
    uint32_t ten = arb_of_number(ctx, 10);
    uint32_t msg = arb_of_string(ctx, "hello");
    uint32_t putStr_action = arb_fork(ctx, ten, msg);

    if (!arb_get_fork_children(ctx, putStr_action, &tag, &payload) ||
        !arb_to_number(ctx, tag, &tag_num) || tag_num != 10) {
        fprintf(stderr, "FAIL: putStr tag should be 10\n");
        arboricx_free(ctx);
        return 1;
    }
    printf("PASS: putStr action shape\n");

    /* Test: getLine action = Fork 11 Leaf */
    uint32_t eleven = arb_of_number(ctx, 11);
    uint32_t getLine_action = arb_fork(ctx, eleven, leaf);

    if (!arb_get_fork_children(ctx, getLine_action, &tag, &payload) ||
        !arb_to_number(ctx, tag, &tag_num) || tag_num != 11 ||
        payload != leaf) {
        fprintf(stderr, "FAIL: getLine tag should be 11 with Leaf payload\n");
        arboricx_free(ctx);
        return 1;
    }
    printf("PASS: getLine action shape\n");

    /* Test: readFile action = Fork 20 path */
    uint32_t twenty = arb_of_number(ctx, 20);
    uint32_t path = arb_of_string(ctx, "/tmp/test.txt");
    uint32_t readFile_action = arb_fork(ctx, twenty, path);

    if (!arb_get_fork_children(ctx, readFile_action, &tag, &payload) ||
        !arb_to_number(ctx, tag, &tag_num) || tag_num != 20) {
        fprintf(stderr, "FAIL: readFile tag should be 20\n");
        arboricx_free(ctx);
        return 1;
    }
    printf("PASS: readFile action shape\n");

    /* Test: writeFile action = Fork 21 (Fork path contents) */
    uint32_t twenty_one = arb_of_number(ctx, 21);
    uint32_t contents = arb_of_string(ctx, "data");
    uint32_t write_pair = arb_fork(ctx, path, contents);
    uint32_t writeFile_action = arb_fork(ctx, twenty_one, write_pair);

    if (!arb_get_fork_children(ctx, writeFile_action, &tag, &payload) ||
        !arb_to_number(ctx, tag, &tag_num) || tag_num != 21) {
        fprintf(stderr, "FAIL: writeFile tag should be 21\n");
        arboricx_free(ctx);
        return 1;
    }

    uint32_t wf_path, wf_contents;
    if (!arb_get_fork_children(ctx, payload, &wf_path, &wf_contents) ||
        wf_path != path || wf_contents != contents) {
        fprintf(stderr, "FAIL: writeFile payload should be Fork path contents\n");
        arboricx_free(ctx);
        return 1;
    }
    printf("PASS: writeFile action shape\n");

    /* Test: ok result = Fork (Stem Leaf) (Fork val Leaf) */
    uint32_t stem_leaf = arb_stem(ctx, leaf);
    uint32_t val_pair = arb_fork(ctx, msg, leaf);
    uint32_t ok_result = arb_fork(ctx, stem_leaf, val_pair);

    if (!arb_is_fork(ctx, ok_result)) {
        fprintf(stderr, "FAIL: ok result should be fork\n");
        arboricx_free(ctx);
        return 1;
    }

    uint32_t ok_tag, ok_rest;
    if (!arb_get_fork_children(ctx, ok_result, &ok_tag, &ok_rest) ||
        !arb_is_stem(ctx, ok_tag)) {
        fprintf(stderr, "FAIL: ok result left should be stem\n");
        arboricx_free(ctx);
        return 1;
    }

    uint32_t ok_val, ok_leaf;
    if (!arb_get_fork_children(ctx, ok_rest, &ok_val, &ok_leaf) ||
        ok_val != msg || ok_leaf != leaf) {
        fprintf(stderr, "FAIL: ok result right should be Fork val Leaf\n");
        arboricx_free(ctx);
        return 1;
    }
    printf("PASS: ok result shape\n");

    /* Test: err result = Fork Leaf (Fork code Leaf) */
    uint32_t err_code = arb_of_number(ctx, 42);
    uint32_t err_pair = arb_fork(ctx, err_code, leaf);
    uint32_t err_result = arb_fork(ctx, leaf, err_pair);

    if (!arb_is_fork(ctx, err_result)) {
        fprintf(stderr, "FAIL: err result should be fork\n");
        arboricx_free(ctx);
        return 1;
    }

    uint32_t err_tag, err_rest;
    if (!arb_get_fork_children(ctx, err_result, &err_tag, &err_rest) ||
        !arb_is_leaf(ctx, err_tag)) {
        fprintf(stderr, "FAIL: err result left should be leaf\n");
        arboricx_free(ctx);
        return 1;
    }

    uint32_t err_c, err_l;
    if (!arb_get_fork_children(ctx, err_rest, &err_c, &err_l) ||
        err_c != err_code || err_l != leaf) {
        fprintf(stderr, "FAIL: err result right should be Fork code Leaf\n");
        arboricx_free(ctx);
        return 1;
    }
    printf("PASS: err result shape\n");

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