#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 *node = (IdNode*)malloc(sizeof(IdNode));
    node->type = type;
    node->name = strdup(name);
    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* FindId(IdNode *root, NodeType targetType, char *targetName) {
    IdNode *result = NULL;
    IdNode **frontier = (IdNode**)malloc(sizeof(IdNode*));
    frontier[0] = root;
    uint32_t frontierCount = 1;

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

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

        IdNode **temp = 
            (IdNode**)malloc(sizeof(IdNode*) * (frontierCount - 1 + current->childCount));
        uint32_t i;
        for (i = 1; i < frontierCount; i++) {
            temp[i - 1] = frontier[i];
        }
        uint32_t offset = i;
        for (i = 0; i < current->childCount; i++) {
            temp[offset + i] = current->children[i];
        }
        frontier = temp;
        frontierCount += current->childCount - 1;
    }

    free(frontier);
    return result;
}

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

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

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

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

            return ifElseNode;
        }

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

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

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

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

        default:
            return NULL;
    }
}

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

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

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