#include <stdbool.h>
#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;
    node->typeTag = 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(OrderedScope, "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(OrderedScope, "if-else", parent);
            IdNode *ifBranch = MakeIdTree(ifNode, ifElseNode);
            IdNode *elseBranch = MakeIdNode(OrderedScope, "else", ifElseNode);

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

            return ifElseNode;
        }

        case ForLoop: {
            Node *loopDecl = astNode->children[0];
            Node *loopBody = astNode->children[3];
            IdNode *loopNode = MakeIdNode(OrderedScope, "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: {
            IdNode *declNode = MakeIdNode(Variable, astNode->children[1]->value.string, parent);
            declNode->typeTag = MakeTypeTag(astNode);
            return declNode;
        }

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

        default:
            return NULL;
    }
}

void PrintIdNode(IdNode *node) {
    if (node == NULL) {
        fprintf(stderr, "wraith: Attempted to call PrintIdNode with null value.\n");
        return;
    }

    switch(node->type) {
        case OrderedScope: 
            printf("OrderedScope (%s)\n", node->name); 
            break;
        case UnorderedScope: 
            printf("UnorderedScope (%s)\n", node->name); 
            break;
        case Struct:
            printf("%s : %s\n", node->name, TypeTagToString(node->typeTag)); 
            break;
        case Function:
            printf("%s : Function<%s>\n", node->name, TypeTagToString(node->typeTag));
            break;
        case Variable: 
            printf("%s : %s\n", node->name, TypeTagToString(node->typeTag)); 
            break;
    }
}

void PrintIdTree(IdNode *tree, uint32_t tabCount) {
    if (tree == NULL) {
        fprintf(stderr, "wraith: Attempted to call PrintIdTree on a null value.\n");
        return;
    }

    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* LookdownId(IdNode *root, NodeType targetType, char *targetName) {
    if (root == NULL) {
        fprintf(stderr, "wraith: Attempted to call LookdownId on a null value.\n");
        return NULL;
    }

    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) == 0) {
            result = current;
            break;
        }

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

    free(frontier);
    return result;
}

bool ScopeHasOrdering(IdNode *node) {
    switch (node->type) {
        case OrderedScope:
        case Function:
        case Variable: // this is only technically true
            return true;
        default:
            return false;
    }
}

IdNode* LookupId(IdNode *node, IdNode *prev, char *target) {
    if (node == NULL) {
        return NULL;
    }

    if (strcmp(node->name, target) == 0) {
        return node;
    }

    // If the current node forms an ordered scope then we want to prevent ourselves from looking up
    // identifiers declared after the scope we have just come from.
    uint32_t idxLimit;
    if (prev == NULL || !ScopeHasOrdering(node)) {
        idxLimit = node->childCount;
    } else {
        uint32_t i;
        for (i = 0, idxLimit = 0; i < node->childCount; i++, idxLimit++) {
            if (node->children[i] == prev) {
                break;
            }
        }
    }

    uint32_t i;
    for (i = 0; i < idxLimit; i++) {
        IdNode *child = node->children[i];
        if (strcmp(child->name, target) == 0) {
            return child;
        } 

        if (child->type == Struct) {
            uint32_t j;
            for (j = 0; j < child->childCount; j++) {
                IdNode *grandchild = child->children[j];
                if (strcmp(grandchild->name, target) == 0) {
                    return grandchild;
                } 
            }
        }
    }

    return LookupId(node->parent, node, target);
}