#include "verify.h"
#include "builtins/builtins.h"

using namespace goose::diagnostics;

namespace goose::verify
{
    bool Func::checkLoop( const llr::BasicBlock& bb, queue< const BasicBlock* >& parentWorkQueue )
    {
        if( ms_TraceMode )
            cout << "  == Checking loop " << bb.index() << endl;

        m_remapper.beginLoopUnrolling( bb.index() );

        bool completed = false;

        uint32_t k = 0;

        // Max K is hardcoded for now. Maybe we'll allow it to be changed per function.
        // Or maybe we won't limit K but rather the number of expressions added to the solver
        // (so that we have a limit that works consistently regardless of loop nesting depth)
        uint32_t maxK = 16;

        while( !completed && k < maxK )
        {
            ++k;

            // Unroll the base case just once for each increasing K: we save the state of the z3 solver
            // before generating the induction case, so we can incrementally add unrolling of the base
            // case one at a time to save some work.
            m_remapper.nextLoopIteration();
            if( !buildZ3Expressions( bb, &parentWorkQueue ) )
            {
                // If a counter example was found while emitting the base case,
                // we know the verification failed and we can stop here.
                completed = true;
                break;
            }

            // Push the solver state before unrolling the induction case,
            // so that we can rollback all of this for the next attempt, if any.
            m_solver.push();

            unordered_set< uint32_t > convertedBBIndices = m_convertedBBIndices;

            // Likewise for the remapper: we don't want to keep the edge infos from the induction case unrolling,
            // plus we need to keep the state to do the next unrolling of the base case correctly.
            RemapperLocalStateGuard rsg( m_remapper );

            // Set a different assertion handler for the induction case: we want to just emit them as asumptions.
            BuilderStateGuard bsg( m_builder, [&]( const z3::expr& expr, const z3::expr& exprToCheck, uint32_t locationId )
            {
                m_builder.assume( expr );
                return true;
            } );

            // Havoc all the variables modified during the current loop and emit
            // the first unrolling for the induction case.
            m_remapper.nextLoopIteration();
            uint32_t firstInductionIterationHeaderIndex = m_remapper.remapBBId( bb );

            m_cfg->forEachVariableModifiedInLoop( bb.index(), [&]( auto&& type, uint32_t varId )
            {
                m_builder.havocVar( firstInductionIterationHeaderIndex, type, varId );
            } );

            buildZ3Expressions( bb, &parentWorkQueue );

            // Emit the rest of the induction case.
            uint32_t i = 1;
            while( i < k )
            {
                m_remapper.nextLoopIteration();
                buildZ3Expressions( bb, &parentWorkQueue );
                ++i;
            }

            // For the final loop instance, we want to resolve the assertions, but without emitting diagnostics:
            // if they fail it means we need to increase k and try again.
            BuilderStateGuard bsg2( m_builder, [&]( const z3::expr& expr, const z3::expr& exprToCheck, uint32_t locationId )
            {
                if( ms_TraceMode )
                    return true;

                z3::expr_vector exprVec( GetZ3Context() );
                exprVec.push_back( exprToCheck );
                bool result = m_solver.check( exprVec ) == z3::check_result::unsat;

                if( ( !result && ms_DumpSolverOnFailure ) || ( result && ms_DumpSolverOnSuccess ) )
                    cout << "Solver dump:\n" << m_solver << "check_unsat_ind: " << exprToCheck << endl << endl;

                return result;
            } );

            m_remapper.nextLoopIteration();
            if( buildZ3Expressions( bb, &parentWorkQueue ) )
            {
                completed = true;
                break;
            }

            m_convertedBBIndices = move( convertedBBIndices );
            m_solver.pop();
        }

        if( !completed )
        {
            // If completed is false, the verification couldn't be completed with the maxK value used.
            // The whole business is undecidable so obviously there will be cases where we can't tell
            // whether the program is good.
            //
            // Perhaps instead of this we should display the assertion that failes at max K as an indication
            // of what might be wrong.
            //
            // TODO: we also need a location for that... Maybe let the loop parser set a proper location id in the loop header?
            // (like at the loop key word + the condition?)
            DiagnosticsManager::GetInstance().emitErrorMessage( 0,
                "couldn't verify this loop." );
        }

        m_remapper.endLoopUnrolling();

        if( ms_TraceMode )
            cout << "  == Finished checking loop " << bb.index() << endl;
        return completed && !m_builder.hasCheckFailed();
    }
}
