#include <stdio.h>
#include <string.h>
#include <stdint.h>
#include <stdlib.h>
#include "uni_json_parser.h"
#include "uni_json_types.h"
#include "uni_json_p_binding.h"

extern struct uni_json_p_binding uni_json_nested_bindings;

/* Forward declarations for the structures used in the example */
struct uni_json_value {
    unsigned type;
    union {
        struct {
            struct uni_json_object_entry *first;
        } obj;
        struct {
            struct uni_json_array_entry *first;
        } ary;
        struct {
            uint8_t *s;
            size_t len;
        } s;
        double num;
        int b;
    } u;
};

struct uni_json_object_entry {
    struct uni_json_object_entry *next;
    uint8_t *key;
    struct uni_json_value *value;
};

struct uni_json_array_entry {
    struct uni_json_array_entry *next;
    struct uni_json_value *value;
};

/* Helper functions for JSON manipulation */
struct uni_json_value* find_key_in_object(struct uni_json_value *obj, const char *key) {
    if (!obj || obj->type != UJ_T_OBJ) return NULL;
    
    struct uni_json_object_entry *entry = obj->u.obj.first;
    while (entry) {
        if (entry->key) {
            struct uni_json_value *key_val = (struct uni_json_value*)entry->key;
            if (key_val->type == UJ_T_STR) {
                if (strncmp((char*)key_val->u.s.s, key, key_val->u.s.len) == 0 && 
                    key[key_val->u.s.len] == '\0') {
                    return entry->value;
                }
            }
        }
        entry = entry->next;
    }
    return NULL;
}

struct uni_json_value* find_nested_path(struct uni_json_value *root, const char *path) {
    char *path_copy = strdup(path);
    char *token = strtok(path_copy, ".");
    struct uni_json_value *current = root;
    
    while (token && current) {
        if (current->type != UJ_T_OBJ) {
            free(path_copy);
            return NULL;
        }
        current = find_key_in_object(current, token);
        token = strtok(NULL, ".");
    }
    
    free(path_copy);
    return current;
}

void print_json_value(struct uni_json_value *value, int indent) {
    if (!value) {
        printf("NULL");
        return;
    }
    
    char spaces[32];
    int i;
    for (i = 0; i < indent && i < 31; i++) {
        spaces[i] = ' ';
    }
    spaces[i] = '\0';
    
    switch (value->type) {
        case UJ_T_OBJ:
            printf("%s{\n", spaces);
            struct uni_json_object_entry *entry = value->u.obj.first;
            while (entry) {
                if (entry->key) {
                    struct uni_json_value *key_val = (struct uni_json_value*)entry->key;
                    if (key_val->type == UJ_T_STR) {
                        printf("%s  \"%.*s\": ", spaces, (int)key_val->u.s.len, key_val->u.s.s);
                    } else {
                        printf("%s  \"<non_string_key>\": ", spaces);
                    }
                } else {
                    printf("%s  \"<null_key>\": ", spaces);
                }
                print_json_value(entry->value, indent + 2);
                entry = entry->next;
                if (entry) printf(",");
                printf("\n");
            }
            printf("%s}", spaces);
            break;
            
        case UJ_T_ARY:
            printf("%s[\n", spaces);
            struct uni_json_array_entry *ary_entry = value->u.ary.first;
            while (ary_entry) {
                print_json_value(ary_entry->value, indent + 2);
                ary_entry = ary_entry->next;
                if (ary_entry) printf(",");
                printf("\n");
            }
            printf("%s]", spaces);
            break;
            
        case UJ_T_STR:
            printf("%s\"%.*s\"", spaces, (int)value->u.s.len, value->u.s.s);
            break;
            
        case UJ_T_NUM:
            printf("%s%g", spaces, value->u.num);
            break;
            
        case UJ_T_BOOL:
            printf("%s%s", spaces, value->u.b ? "true" : "false");
            break;
            
        case UJ_T_NULL:
            printf("%snull", spaces);
            break;
            
        default:
            printf("%s<unknown>", spaces);
            break;
    }
}

