#include "analyser.hpp"
#include "llvm-helpers.hpp"

#include <iostream>
#include <string>

DIVINE_RELAX_WARNINGS
#include <llvm/IR/Type.h>
#include <llvm/IR/DerivedTypes.h>
#include <llvm/IR/Constants.h>
#include <llvm/IR/BasicBlock.h>
#include <llvm/IR/IRBuilder.h>
#include <llvm/IR/Function.h>
#include <llvm/Transforms/Utils/Cloning.h>
DIVINE_UNRELAX_WARNINGS

#include <divine/mc/job.hpp>
#include <divine/mc/safety.hpp>

#include <divine/vm/divm.h>
#include <divine/vm/memory.hpp>

#include <divine/ui/log.hpp>

#include <divine/dbg/context.hpp>
#include <divine/dbg/stepper.hpp>

namespace shoop
{

using namespace std::string_literals;

void function_analyser::dump( const divine::mc::BitCode &bc, const std::string &name )
{
    const auto file = std::to_string( _dump_no++ ) + name;
    _owner->dump( bc, file );
}

formula lattice_climb_analyser::run()
{
    auto lattice = invocation_lattice{ _function->getFunctionType()->getNumParams() };

    std::cout << "\tClimbing lattice (" << lattice._bottom.to_string( "_", "_", "_", ", " ) << ")...\n";

    lattice.traverse( [ & ]( const invocation_lattice_node &node )
    {
        std::cout << "\tEntering node (" << node.to_string( "+", "-", "t", ", " ) << ")...\n";

        // First try with the least neccessary amount of terms.

        auto precondition = lattice.build_precondition( node );
        auto domains = node.promote( precondition );

        std::cout << "\t\t[Unit run]\n";
        std::cout << "\t\tPrecondition: " << precondition.to_string() << "\n";
        std::cout << "\t\tPromoted node: (" << domains.to_string( "+", "-", "t", ", " ) << ")\n";

        auto res = find_error( domains, precondition );

        if ( res.found_error )
        {
            // Have an error with units. Now swap units for terms, see if the error still
            // appears.

            domains = domains.promote();

            std::cout << "\t\t[Initial term run]\n";
            std::cout << "\t\tPrecondition: " << precondition.to_string() << "\n";
            std::cout << "\t\tPromoted node: (" << domains.to_string( "+", "-", "t", ", " ) << ")\n";

            res = find_error( domains, precondition );

            if ( res.found_error )
            {
                std::cout << "\t\tError path condition: " << res.path_condition.to_string() << "\n";
                lattice.add_precondition_conjunct( node, formula::negation( res.path_condition ) );

                while ( res.found_error )
                {
                    precondition = lattice.build_precondition( node );
                    res = find_error( domains, precondition );

                    std::cout << "\t\t[Ordinary term run]\n";
                    std::cout << "\t\tPrecondition: " << precondition.to_string() << "\n";

                    if ( res.found_error )
                    {
                        std::cout << "\t\tError path condition: " << res.path_condition.to_string() << "\n";
                        lattice.add_precondition_conjunct( node, formula::negation( res.path_condition ) );
                    }
                }
            }
            else
            {
                // Spurious error!
                // TODO: Do something about this
                std::cout << "\t\tSpurious error!\n";
            }
        }

        std::cout << "\t\tFinished.\n";

        return precondition;
    } );

    std::cout << "\tComputed precondition: " << lattice.top_precondition().to_string() << "\n";

    return lattice.top_precondition();
}

formula all_terms_analyser::run()
{
    auto top = invocation_lattice_node::top( _function->getFunctionType()->getNumParams() ).promote();

    std::cout << "\tEntering node (" << top.to_string( "+", "+", "+", ", " ) << ")...\n";

    auto precondition = formula::constant( true );

    std::cout << "\t\t[Initial term run]\n";
    std::cout << "\t\tPrecondition: " << precondition.to_string() << "\n";

    auto res = find_error( top, precondition );

    while ( res.found_error )
    {
        std::cout << "\t\tError path condition: " << res.path_condition.to_string() << "\n";

        const auto conj = std::vector< formula >{ precondition, formula::negation( res.path_condition ) };
        precondition = formula::conjunction( conj.cbegin(), conj.cend() );

        std::cout << "\t\t[Ordinary term run]\n";
        std::cout << "\t\tPrecondition: " << precondition.to_string() << "\n";

        res = find_error( top, precondition );
    }

    std::cout << "\tComputed precondition: " << precondition.to_string() << "\n";

    return precondition;
}

static llvm::Constant *create_formula_data( llvm::Module &mod, const expr_t &expr )
{
    auto &ctx = mod.getContext();
    const auto i8_type = get_int_type< uint8_t >( ctx );

    auto data = std::vector< llvm::Constant * >();

    for ( const auto v : bytes( expr ) )
        data.push_back( llvm::ConstantInt::get( i8_type, v ) );

    auto value_type = llvm::ArrayType::get( i8_type, data.size() );
    auto value = llvm::ConstantArray::get( value_type, data );

    auto variable = llvm::dyn_cast< llvm::GlobalVariable >( mod.getOrInsertGlobal( precondition_var_name, value_type ) );
    if ( !variable )
        die( "Could not insert precondition data." );

    variable->setInitializer( value );
    variable->setConstant( true );
    variable->setLinkage( llvm::GlobalValue::LinkageTypes::PrivateLinkage );
    variable->setUnnamedAddr( llvm::GlobalValue::UnnamedAddr::Global );

    return llvm::ConstantExpr::getBitCast( variable, i8_type->getPointerTo() );
}

// Create a __tester which calls the tested function with arguments of the specified domains.
// We also assume the precondition at the beginning.
static void generate_test_wrapper( llvm::Function *fun, const invocation_lattice_node &params, const formula &pre )
{
    auto &ctx = fun->getContext();
    auto mod  = fun->getParent();

    const auto type = llvm::FunctionType::get( llvm::Type::getVoidTy( ctx ), false );
    auto wrapper    = as_function( mod->getOrInsertFunction( test_runner_name, type ) );

    auto entry = llvm::BasicBlock::Create( ctx, std::string( test_runner_name ) + ".0", wrapper );
    llvm::IRBuilder<> builder{ entry };

    // Assuming the precondition holds...

    // We add assumptions using __set_precondition( ptr, size ), where 'ptr' points to representation of
    // the SMT expression in memory and 'size' is its size in bytes.

    const auto &expr = pre.representation();
    const auto formula_data = create_formula_data( *mod, expr );
    const auto formula_size = llvm::ConstantInt::get( get_int_type< int >( ctx ), bytes( expr ).size() );

    builder.CreateCall( get_set_precondition( *mod ), { formula_data, formula_size } );

    // ... call the tested function.

    auto call_args = std::vector< llvm::Value * >();
    auto parameters_it = params.parameters().begin();

    for ( auto &t : fun->getFunctionType()->params() )
    {
        auto sym = builder.CreateCall( get_abstract_constructor( *parameters_it++, mod, t ) );
        call_args.push_back( sym );
    }

    builder.CreateCall( fun, call_args );
    builder.CreateRetVoid();
}

static std::string mk_dump_name(
    const llvm::Function &f,
    const invocation_lattice_node &params,
    const std::string &prefix,
    const std::string &sep
)
{
    return prefix + sep + f.getName().str() + sep + params.to_string( "unit", "sunit", "term", sep ) + ".bc";
}

divine_result function_analyser::find_error(
    const invocation_lattice_node &param_domains,
    const formula &precondition
)
{
    const auto source_module = _function->getParent();
    auto tested_module = llvm::CloneModule( *source_module );

    auto tested_fun = tested_module->getFunction( _function->getName() );
    ASSERT( tested_fun ); // The function must exist.

    // Prepare the module for verification of our given function with the correct domains and preconditions.
    ensure_main_exists( *tested_module );
    generate_test_wrapper( tested_fun, param_domains, precondition );

    auto bc = std::make_shared< divine::mc::BitCode >( std::move( tested_module ), _owner->_context );
    dump( *bc, mk_dump_name( *_function, param_domains, "", "_" ) );

    // Now run Divine...

    auto opts = divine::mc::BCOptions();
    opts.dios_config = "static"; // See `dios/config/static.cpp`
    opts.lamp_config = "shoop";  // See `dios/lamp/shoop.cpp`
    opts.symbolic    = true;

    bc->set_options( opts );
    bc->solver( "stp" );
    bc->init();

    dump( *bc, mk_dump_name( *_function, param_domains, "post_lart", "_" ) );

    auto safety = divine::mc::make_job< divine::mc::Safety >( bc, divine::ss::passive_listen() );
    safety->start( 1 );
    safety->wait();

    if ( safety->result() == divine::mc::Result::Valid )
        return { false, formula::constant( true ) };
    if ( safety->result() == divine::mc::Result::BootError )
        die( "Could not boot DiOS" );
    if ( safety->result() == divine::mc::Result::None )
        die( "Divine returned result 'None'" );

    // Otherwise, we must have ended with error (exhaustiveness check).
    ASSERT( safety->result() == divine::mc::Result::Error );

    divine::dbg::Context< divine::vm::CowHeap > dbg {  bc->program(), bc->debug() };
    auto trace = safety->ce_trace();
    auto log = divine::ui::make_yaml( std::cerr, false );

    if ( _owner->_opts._divine_log )
    {
        log->info( "\n" );
        log->result( safety->result(), trace );
    }

    safety->dbg_fill( dbg );
    dbg.load( trace.final );
    dbg._lock = trace.steps.back();
    dbg._lock_mode = decltype( dbg )::LockBoth;

    divine::vm::setup::scheduler( dbg );
    using Stepper = divine::dbg::Stepper< decltype( dbg ) >;
    Stepper step;
    step._stop_on_error = true;
    step._stop_on_accept = true;
    step._ff_components = divine::dbg::component::kernel;
    step.run( dbg, Stepper::Quiet );

    if ( _owner->_opts._divine_log )
        log->backtrace( dbg, 10 ); // TODO: How detailed we want the log to be?

    return { true, get_path_condition( trace, precondition ) };
}

}