#include "validation.h"

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

Node *GetIdFromStruct(Node *structDecl)
{
    if (structDecl->syntaxKind != StructDeclaration)
    {
        fprintf(
            stderr,
            "wraith: Attempted to call GetIdFromStruct on node with kind: "
            "%s.\n",
            SyntaxKindString(structDecl->syntaxKind));
        return NULL;
    }

    return structDecl->structDeclaration.identifier;
}

Node *GetIdFromFunction(Node *funcDecl)
{
    if (funcDecl->syntaxKind != FunctionDeclaration)
    {
        fprintf(
            stderr,
            "wraith: Attempted to call GetIdFromFunction on node with kind: "
            "%s.\n",
            SyntaxKindString(funcDecl->syntaxKind));
        return NULL;
    }

    Node *sig = funcDecl->functionDeclaration.functionSignature;
    return sig->functionSignature.identifier;
}

Node *GetIdFromDeclaration(Node *decl)
{
    if (decl->syntaxKind != Declaration)
    {
        fprintf(
            stderr,
            "wraith: Attempted to call GetIdFromDeclaration on node with kind: "
            "%s.\n",
            SyntaxKindString(decl->syntaxKind));
    }

    return decl->declaration.identifier;
}

bool AssignmentHasDeclaration(Node *assign)
{
    return (
        assign->syntaxKind == Assignment &&
        assign->assignmentStatement.left->syntaxKind == Declaration);
}

Node *GetIdFromAssignment(Node *assign)
{
    if (assign->syntaxKind != Assignment)
    {
        fprintf(
            stderr,
            "wraith: Attempted to call GetIdFromAssignment on node with kind: "
            "%s.\n",
            SyntaxKindString(assign->syntaxKind));
    }

    if (AssignmentHasDeclaration(assign))
    {
        return GetIdFromDeclaration(assign->assignmentStatement.left);
    }

    return NULL;
}

bool NodeMayHaveId(Node *node)
{
    switch (node->syntaxKind)
    {
    case StructDeclaration:
    case FunctionDeclaration:
    case Declaration:
    case Assignment:
        return true;
    default:
        return false;
    }
}

Node *TryGetId(Node *node)
{
    switch (node->syntaxKind)
    {
    case Assignment:
        return GetIdFromAssignment(node);
    case Declaration:
        return GetIdFromDeclaration(node);
    case FunctionDeclaration:
        return GetIdFromFunction(node);
    case StructDeclaration:
        return GetIdFromStruct(node);
    default:
        return NULL;
    }
}

Node *LookupFunctionArgId(Node *funcDecl, char *target)
{
    Node *args = funcDecl->functionDeclaration.functionSignature
                     ->functionSignature.arguments;

    uint32_t i;
    for (i = 0; i < args->functionArgumentSequence.count; i += 1)
    {
        Node *arg = args->functionArgumentSequence.sequence[i];
        if (arg->syntaxKind != Declaration)
        {
            fprintf(
                stderr,
                "wraith: Encountered %s node in function signature args "
                "list.\n",
                SyntaxKindString(arg->syntaxKind));
            continue;
        }

        Node *argId = GetIdFromDeclaration(arg);
        if (argId != NULL && strcmp(target, argId->identifier.name) == 0)
            return argId;
    }

    return NULL;
}

Node *LookupStructInternalId(Node *structDecl, char *target)
{
    Node *decls = structDecl->structDeclaration.declarationSequence;

    uint32_t i;
    for (i = 0; i < decls->declarationSequence.count; i += 1)
    {
        Node *match = TryGetId(decls->declarationSequence.sequence[i]);
        if (match != NULL && strcmp(target, match->identifier.name) == 0)
            return match;
    }

    return NULL;
}