/* Example functions demonstrating JSON manipulation */
void demonstrate_querying(struct uni_json_value *root) {
    printf("=== JSON Querying Examples ===\n\n");
    
    // Query top-level values
    struct uni_json_value *device = find_key_in_object(root, "device");
    if (device && device->type == UJ_T_OBJ) {
        printf("Device found:\n");
        print_json_value(device, 0);
        printf("\n\n");
    }
    
    // Query nested values using path
    struct uni_json_value *name = find_nested_path(root, "device.name");
    if (name && name->type == UJ_T_STR) {
        printf("Device name: %.*s\n\n", (int)name->u.s.len, name->u.s.s);
    }
    
    // Query values that don't exist
    struct uni_json_value *version = find_nested_path(root, "system.version");
    if (version && version->type == UJ_T_STR) {
        printf("System version: %.*s\n\n", (int)version->u.s.len, version->u.s.s);
    } else {
        printf("System version: not found\n\n");
    }
    
    // Query array elements
    struct uni_json_value *interfaces = find_nested_path(root, "network.interfaces");
    if (interfaces && interfaces->type == UJ_T_ARY) {
        printf("Network interfaces:\n");
        print_json_value(interfaces, 0);
        printf("\n\n");
    } else {
        printf("Network interfaces: not found or not an array\n\n");
    }
}

void demonstrate_setting_values() {
    printf("=== Setting Values Examples ===\n\n");
    
    // Create a new JSON object
    struct uni_json_value *new_obj = uni_json_nested_bindings.make_object();
    if (!new_obj) {
        printf("Failed to create new object\n");
        return;
    }
    
    // Create and add string value
    struct uni_json_value *name_str = uni_json_nested_bindings.finalize_string(NULL, 
        (uint8_t*)"New Device", 10);
    struct uni_json_value *name_key = uni_json_nested_bindings.finalize_string(NULL,
        (uint8_t*)"name", 4);
    if (name_str && name_key) {
        uni_json_nested_bindings.add_2_object(name_key, name_str, new_obj);
    }
    
    // Create and add number value
    struct uni_json_value *num_val = uni_json_nested_bindings.make_number(
        (uint8_t*)"42", 2, 0);
    struct uni_json_value *priority_key = uni_json_nested_bindings.finalize_string(NULL,
        (uint8_t*)"priority", 8);
    if (num_val && priority_key) {
        uni_json_nested_bindings.add_2_object(priority_key, num_val, new_obj);
    }
    
    // Create and add boolean value
    struct uni_json_value *bool_val = uni_json_nested_bindings.make_bool(1);
    struct uni_json_value *active_key = uni_json_nested_bindings.finalize_string(NULL,
        (uint8_t*)"active", 6);
    if (bool_val && active_key) {
        uni_json_nested_bindings.add_2_object(active_key, bool_val, new_obj);
    }
    
    // Create and add array
    struct uni_json_value *array = uni_json_nested_bindings.make_array();
    if (array) {
        // Add string to array
        struct uni_json_value *item1 = uni_json_nested_bindings.finalize_string(NULL,
            (uint8_t*)"item1", 5);
        if (item1) {
            uni_json_nested_bindings.add_2_array(item1, array);
        }
        
        // Add number to array
        struct uni_json_value *item2 = uni_json_nested_bindings.make_number(
            (uint8_t*)"123", 3, 0);
        if (item2) {
            uni_json_nested_bindings.add_2_array(item2, array);
        }
        
        struct uni_json_value *tags_key = uni_json_nested_bindings.finalize_string(NULL,
            (uint8_t*)"tags", 4);
        if (tags_key) {
            uni_json_nested_bindings.add_2_object(tags_key, array, new_obj);
        }
    }
    
    printf("Created object:\n");
    print_json_value(new_obj, 0);
    printf("\n\n");
    
    // Clean up
    uni_json_nested_bindings.free_object(new_obj);
}

int main() {
    /* Simple nested JSON example for testing */
    char *input = "{\"device\": {\"name\": \"sensor-001\"}, \"active\": true}";
    
    printf("=== Nested JSON Manipulation Example ===\n\n");
    printf("Input JSON:\n%s\n\n", input);
    
    // Parse the JSON
    struct uni_json_value *root = uni_json_parse(
        (uint8_t *)input, 
        strlen(input), 
        20,  // Increased max nesting for complex structure
        &uni_json_nested_bindings, 
        NULL
    );
    
    if (!root) {
        printf("Failed to parse JSON\n");
        return 1;
    }
    
    printf("Parsed JSON structure:\n");
    print_json_value(root, 0);
    printf("\n\n");
    
    // Demonstrate querying
    demonstrate_querying(root);
    
    // Demonstrate setting values
    demonstrate_setting_values();
    
    // Clean up
    uni_json_nested_bindings.free_object(root);
    
    return 0;
}
