using Antlr4.Runtime;
using meowlang.parser.antlr;

namespace meowlang.parser;

public class ExpressionVisitorNya : MeowBaseVisitorNya<ExpressionModel>
{
    public override ExpressionModel VisitAssignmentExpression(MeowParser.AssignmentExpressionContext context)
    {
        var expression = Visit(context.logicalOrExpression());
        if (context.op == null)
        {
            return expression;
        }

        var op = context.op.Text.ToAssignmentOperator();
        var pattern = new PatternModel(context.pattern().GetSpan(),
            new List<ExpressionModel>(context.pattern().pointerDereferenceExpression()
                .Select(Visit).ToList()));
        return new AssignmentExpressionModel(context.GetSpan(), pattern, op, expression);
    }

    private ExpressionModel MakeBinOpTree(List<ExpressionModel> expressions, BinaryOperator op)
    {
        if (expressions.Count == 1) return expressions[0];
        var left = MakeBinOpTree(expressions.Take(expressions.Count - 1).ToList(), op);
        var right = expressions[^1];
        var span = new Span(left.Span.Filename, left.Span.From, right.Span.To, left.Span.Line);
        return new BinaryOperationExpressionModel(span, left, op, right);
    }

    private ExpressionModel MakeBinOpTree(List<ExpressionModel> expressions, List<BinaryOperator> ops)
    {
        if (expressions.Count == 1) return expressions[0];
        var op = ops[^1];
        var right = expressions[^1];
        var remainingExpressions = expressions.Take(expressions.Count - 1).ToList();
        var remainingOps = ops.Take(ops.Count - 1).ToList();
        var left = MakeBinOpTree(remainingExpressions, remainingOps);
        var span = new Span(left.Span.Filename, left.Span.From, right.Span.To, left.Span.Line);
        return new BinaryOperationExpressionModel(span, left, op, right);
    }

    public override ExpressionModel VisitLogicalOrExpression(MeowParser.LogicalOrExpressionContext context)
    {
        var expressions = context.logicalXorExpression().Select(Visit).ToList();
        return MakeBinOpTree(expressions, BinaryOperator.LogicalOr);
    }

    public override ExpressionModel VisitLogicalXorExpression(MeowParser.LogicalXorExpressionContext context)
    {
        var expressions = context.logicalAndExpression().Select(Visit).ToList();
        return MakeBinOpTree(expressions, BinaryOperator.LogicalXor);
    }

    public override ExpressionModel VisitLogicalAndExpression(MeowParser.LogicalAndExpressionContext context)
    {
        var expressions = context.bitwiseOrExpression().Select(Visit).ToList();
        return MakeBinOpTree(expressions, BinaryOperator.LogicalAnd);
    }

    public override ExpressionModel VisitBitwiseOrExpression(MeowParser.BitwiseOrExpressionContext context)
    {
        var expressions = context.bitwiseXorExpression().Select(Visit).ToList();
        return MakeBinOpTree(expressions, BinaryOperator.BitwiseOr);
    }

    public override ExpressionModel VisitBitwiseXorExpression(MeowParser.BitwiseXorExpressionContext context)
    {
        var expressions = context.bitwiseAndExpression().Select(Visit).ToList();
        return MakeBinOpTree(expressions, BinaryOperator.BitwiseXor);
    }

    public override ExpressionModel VisitBitwiseAndExpression(MeowParser.BitwiseAndExpressionContext context)
    {
        var expressions = context.equalityExpression().Select(Visit).ToList();
        return MakeBinOpTree(expressions, BinaryOperator.BitwiseAnd);
    }

    public override ExpressionModel VisitEqualityExpression(MeowParser.EqualityExpressionContext context)
    {
        var expressions = context.relationalExpression().Select(Visit).ToList();
        var operators = context._op.Select(x => x.Text.ToBinaryOperator()).ToList();
        return MakeBinOpTree(expressions, operators);
    }

    public override ExpressionModel VisitRelationalExpression(MeowParser.RelationalExpressionContext context)
    {
        var expressions = context.shiftExpression().Select(Visit).ToList();
        var operators = context._op.Select(x => x.Text.ToBinaryOperator()).ToList();
        return MakeBinOpTree(expressions, operators);
    }

    public override ExpressionModel VisitShiftExpression(MeowParser.ShiftExpressionContext context)
    {
        var expressions = context.addExpression().Select(Visit).ToList();
        var operators = context._op.Select(x => x.Text.Replace(">", ">>").ToBinaryOperator()).ToList();
        return MakeBinOpTree(expressions, operators);
    }

    public override ExpressionModel VisitAddExpression(MeowParser.AddExpressionContext context)
    {
        var expressions = context.multExpression().Select(Visit).ToList();
        var operators = context._op.Select(x => x.Text.ToBinaryOperator()).ToList();
        return MakeBinOpTree(expressions, operators);
    }

    public override ExpressionModel VisitMultExpression(MeowParser.MultExpressionContext context)
    {
        var expressions = context.unaryExpression().Select(Visit).ToList();
        var operators = context._op.Select(x => x.Text.ToBinaryOperator()).ToList();
        return MakeBinOpTree(expressions, operators);
    }

    public override ExpressionModel VisitUnaryExpression(MeowParser.UnaryExpressionContext context)
    {
        var expression = Visit(context.pointerDereferenceExpression());
        foreach (var unaryOperator in context._op.Reverse())
        {
            var span = new Span(expression.Span.Filename, unaryOperator.StartIndex, expression.Span.To, expression.Span.Line);
            if (unaryOperator.Text == "&")
            {
                expression = new PointerReferenceExpressionModel(span, expression);
            }
            else
            {
                expression = new UnaryOperationExpressionModel(span, unaryOperator.Text.ToUnaryOperator(), expression);
            }
        }

        return expression;
    }