Node *InspectNode(Node *node, char *target)
{
    /* If this node may have an identifier declaration inside it, attempt to
     * look up the identifier
     * node itself, returning it if it matches the given target name. */
    if (NodeMayHaveId(node))
    {
        Node *candidateId = TryGetId(node);
        if (candidateId != NULL &&
            strcmp(target, candidateId->identifier.name) == 0)
            return candidateId;
    }

    /* If the candidate node was not the one we wanted, but the node node is a
     * function declaration, it's possible that the identifier we want is one of
     * the function's parameters rather than the function's name itself. */
    if (node->syntaxKind == FunctionDeclaration)
    {
        Node *match = LookupFunctionArgId(node, target);
        if (match != NULL)
            return match;
    }

    /* Likewise if the node node is a struct declaration, inspect the struct's
     * internals
     * to see if a top-level definition is the one we're looking for. */
    if (node->syntaxKind == StructDeclaration)
    {
        Node *match = LookupStructInternalId(node, target);
        if (match != NULL)
            return match;
    }

    return NULL;
}

/* FIXME: Handle staged lookups for AccessExpressions. */
/* FIXME: Similar to above, disallow inspection of struct internals outside of
 *        AccessExpressions. */
Node *LookupId(Node *current, Node *prev, char *target)
{
    if (current == NULL)
        return NULL;

    Node *match;

    /* First inspect the current node to see if it contains the target
     * identifier. */
    match = InspectNode(current, target);
    if (match != NULL)
        return match;

    /* If this is the start of our search, we should not attempt to look at
     * child nodes. Only looking up the AST 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. The same is true for functions. */
    if (prev == NULL)
        return LookupId(current->parent, current, target);

    uint32_t i;
    uint32_t idxLimit;
    switch (current->syntaxKind)
    {
    case DeclarationSequence:
        for (i = 0; i < current->declarationSequence.count; i += 1)
        {
            Node *decl = current->declarationSequence.sequence[i];
            match = InspectNode(decl, target);
            if (match != NULL)
                return match;
        }
        break;
    case StatementSequence:
        idxLimit = current->statementSequence.count;
        for (i = 0; i < current->statementSequence.count; i += 1)
        {
            if (current->statementSequence.sequence[i] == prev)
            {
                idxLimit = i;
                break;
            }
        }

        for (i = 0; i < idxLimit; i += 1)
        {
            Node *stmt = current->statementSequence.sequence[i];
            if (stmt == prev)
                break;

            match = InspectNode(stmt, target);
            if (match != NULL)
                return match;
        }
        break;
    }

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

/* FIXME: This function should be extended to handle multi-stage ID lookups for
 *        AccessExpression nodes. */
/* FIXME: Make this function return an error status object of some kind.
 *        A non-OK status should halt compilation. */
void ValidateIdentifiers(Node *node)
{
    if (node == NULL)
        return;

    /* Skip over generic arguments. They contain Identifiers but are not
     * actually identifiers, they declare types. */
    if (node->syntaxKind == GenericArguments)
        return;

    if (node->syntaxKind != Identifier)
    {
        Recurse(node, *ValidateIdentifiers);
        return;
    }

    char *name = node->identifier.name;
    Node *decl = LookupId(node, NULL, name);
    if (decl == NULL)
    {
        /* FIXME: Express this case as an error with AST information, see the
         *        FIXME comment above. */
        fprintf(
            stderr,
            "wraith: Could not find definition of identifier %s.\n",
            name);
    }
}

/* FIXME: This function should be extended to handle multi-stage ID lookups for
 *        AccessExpression nodes. */
void TagIdentifierTypes(Node *node)
{
    if (node == NULL)
        return;

    switch (node->syntaxKind)
    {
    case AllocExpression:
        node->typeTag = MakeTypeTag(node);
        break;

    case Declaration:
        node->declaration.identifier->typeTag = MakeTypeTag(node);
        break;

    case FunctionDeclaration:
        node->functionDeclaration.functionSignature->functionSignature
            .identifier->typeTag = MakeTypeTag(node);
        break;

    case StructDeclaration:
        node->structDeclaration.identifier->typeTag = MakeTypeTag(node);
        break;

    case GenericArgument:
        node->genericArgument.identifier->typeTag = MakeTypeTag(node);
        break;

    case Identifier:
    {
        if (node->typeTag != NULL)
            return;

        char *name = node->identifier.name;
        Node *declaration = LookupId(node, NULL, name);
        /* FIXME: Remove this case once ValidateIdentifiers returns error status
         * info and halts compilation. See ValidateIdentifiers FIXME. */
        if (declaration == NULL)
        {
            TypeTag *tag = (TypeTag *)malloc(sizeof(TypeTag));
            tag->type = Unknown;
            node->typeTag = tag;
        }
        else
        {
            node->typeTag = declaration->typeTag;
        }
        break;
    }
    }

    Recurse(node, *TagIdentifierTypes);
}

Node *LookupType(Node *current, char *target)
{
    if (current == NULL)
        return NULL;

    switch (current->syntaxKind)
    {
    /* If we've encountered a function declaration, check to see if it's generic
     * and, if so, if one of its type parameters is the target. */
    case FunctionDeclaration:
    {
        Node *typeArgs = current->functionDeclaration.functionSignature
                             ->functionSignature.genericArguments;
        uint32_t i;
        for (i = 0; i < typeArgs->genericArguments.count; i += 1)
        {
            Node *arg = typeArgs->genericArguments.arguments[i];
            Node *argId = arg->genericArgument.identifier;
            char *argName = argId->identifier.name;
            /* note: return the GenericArgument, not the Identifier, so that
             * the caller can differentiate between generics and customs. */
            if (strcmp(target, argName) == 0)
                return arg;
        }

        return LookupType(current->parent, target);
    }

    case StructDeclaration:
    {
        Node *structId = GetIdFromStruct(current);
        if (strcmp(target, structId->identifier.name) == 0)
            return structId;

        return LookupType(current->parent, target);
    }

    /* If we encounter a declaration sequence, search each of its children for
     * struct definitions in case one of them is the target. */
    case DeclarationSequence:
    {
        uint32_t i;
        for (i = 0; i < current->declarationSequence.count; i += 1)
        {
            Node *decl = current->declarationSequence.sequence[i];
            if (decl->syntaxKind == StructDeclaration)
            {
                Node *structId = GetIdFromStruct(decl);
                if (strcmp(target, structId->identifier.name) == 0)
                    return structId;
            }
        }

        return LookupType(current->parent, target);
    }

    default:
        return LookupType(current->parent, target);
    }
}

/* FIXME: This function should be modified to handle type parameters over
 * structs. */
void ConvertCustomsToGenerics(Node *node)
{
    if (node == NULL)
        return;

    switch (node->syntaxKind)
    {
    case Declaration:
    {
        Node *id = node->declaration.identifier;
        Node *type = node->declaration.type->type.typeNode;
        if (type->syntaxKind == CustomTypeNode)
        {
            char *target = id->typeTag->value.customType;
            Node *typeLookup = LookupType(node, target);
            if (typeLookup != NULL && typeLookup->syntaxKind == GenericArgument)
            {
                id->typeTag->type = Generic;
                free(node->declaration.type);
                node->declaration.type =
                    MakeGenericTypeNode(id->typeTag->value.genericType);
            }
        }
        break;
    }

    case FunctionSignature:
    {
        Node *id = node->functionSignature.identifier;
        Node *type = node->functionSignature.type->type.typeNode;
        if (type->syntaxKind == CustomTypeNode)
        {
            char *target = id->typeTag->value.customType;
            Node *typeLookup = LookupType(node, target);
            if (typeLookup != NULL && typeLookup->syntaxKind == GenericArgument)
            {
                id->typeTag->type = Generic;
                free(node->functionSignature.type);
                node->functionSignature.type =
                    MakeGenericTypeNode(id->typeTag->value.genericType);
            }
        }
        break;
    }
    }

    Recurse(node, *ConvertCustomsToGenerics);
}