#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.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 = (IdNode **)realloc(
            node->children,
            sizeof(IdNode *) * node->childCapacity);
    }

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

IdNode *MakeIdTree(Node *astNode, IdNode *parent)
{
    uint32_t i;
    IdNode *mainNode;
    switch (astNode->syntaxKind)
    {
    case AccessExpression:
        AddChildToNode(
            parent,
            MakeIdTree(astNode->accessExpression.accessee, parent));
        AddChildToNode(
            parent,
            MakeIdTree(astNode->accessExpression.accessor, parent));
        return NULL;

    case AllocExpression:
        astNode->typeTag = MakeTypeTag(astNode);
        return NULL;

    case Assignment:
    {
        if (astNode->assignmentStatement.left->syntaxKind == Declaration)
        {
            return MakeIdTree(astNode->assignmentStatement.left, parent);
        }
        else
        {
            AddChildToNode(
                parent,
                MakeIdTree(astNode->assignmentStatement.left, parent));
            AddChildToNode(
                parent,
                MakeIdTree(astNode->assignmentStatement.right, parent));
            return NULL;
        }
    }

    case BinaryExpression:
        AddChildToNode(
            parent,
            MakeIdTree(astNode->binaryExpression.left, parent));
        AddChildToNode(
            parent,
            MakeIdTree(astNode->binaryExpression.right, parent));
        return NULL;

    case Declaration:
    {
        Node *idNode = astNode->declaration.identifier;
        mainNode = MakeIdNode(Variable, idNode->identifier.name, parent);
        mainNode->typeTag = MakeTypeTag(astNode);
        idNode->typeTag = mainNode->typeTag;
        break;
    }

    case DeclarationSequence:
    {
        mainNode = MakeIdNode(UnorderedScope, "", parent);
        for (i = 0; i < astNode->declarationSequence.count; i++)
        {
            AddChildToNode(
                mainNode,
                MakeIdTree(astNode->declarationSequence.sequence[i], mainNode));
        }
        break;
    }

    case ForLoop:
    {
        Node *loopDecl = astNode->forLoop.declaration;
        Node *loopBody = astNode->forLoop.statementSequence;
        mainNode = MakeIdNode(OrderedScope, "for-loop", parent);
        AddChildToNode(mainNode, MakeIdTree(loopDecl, mainNode));
        AddChildToNode(mainNode, MakeIdTree(loopBody, mainNode));
        break;
    }

    case FunctionArgumentSequence:
        for (i = 0; i < astNode->functionArgumentSequence.count; i++)
        {
            AddChildToNode(
                parent,
                MakeIdTree(
                    astNode->functionArgumentSequence.sequence[i],
                    parent));
        }
        return NULL;

    case FunctionCallExpression:
        AddChildToNode(
            parent,
            MakeIdTree(astNode->functionCallExpression.identifier, parent));
        AddChildToNode(
            parent,
            MakeIdTree(
                astNode->functionCallExpression.argumentSequence,
                parent));
        return NULL;

    case FunctionDeclaration:
    {
        Node *sigNode = astNode->functionDeclaration.functionSignature;
        Node *idNode = sigNode->functionSignature.identifier;
        char *funcName = idNode->identifier.name;
        mainNode = MakeIdNode(Function, funcName, parent);
        mainNode->typeTag = MakeTypeTag(astNode);
        idNode->typeTag = mainNode->typeTag;
        MakeIdTree(sigNode->functionSignature.genericArguments, mainNode);
        MakeIdTree(sigNode->functionSignature.arguments, mainNode);
        MakeIdTree(astNode->functionDeclaration.functionBody, mainNode);
        break;
    }

    case FunctionSignatureArguments:
    {
        for (i = 0; i < astNode->functionSignatureArguments.count; i++)
        {
            Node *argNode = astNode->functionSignatureArguments.sequence[i];
            AddChildToNode(parent, MakeIdTree(argNode, parent));
        }
        return NULL;
    }

    case GenericArgument:
    {
        char *name = astNode->genericArgument.identifier->identifier.name;
        mainNode = MakeIdNode(GenericType, name, parent);
        break;
    }

    case GenericArguments:
    {
        for (i = 0; i < astNode->genericArguments.count; i += 1)
        {
            Node *argNode = astNode->genericArguments.arguments[i];
            AddChildToNode(parent, MakeIdTree(argNode, parent));
        }
        return NULL;
    }

    case Identifier:
    {
        char *name = astNode->identifier.name;
        mainNode = MakeIdNode(Placeholder, name, parent);
        IdNode *lookupNode = LookupId(mainNode, NULL, name);
        if (lookupNode == NULL)
        {
            fprintf(stderr, "wraith: Could not find IdNode for id %s\n", name);
            TypeTag *tag = (TypeTag *)malloc(sizeof(TypeTag));
            tag->type = Unknown;
            astNode->typeTag = tag;
        }
        else
        {
            astNode->typeTag = lookupNode->typeTag;
        }
        break;
    }

    case IfStatement:
    {
        Node *clause = astNode->ifStatement.expression;
        Node *stmtSeq = astNode->ifStatement.statementSequence;
        mainNode = MakeIdNode(OrderedScope, "if", parent);
        MakeIdTree(clause, mainNode);
        MakeIdTree(stmtSeq, mainNode);
        break;
    }

    case IfElseStatement:
    {
        Node *ifNode = astNode->ifElseStatement.ifStatement;
        Node *elseStmts = astNode->ifElseStatement.elseStatement;
        mainNode = MakeIdNode(OrderedScope, "if-else", parent);
        IdNode *ifBranch = MakeIdTree(ifNode, mainNode);
        AddChildToNode(mainNode, ifBranch);
        IdNode *elseScope = MakeIdNode(OrderedScope, "else", mainNode);
        MakeIdTree(elseStmts, elseScope);
        AddChildToNode(mainNode, elseScope);
        break;
    }

    case ReferenceTypeNode:
        AddChildToNode(parent, MakeIdTree(astNode->referenceType.type, parent));
        return NULL;

    case Return:
        AddChildToNode(
            parent,
            MakeIdTree(astNode->returnStatement.expression, parent));
        return NULL;

    case StatementSequence:
    {
        for (i = 0; i < astNode->statementSequence.count; i++)
        {
            Node *argNode = astNode->statementSequence.sequence[i];
            AddChildToNode(parent, MakeIdTree(argNode, parent));
        }
        return NULL;
    }

    case StructDeclaration:
    {
        Node *idNode = astNode->structDeclaration.identifier;
        Node *declsNode = astNode->structDeclaration.declarationSequence;
        mainNode = MakeIdNode(Struct, idNode->identifier.name, parent);
        mainNode->typeTag = MakeTypeTag(astNode);
        for (i = 0; i < declsNode->declarationSequence.count; i++)
        {
            Node *decl = declsNode->declarationSequence.sequence[i];
            AddChildToNode(mainNode, MakeIdTree(decl, mainNode));
        }
        break;
    }

    case Type:
        AddChildToNode(parent, MakeIdTree(astNode->type.typeNode, parent));
        return NULL;

    case UnaryExpression:
        AddChildToNode(
            parent,
            MakeIdTree(astNode->unaryExpression.child, parent));
        return NULL;

    case Comment:
    case CustomTypeNode:
    case FunctionModifiers:
    case FunctionSignature:
    case Number:
    case PrimitiveTypeNode:
    case ReturnVoid:
    case StaticModifier:
    case StringLiteral:
        return NULL;
    }

    astNode->idLink = mainNode;
    return mainNode;
}

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

    switch (node->type)
    {
    case Placeholder:
        printf("Placeholder (%s)\n", node->name);
        break;
    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;
    case GenericType:
        printf("Generic type: %s\n", node->name);
        break;
    case Alloc:
        printf("Alloc: %s\n", 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 = (IdNode **)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 && node->type != Placeholder)
    {
        return node;
    }

    /* If this is the start of our search, we should not attempt to look at
     * child nodes. Only looking up the scope tree is valid at this point.
     *
     * This has the notable side-effect that this function will return NULL if
     * you attempt to look up a struct's internals starting from the node
     * representing the struct itself. This is because an IdNode corresponds to
     * the location *where an identifier is first declared.* Thus, an identifier
     * has no knowledge of identifiers declared "inside" of it.
     */
    if (prev == NULL)
    {
        return LookupId(node->parent, node, target);
    }

    /* 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 (ScopeHasOrdering(node))
    {
        uint32_t i;
        for (i = 0, idxLimit = 0; i < node->childCount; i++, idxLimit++)
        {
            if (node->children[i] == prev)
            {
                break;
            }
        }
    }
    else
    {
        idxLimit = node->childCount;
    }

    uint32_t i;
    for (i = 0; i < idxLimit; i++)
    {
        IdNode *child = node->children[i];
        if (child == prev || child->type == Placeholder)
        {
            /* Do not inspect the node we just came from or placeholders. */
            continue;
        }

        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);
}