    public override ExpressionModel VisitPointerDereferenceExpression(
        MeowParser.PointerDereferenceExpressionContext context)
    {
        var expression = Visit(context.accessExpression());
        foreach (var op in context._op.Reverse())
        {
            var span = new Span(expression.Span.Filename, op.StartIndex, expression.Span.To, expression.Span.Line);
            expression = new PointerDereferenceExpressionModel(span, expression);
        }

        return expression;
    }

    public override ExpressionModel VisitAccessExpression(MeowParser.AccessExpressionContext context)
    {
        var expression = Visit(context.terminalExpression());
        var current = context.accessExpression1();
        while (current != null)
        {
            if (current.subscript != null) // array access
            {
                var span = new Span(expression.Span.Filename, expression.Span.From, current.end.StopIndex, expression.Span.Line);
                var indexer = Visit(current.subscript);
                expression = new ArrayAccessExpressionModel(span, expression, indexer);
            }
            else // struct access
            {
                var span = new Span(expression.Span.Filename, expression.Span.From, current.name.StopIndex, expression.Span.Line);
                var name = current.name.Text;
                if (name.Contains('.') || (char.IsDigit(name[0]) && !char.IsDigit(name.TrimStart('0')[0])))
                {
                    throw new InvalidStructAccessorException(current.name, "");
                }

                expression = new StructAccessExpressionModel(span, expression, name);
            }

            current = current.accessExpression1();
        }

        return expression;
    }

    public override ExpressionModel VisitLiteralExpression(MeowParser.LiteralExpressionContext context)
    {
        var literal = new LiteralVisitorNya().Visit(context.literal());
        return new LiteralExpressionModel(context.GetSpan(), literal);
    }


    public override ExpressionModel VisitNameExpression(MeowParser.NameExpressionContext context)
    {
        var name = new IdentifierNameVisitorNya().Visit(context.identifierName());
        return new NameExpressionModel(context.GetSpan(), name);
    }

    public override ExpressionModel VisitStructConstructor(MeowParser.StructConstructorContext context)
    {
        var typeName = new TypeReferenceVisitorNya().TryVisit(context.typeName()) as TypeNameModel;
        var members = context.structConstructorMember().Select(x =>
            new StructConstructorMemberModel(x.GetSpan(), x.name.Text, Visit(x.expression()))).ToList();
        return new StructConstructorExpressionModel(context.GetSpan(), typeName, members);
    }

    public override ExpressionModel VisitArrayConstructor(MeowParser.ArrayConstructorContext context)
    {
        var type = new TypeReferenceVisitorNya().TryVisit(context.arrayType() ?? context.typeName() as ParserRuleContext);
        var members = context.expression().Select(Visit).ToList();
        return new ArrayConstructorExpressionModel(context.GetSpan(), type, members);
    }

    public override ExpressionModel VisitEnumConstructor(MeowParser.EnumConstructorContext context)
    {
        var typeName = new TypeReferenceVisitorNya().TryVisit(context.typeName()) as TypeNameModel;
        var name = context.name.Text;
        return new EnumConstructorExpressionModel(context.GetSpan(), typeName!, name);
    }

    public override ExpressionModel VisitTupleConstructor(MeowParser.TupleConstructorContext context)
    {
        var values = context.expression().Select(Visit).ToList();
        return new TupleConstructorExpressionModel(context.GetSpan(), values);
    }

    public override ExpressionModel VisitFunctionCall(MeowParser.FunctionCallContext context)
    {
        var name = new IdentifierNameVisitorNya().Visit(context.identifierName());
        var arguments = context.expression().Select(Visit).ToList();
        return new FunctionCallExpressionModel(context.GetSpan(), name, arguments);
    }

    public override ExpressionModel VisitIfExpression(MeowParser.IfExpressionContext context)
    {
        var cond = Visit(context.cond);
        var then = Visit(context.then) as BlockExpressionModel;
        var @else = TryVisit(context.@else) as BlockExpressionModel;
        var ifElses = context.elseIfExpression().Select(x => (Visit(x.cond), (Visit(x.then) as BlockExpressionModel)!))
            .ToList();
        return new IfExpressionModel(context.GetSpan(), cond, then!, ifElses, @else);
    }

    public override ExpressionModel VisitSwitchExpression(MeowParser.SwitchExpressionContext context)
    {
        var expression = Visit(context.expression());
        var cases = context.switchCase().Select(x => (new LiteralVisitorNya().Visit(x.literal()), TryVisit(x.expression()))).ToList();
        var defaultCase = TryVisit(context.switchDefault()?.expression());
        return new SwitchExpressionModel(context.GetSpan(), expression, cases, defaultCase);
    }

    public override ExpressionModel VisitBlockExpression(MeowParser.BlockExpressionContext context)
    {
        var statements = context.terminatedStatement().Select(x => new StatementVisitorNya().Visit(x)).ToList();
        var value = TryVisit(context.expression());
        return new BlockExpressionModel(context.GetSpan(), statements, value);
    }

    public override ExpressionModel VisitParenthesisExpression(MeowParser.ParenthesisExpressionContext context)
    {
        var expression = Visit(context.expression());
        return new ParenthesisExpressionModel(context.GetSpan(), expression);
    }
}