#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>

#include "ast.h"
#include "identcheck.h"

IdNode* MakeIdNode(NodeType type, char *name, IdNode *parent) {
    IdNode *node = (IdNode*)malloc(sizeof(IdNode));
    node->type = type;
    node->name = strdup(name);
    node->parent = parent;
    node->childCount = 0;
    node->childCapacity = 0;
    node->children = NULL;
    return node;
}

void AddChildToNode(IdNode *node, IdNode *child) {
    if (child == NULL) return;

    if (node->children == NULL) {
        node->childCapacity = 2;
        node->children = (IdNode**)malloc(sizeof(IdNode*) * node->childCapacity);
    } else if (node->childCount == node->childCapacity) {
        node->childCapacity *= 2;
        node->children = realloc(node->children, sizeof(IdNode*) * node->childCapacity);
    }

    node->children[node->childCount] = child;
    node->childCount += 1;
}

IdNode* MakeIdTree(Node *astNode, IdNode *parent) {
    uint32_t i;
    switch (astNode->syntaxKind) {
        case Assignment:
            return (astNode->children[0]->syntaxKind == Declaration)
                ? MakeIdTree(astNode->children[0], parent)
                : NULL;

        case IfStatement: {
            Node *stmtSeq = astNode->children[1];
            IdNode *ifNode = MakeIdNode(LexicalScope, "if", parent);
            for (i = 0; i < stmtSeq->childCount; i++) {
                AddChildToNode(ifNode, MakeIdTree(stmtSeq->children[i], ifNode));
            }
            return ifNode;
        }

        case IfElseStatement: {
            Node *ifNode = astNode->children[0];
            Node *elseStmts = astNode->children[1];
            IdNode *ifElseNode = MakeIdNode(LexicalScope, "if-else", parent);
            IdNode *ifBranch = MakeIdTree(ifNode, ifElseNode);
            IdNode *elseBranch = MakeIdNode(LexicalScope, "else", ifElseNode);

            AddChildToNode(ifElseNode, ifBranch);
            AddChildToNode(ifElseNode, elseBranch);
            for (i = 0; i < elseStmts->childCount; i++) {
                AddChildToNode(elseBranch, MakeIdTree(elseStmts->children[i], elseBranch));
            }

            return ifElseNode;
        }

        case ForLoop: {
            Node *loopDecl = astNode->children[0];
            Node *loopBody = astNode->children[3];
            IdNode *loopNode = MakeIdNode(LexicalScope, "for-loop", parent);
            AddChildToNode(loopNode, MakeIdTree(loopDecl, loopNode));
            for (i = 0; i < loopBody->childCount; i++) {
                AddChildToNode(loopNode, MakeIdTree(loopBody->children[i], loopNode));
            }
            return loopNode;
        }

        case Declaration:
            return MakeIdNode(Variable, astNode->children[1]->value.string, parent);

        case StructDeclaration: {
            Node *idNode = astNode->children[0];
            Node *declsNode = astNode->children[1];
            IdNode *structNode = MakeIdNode(Struct, idNode->value.string, parent);
            for (i = 0; i < declsNode->childCount; i++) {
                AddChildToNode(structNode, MakeIdTree(declsNode->children[i], structNode));
            }
            return structNode;
        }

        case FunctionDeclaration: {
            Node *sigNode = astNode->children[0];
            Node *funcNameNode = sigNode->children[0];
            Node *funcArgsNode = sigNode->children[2];
            Node *bodyStatementsNode = astNode->children[1];
            IdNode *funcNode = MakeIdNode(Function, funcNameNode->value.string, parent);
            for (i = 0; i < funcArgsNode->childCount; i++) {
                AddChildToNode(funcNode, MakeIdTree(funcArgsNode->children[i], funcNode));
            }
            for (i = 0; i < bodyStatementsNode->childCount; i++) {
                AddChildToNode(funcNode, MakeIdTree(bodyStatementsNode->children[i], funcNode));
            }
            return funcNode;
        }

        case DeclarationSequence: {
            IdNode *declSeqNode = MakeIdNode(LexicalScope, "", parent);
            for (i = 0; i < astNode->childCount; i++) {
                AddChildToNode(declSeqNode, MakeIdTree(astNode->children[i], declSeqNode));
            }
            return declSeqNode;
        }

        default:
            return NULL;
    }
}

void PrintIdNode(IdNode *node) {
    switch(node->type) {
        case LexicalScope: printf("Scope (%s)\n", node->name); break;
        case Struct: printf("%s : Struct\n", node->name); break;
        case Function: printf("%s : Function\n", node->name); break;
        case Variable: printf("%s : Variable\n", node->name); break;
    }
}

void PrintIdTree(IdNode *tree, uint32_t tabCount) {
    uint32_t i;
    for (i = 0; i < tabCount; i++) {
        printf("| ");
    }

    PrintIdNode(tree);

    for (i = 0; i < tree->childCount; i++) {
        PrintIdTree(tree->children[i], tabCount + 1);
    }
}


int PrintAncestors(IdNode *node) {
    if (node == NULL) return -1;

    int i;
    int indent = 1;
    indent += PrintAncestors(node->parent);
    for (i = 0; i < indent; i++) {
        printf(" ");
    }
    PrintIdNode(node);
    return indent;
}

IdNode* FindId(IdNode *root, NodeType targetType, char *targetName) {
    if (root == NULL) {
        fprintf(stderr, "wraith: Attempted to call FindId on a null value.\n");
        return NULL;
    }

    IdNode *result = NULL;
    IdNode **frontier = (IdNode**)malloc(sizeof(IdNode*));
    frontier[0] = root;
    uint32_t frontierCount = 1;
    uint32_t cursor = 0;

    while (frontierCount > 0 && cursor < frontierCount) {
        IdNode *current = frontier[cursor];

        if (current->type == targetType && strcmp(current->name, targetName) == 0) {
            result = current;
            break;
        }

        int newSize = frontierCount + current->childCount;
        frontier = realloc(frontier, sizeof(IdNode*) * newSize);
        uint32_t i;
        for (i = 0; i < current->childCount; i++) {
            frontier[frontierCount + i] = current->children[i];
        }
        frontierCount += current->childCount;
        cursor += 1;
    }

    free(frontier);
    return result;
}