using System.Diagnostics.CodeAnalysis;

namespace Compiler;

public sealed class Parser(List<Token> tokens)
{
    public static List<NodeDefinition> Parse(List<Token> tokens)
    {
        return new Parser(tokens).Parse();
    }

    private int index;

    private List<NodeDefinition> Parse()
    {
        var nodes = new List<NodeDefinition>();

        while (Peek() != null)
            nodes.Add(ParseDefinition());

        return nodes;
    }

    private NodeDefinition ParseDefinition()
    {
        var startIndex = index;

        if (TryExpectKeyword(Keyword.Func))
        {
            var name = ExpectIdent();
            var parameters = new List<NodeDefinitionFunc.Param>();

            ExpectSymbol(Symbol.OpenParen);
            while (!TryExpectSymbol(Symbol.CloseParen))
            {
                var paramStartIndex = index;
                var parameterName = ExpectIdent();
                ExpectSymbol(Symbol.Colon);
                var parameterType = ParseType();
                parameters.Add(new NodeDefinitionFunc.Param(TokensFrom(paramStartIndex), parameterName, parameterType));
            }

            ExpectSymbol(Symbol.Colon);
            var returnType = ParseType();

            var body = ParseStatement();

            return new NodeDefinitionFunc(TokensFrom(startIndex), name, parameters, body, returnType);
        }

        throw new Exception("Not a valid definition");
    }

    private NodeStatement ParseStatement()
    {
        var startIndex = index;

        if (TryExpectSymbol(Symbol.OpenCurly))
        {
            var statements = new List<NodeStatement>();
            while (!TryExpectSymbol(Symbol.CloseCurly))
                statements.Add(ParseStatement());

            return new NodeStatementBlock(TokensFrom(startIndex), statements);
        }

        if (TryExpectKeyword(Keyword.Return))
        {
            var value = ParseExpression();
            return new NodeStatementReturn(TokensFrom(startIndex), value);
        }

        if (TryExpectKeyword(Keyword.Let))
        {
            var name = ExpectIdent();
            ExpectSymbol(Symbol.Colon);
            var type = ParseType();
            ExpectSymbol(Symbol.Equal);
            var value = ParseExpression();
            return new NodeStatementVariableDeclaration(TokensFrom(startIndex), name, type, value);
        }

        if (TryExpectKeyword(Keyword.If))
        {
            var condition = ParseExpression();
            var thenBlock = ParseStatement();
            NodeStatement? elseBlock = null;

            if (TryExpectKeyword(Keyword.Else))
                elseBlock = ParseStatement();

            return new NodeStatementIf(TokensFrom(startIndex), condition, thenBlock, elseBlock);
        }

        if (TryExpectKeyword(Keyword.While))
        {
            var condition = ParseExpression();
            var thenBlock = ParseStatement();
            return new NodeStatementWhile(TokensFrom(startIndex), condition, thenBlock);
        }

        var target = ParseExpression();

        if (TryExpectSymbol(Symbol.OpenParen))
        {
            var parameters = new List<NodeExpression>();

            while (!TryExpectSymbol(Symbol.CloseParen))
                parameters.Add(ParseExpression());

            return new NodeStatementFuncCall(TokensFrom(startIndex), target, parameters);
        }

        if (TryExpectSymbol(Symbol.Equal))
        {
            var value = ParseExpression();
            return new NodeStatementAssignment(TokensFrom(startIndex), target, value);
        }

        throw new Exception("Not a valid followup for expression statement");
    }

    private NodeExpression ParseExpression()
    {
        var startIndex = index;

        if (TryExpectIntLiteral(out var intLiteral))
        {
            return new NodeExpressionIntLiteral(TokensFrom(startIndex), intLiteral);
        }

        if (TryExpectStringLiteral(out var stringLiteral))
        {
            return new NodeExpressionStringLiteral(TokensFrom(startIndex), stringLiteral);
        }

        if (TryExpectBoolLiteral(out var boolLiteral))
        {
            return new NodeExpressionBoolLiteral(TokensFrom(startIndex), boolLiteral);
        }

        if (TryExpectIdent(out var ident))
        {
            return new NodeExpressionIdent(TokensFrom(startIndex), ident);
        }

        throw new Exception("Not a valid expression");
    }

