#include "builtins/builtins.h"
#include "parse/parse.h"
#include "precedence.h"
#include "builtins/helpers.h"

using namespace goose;
using namespace goose::ir;
using namespace goose::parse;

namespace goose::builtins
{
    void SetupWhileStmt( Env& e )
    {
        auto handleWhile = []( Parser& p, uint32_t locationId, uint32_t prec )
        {
            auto& dm = DiagnosticsManager::GetInstance();

            if( p.isInParenExpr() || !p.cfg() )
            {
                dm.emitSyntaxErrorMessage( locationId, "the while statement is not allowed here.", 0 );
                return false;
            }

            auto pPrecBB = p.cfg()->currentBB();

            if( !pPrecBB || pPrecBB->terminator() )
            {
                DiagnosticsManager::GetInstance().emitSyntaxErrorMessage(
                    locationId, "unreachable code.", 0 );
                p.cfg()->poison();
            }

            // Create a scope for the entire while, so that any var declared
            // inside of the condition is alive and visible throughout the while.
            Scope s( p );

            auto np = p.makeNestedParser();
            if( !np.parseExpression( precedence::IfStmt ) )
            {
                dm.emitSyntaxErrorMessage( locationId, "expected an expression following the while statement.", 0 );
                return false;
            }

            auto condVal = np.popValue();
            if( !condVal )
            {
                dm.emitSyntaxErrorMessage( locationId, "expected an expression following the while statement.", 0 );
                return false;
            }

            const auto& context = p.resolver()->context();
            auto converted = ConvertValueToType( context, *condVal, GetValueType< bool >() );
            if( holds_alternative< ValUnifyError >( converted ) )
            {
                switch( get< ValUnifyError >( converted ) )
                {
                    // If the condition is invalid, bail out and mark the current
                    // diagnostics context as bust to avoid spamming cascading errors.
                    case ValUnifyError::NoSolution:
                        dm.emitSyntaxErrorMessage( condVal->locationId(), "the while condition can't be converted to a bool." );
                        break;

                    case ValUnifyError::Ambiguous:
                        dm.emitSyntaxErrorMessage( condVal->locationId(), "ambiguous while condition bool conversion." );
                        break;
                }

                return false;
            }

            if( get< Value >( converted ).isPoison() )
                p.cfg()->poison();

            // We need to know if the condition is constant true and this is an infinite loop:
            // it means that no return is required after the while. In practice, that means
            // we don't need to create a successor block and therefore won't leave it unterminated,
            // which will prevent the higher level parsing code to complain about a missing return.
            bool infiniteLoop = false;

            // Try to eagerly evaluate the condition, if it happens to be constant true
            // we need to know about it.
            if( llr::CanValueBeEagerlyEvaluated( get< Value >( converted ) ) )
            {
                execute::VM vm;
                auto cond = Evaluate( get< Value >( converted ), vm );
                if( cond.isPoison() )
                    p.cfg()->poison();
                else if( cond.isConstant() && *FromValue< bool >( cond ) )
                    infiniteLoop = true;
            }

            ptr< llr::BasicBlock > pHeaderBB;

            // If this is an infinite loop, we don't need a header block: we can jump straight to the
            // body.
            if( !infiniteLoop )
                pHeaderBB = p.cfg()->createBB();

            Parser::BreakableScopeGuard bsg( p );
            Parser::ContinuableScopeGuard csg( p );

            auto pBodyBB = ParseSubStatement( p, precedence::IfStmt );
            if( !pBodyBB )
            {
                dm.emitSyntaxErrorMessage( p.resolver()->getCurrentLocation(), "expected a statement after the while condition.", 0 );
                return false;
            }

            if( !pPrecBB )
                return true;

            // If this is an infinite loop, we don't need a header block: we can jump straight to the
            // body.
            if( infiniteLoop )
                pHeaderBB = pBodyBB;

            // Retrieve the final block of the loop body, if any
            if( auto pBodySuccBB = p.cfg()->currentBB(); pBodySuccBB && !pBodySuccBB->terminator() )
            {
                // Jump from the end of the body BB back to the loop header
                pBodySuccBB->setTerminator( llr::Branch( pHeaderBB ) );
            }

            // Jump unconditionally from the pred block to the loop header.
            pPrecBB->setTerminator( llr::Branch( pHeaderBB ) );

            ptr< llr::BasicBlock > pSuccBB;
            if( !infiniteLoop )
                pSuccBB = p.cfg()->createBB();

            auto breakLevel = p.breakableScopeLevels();
            auto continueLevel = p.continuableScopeLevels();

            // Go through all basic blocks, find all break and continue terminators
            // for our scope level and replace them with branches respectively to
            // the successor BB or to the loop header BB.
            p.cfg()->ForEachBB( [&]( auto&& bb )
            {
                const auto& t = bb->terminator();
                if( !t )
                    return;

                if( const auto* pBreak = get_if< llr::Break >( &t->content() ) )
                {
                    if( pBreak->level() == breakLevel )
                    {
                        // If there is a break somewhere, we are not an infinite loop after all.
                        // this doesn't affect that we don't need a header, but it means that we do need
                        // a successor block.
                        if( !pSuccBB )
                            pSuccBB = p.cfg()->createBB();

                        bb->setTerminator( llr::Branch( pSuccBB ) );
                    }
                    return;
                }

                if( const auto* pCont = get_if< llr::Continue >( &t->content() ) )
                {
                    if( pCont->level() == continueLevel )
                        bb->setTerminator( llr::Branch( pHeaderBB ) );
                    return;
                }
            } );

            if( pSuccBB && !infiniteLoop )
            {
                // Emit the conditional branch that will either run an iteration of the loop or exit to the succ bb.
                pHeaderBB->setTerminator( llr::CondBranch(
                    get< Value >( converted ),
                    pBodyBB, pSuccBB ) );
            }

            // If we have no successor BB, this will intetionally set the current BB to null.
            p.cfg()->setCurrentBB( pSuccBB );
            return true;
        };

        Rule r( handleWhile );
        auto ruleVal = ToValue( move( r ) );
        auto ruleTerm = ValueToIRExpr( ruleVal );
        e.storeValue( AppendToVectorTerm( RootIdentity(), TSID( while ) ), ANYTERM( _ ), ruleTerm );
    }
}
