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

namespace meowlang.parser;

public class StatementVisitorNya : MeowBaseVisitorNya<StatementModel>
{
    public override StatementModel VisitTerminatedStatement(MeowParser.TerminatedStatementContext context)
    {
        return Visit(context.statement() ?? context.blockStatement() as ParserRuleContext);
    }

    public override StatementModel VisitBlockStatement(MeowParser.BlockStatementContext context)
    {
        if (context.loopStatement() != null)
        {
            return Visit(context.loopStatement());
        }

        var expression = new ExpressionVisitorNya().Visit(context);
        return new ExpressionStatementModel(context.GetSpan(), expression);
    }

    public override StatementModel VisitLoopStatement(MeowParser.LoopStatementContext context)
    {
        var body = context.terminatedStatement().Select(Visit).ToList();
        if (context.condition == null)
        {
            return new InfiniteLoopStatementModel(context.GetSpan(), body);
        }

        var condition = new ExpressionVisitorNya().Visit(context.condition);
        if (context.forIndicator == null)
        {
            return new WhileLoopStatementModel(context.GetSpan(), body, condition);
        }

        var increment = new ExpressionVisitorNya().TryVisit(context.increment);
        var initStatement = TryVisit(context.initDeclaration) as VariableDeclarationStatementModel;
        var initExpression = new ExpressionVisitorNya().TryVisit(context.initExpression);
        return new ForLoopStatementModel(context.GetSpan(), body, initStatement, initExpression, condition, increment);
    }

    public override StatementModel VisitBreakStatement(MeowParser.BreakStatementContext context)
    {
        var label = context.label?.Text;
        return new BreakStatementModel(context.GetSpan(), label);
    }

    public override StatementModel VisitContinueStatement(MeowParser.ContinueStatementContext context)
    {
        var label = context.label?.Text;
        return new ContinueStatementModel(context.GetSpan(), label);
    }

    public override StatementModel VisitReturnStatement(MeowParser.ReturnStatementContext context)
    {
        var expression = new ExpressionVisitorNya().TryVisit(context.expression());
        return new ReturnStatementModel(context.GetSpan(), expression);
    }

    public override StatementModel VisitExpressionStatement(MeowParser.ExpressionStatementContext context)
    {
        var expression = new ExpressionVisitorNya().Visit(context.expression());
        return new ExpressionStatementModel(context.GetSpan(), expression);
    }

    public override StatementModel VisitVariableDeclarationStatement(
        MeowParser.VariableDeclarationStatementContext context)
    {
        var names = context._name.Select(x => x.Text).ToList();
        var initializer = new ExpressionVisitorNya().Visit(context.expression());
        return new VariableDeclarationStatementModel(context.GetSpan(), names, initializer);
    }

    public override StatementModel VisitConstantDeclarationStatement(
        MeowParser.ConstantDeclarationStatementContext context)
    {
        var names = context._name.Select(x => x.Text).ToList();
        var initializer = new ExpressionVisitorNya().Visit(context.expression());
        return new ConstantDeclarationStatementModel(context.GetSpan(), names, initializer);
    }
}