    private NodeType ParseType()
    {
        var startIndex = index;

        if (TryExpectSymbol(Symbol.Caret))
        {
            var to = ParseType();
            return new NodeTypePointer(TokensFrom(startIndex), to);
        }

        if (TryExpectKeyword(Keyword.Func))
        {
            var parameters = new List<NodeType>();

            ExpectSymbol(Symbol.OpenParen);
            while (!TryExpectSymbol(Symbol.CloseParen))
            {
                parameters.Add(ParseType());
            }

            ExpectSymbol(Symbol.Colon);
            var returnType = ParseType();

            return new NodeTypeFunc(TokensFrom(startIndex), parameters, returnType);
        }

        if (TryExpectIdent(out var ident))
        {
            switch (ident.Ident)
            {
                case "void":
                    return new NodeTypeVoid(TokensFrom(startIndex));
                case "string":
                    return new NodeTypeString(TokensFrom(startIndex));
                case "bool":
                    return new NodeTypeBool(TokensFrom(startIndex));
                case "i8":
                    return new NodeTypeSInt(TokensFrom(startIndex), 8);
                case "i16":
                    return new NodeTypeSInt(TokensFrom(startIndex), 16);
                case "i32":
                    return new NodeTypeSInt(TokensFrom(startIndex), 32);
                case "i64":
                    return new NodeTypeSInt(TokensFrom(startIndex), 64);
                case "u8":
                    return new NodeTypeUInt(TokensFrom(startIndex), 8);
                case "u16":
                    return new NodeTypeUInt(TokensFrom(startIndex), 16);
                case "u32":
                    return new NodeTypeUInt(TokensFrom(startIndex), 32);
                case "u64":
                    return new NodeTypeUInt(TokensFrom(startIndex), 64);
                default:
                    return new NodeTypeCustom(TokensFrom(startIndex), ident);
            }
        }

        throw new Exception("Not a valid type");
    }

    private List<Token> TokensFrom(int startIndex)
    {
        return tokens.GetRange(startIndex, index - startIndex);
    }

    private bool TryExpectKeyword(Keyword keyword)
    {
        if (Peek() is TokenKeyword token && token.Keyword == keyword)
        {
            Consume();
            return true;
        }

        return false;
    }

    private void ExpectSymbol(Symbol symbol)
    {
        if (Peek() is TokenSymbol token && token.Symbol == symbol)
        {
            Consume();
            return;
        }

        throw new Exception($"Expected symbol '{symbol}'");
    }

    private bool TryExpectSymbol(Symbol symbol)
    {
        if (Peek() is TokenSymbol token && token.Symbol == symbol)
        {
            Consume();
            return true;
        }

        return false;
    }

    private TokenIdent ExpectIdent()
    {
        if (Peek() is TokenIdent token)
        {
            Consume();
            return token;
        }

        throw new Exception("Expected ident");
    }

    private bool TryExpectIdent([NotNullWhen(true)] out TokenIdent? ident)
    {
        if (Peek() is TokenIdent token)
        {
            Consume();
            ident = token;
            return true;
        }

        ident = null;
        return false;
    }

    private bool TryExpectIntLiteral([NotNullWhen(true)] out TokenIntLiteral? intLiteral)
    {
        if (Peek() is TokenIntLiteral token)
        {
            Consume();
            intLiteral = token;
            return true;
        }

        intLiteral = null;
        return false;
    }

    private bool TryExpectStringLiteral([NotNullWhen(true)] out TokenStringLiteral? stringLiteral)
    {
        if (Peek() is TokenStringLiteral token)
        {
            Consume();
            stringLiteral = token;
            return true;
        }

        stringLiteral = null;
        return false;
    }

    private bool TryExpectBoolLiteral([NotNullWhen(true)] out TokenBoolLiteral? boolLiteral)
    {
        if (Peek() is TokenBoolLiteral token)
        {
            Consume();
            boolLiteral = token;
            return true;
        }

        boolLiteral = null;
        return false;
    }

    private Token Consume()
    {
        if (index >= tokens.Count)
            throw new Exception("End of tokens");

        return tokens[index++];
    }

    private Token? Peek(int offset = 0)
    {
        if (index + offset >= tokens.Count)
            return null;

        return tokens[index + offset];
    }
}

public abstract class Node(List<Token> tokens)
{
    public readonly List<Token> Tokens = tokens;
}

public abstract class NodeDefinition(List<Token> tokens) : Node(tokens);

public sealed class NodeDefinitionFunc(List<Token> tokens, TokenIdent name, List<NodeDefinitionFunc.Param> parameters, NodeStatement body, NodeType returnType) : NodeDefinition(tokens)
{
    public readonly TokenIdent Name = name;
    public readonly List<Param> Parameters = parameters;
    public readonly NodeStatement Body = body;
    public readonly NodeType ReturnType = returnType;

    public sealed class Param(List<Token> tokens, TokenIdent name, NodeType type) : Node(tokens)
    {
        public readonly TokenIdent Name = name;
        public readonly NodeType Type = type;
    }
}

public abstract class NodeStatement(List<Token> tokens) : Node(tokens);

public sealed class NodeStatementBlock(List<Token> tokens, List<NodeStatement> statements) : NodeStatement(tokens)
{
    public readonly List<NodeStatement> Statements = statements;
}

public sealed class NodeStatementFuncCall(List<Token> tokens, NodeExpression target, List<NodeExpression> parameters) : NodeStatement(tokens)
{
    public readonly NodeExpression Target = target;
    public readonly List<NodeExpression> Parameters = parameters;
}

internal class NodeStatementReturn(List<Token> tokens, NodeExpression value) : NodeStatement(tokens)
{
    public readonly NodeExpression Value = value;
}

internal class NodeStatementVariableDeclaration(List<Token> tokens, TokenIdent name, NodeType type, NodeExpression value) : NodeStatement(tokens)
{
    public readonly TokenIdent Name = name;
    public readonly NodeType Type = type;
    public readonly NodeExpression Value = value;
}

internal class NodeStatementAssignment(List<Token> tokens, NodeExpression target, NodeExpression value) : NodeStatement(tokens)
{
    public readonly NodeExpression Target = target;
    public readonly NodeExpression Value = value;
}

internal class NodeStatementIf(List<Token> tokens, NodeExpression condition, NodeStatement thenBlock, NodeStatement? elseBlock) : NodeStatement(tokens)
{
    public readonly NodeExpression Condition = condition;
    public readonly NodeStatement ThenBlock = thenBlock;
    public readonly NodeStatement? ElseBlock = elseBlock;
}

internal class NodeStatementWhile(List<Token> tokens, NodeExpression condition, NodeStatement block) : NodeStatement(tokens)
{
    public readonly NodeExpression Condition = condition;
    public readonly NodeStatement Block = block;
}

public abstract class NodeExpression(List<Token> tokens) : Node(tokens);

public sealed class NodeExpressionIntLiteral(List<Token> tokens, TokenIntLiteral value) : NodeExpression(tokens)
{
    public readonly TokenIntLiteral Value = value;
}

public sealed class NodeExpressionStringLiteral(List<Token> tokens, TokenStringLiteral value) : NodeExpression(tokens)
{
    public readonly TokenStringLiteral Value = value;
}

public sealed class NodeExpressionBoolLiteral(List<Token> tokens, TokenBoolLiteral value) : NodeExpression(tokens)
{
    public readonly TokenBoolLiteral Value = value;
}

public sealed class NodeExpressionIdent(List<Token> tokens, TokenIdent value) : NodeExpression(tokens)
{
    public readonly TokenIdent Value = value;
}

public abstract class NodeType(List<Token> tokens) : Node(tokens);

public sealed class NodeTypeVoid(List<Token> tokens) : NodeType(tokens);

public sealed class NodeTypeUInt(List<Token> tokens, int width) : NodeType(tokens)
{
    public readonly int Width = width;
}

public sealed class NodeTypeSInt(List<Token> tokens, int width) : NodeType(tokens)
{
    public readonly int Width = width;
}

public sealed class NodeTypeBool(List<Token> tokens) : NodeType(tokens);

public sealed class NodeTypeString(List<Token> tokens) : NodeType(tokens);

public sealed class NodeTypeCustom(List<Token> tokens, TokenIdent name) : NodeType(tokens)
{
    public readonly TokenIdent Name = name;
}

public sealed class NodeTypePointer(List<Token> tokens, NodeType to) : NodeType(tokens)
{
    public readonly NodeType To = to;
}

public sealed class NodeTypeFunc(List<Token> tokens, List<NodeType> parameters, NodeType returnType) : NodeType(tokens)
{
    public readonly List<NodeType> Parameters = parameters;
    public readonly NodeType ReturnType = returnType;
}