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

#include <brick-fs>
#include <brick-llvm>

// TODO: Delete unneeded (or move to function_analyser.cpp)
#include <algorithm>
#include <iterator>
#include <iostream> // TODO: Better, more configurable logging with verbosity settings
#include <sstream>  // TODO: Bricks stringbuilder instead? (see also lattice_climb.cpp)
#include <string>
#include <vector>
#include <queue>
#include <set>
#include <unordered_map>
#include <functional>
#include <optional>

#include <divine/smt/rpn.hpp>
#include <divine/smt/builder-llvm.hpp>

DIVINE_RELAX_WARNINGS
#include <llvm/IR/BasicBlock.h>
#include <llvm/IR/IRBuilder.h>
#include <llvm/IRReader/IRReader.h>
#include <llvm/Support/SourceMgr.h>
#include <llvm/Support/raw_ostream.h>
#include <llvm/Analysis/CallGraph.h>
DIVINE_UNRELAX_WARNINGS

namespace shoop
{

using namespace std::string_literals;

// TODO: We also want to support loading a source file and compiling it to LLVM bitcode here
std::unique_ptr< module_analyser > make_analyser( const char *path, analyser_opts opts )
{
    ASSERT( path );

    auto context = std::make_shared< llvm::LLVMContext >();

    llvm::SMDiagnostic diag;
    auto mod = llvm::parseIRFile( path, diag, *context );

    if ( !mod )
    {
        auto error_string = std::string();
        auto stream = llvm::raw_string_ostream( error_string );

        diag.print( path, stream );

        // TODO: Is the error message from SMDiagnostic sufficient, or do we want to prepend/append
        // something to it?
        die( stream.str() );
    }

    mod->setTargetTriple( "x86_64-unknown-none-elf" );

    if ( opts._traversal_strategy == module_strategy::topsort_traversal )
        return std::make_unique< topsort_analyser >( context, std::move( mod ), std::move( opts ) );
    if ( opts._traversal_strategy == module_strategy::outer_only )
        return std::make_unique< outer_only_analyser >( context, std::move( mod ), std::move( opts ) );

    UNREACHABLE( "make (module) analyser: match not exhaustive" );
}

void module_analyser::run()
{
    std::cout << "Bitcode loaded\n"
                 "Loading functions... ";

    load_function_queue();

    std::cout << "ok\n"
                 "Checking function signatures... ";

    check_unsupported_functions();

    std::cout << "ok\n";

    std::cout << "Functions will be analysed in the following order: ";
    for ( const auto &fun : _functions )
        std::cout << fun->getName().str() << " ";
    std::cout << "\n";

    std::cout << "Starting analysis loop...\n";

    analyse_queue();

    std::cout << "SHOOP finished.\n";
}

void topsort_analyser::load_function_queue()
{
    ASSERT_EQ( _functions.size(), 0 );

    auto graph = llvm::CallGraph( *_module );

    // Function -> finish time (or nullopt if not yet closed)
    auto visited = std::unordered_map< llvm::Function *, std::optional< int > >{};
    int counter  = 0;

    std::function< void( llvm::CallGraphNode * ) > visit = [ & ]( llvm::CallGraphNode *node )
    {
        // The node representing a call out of the callgraph has a null function.
        if ( node == graph.getCallsExternalNode() )
            return;

        const auto u = node->getFunction();

        visited[ u ] = std::nullopt;

        for ( const auto [ inst, succ ] : *node )
        {
            const auto v = succ->getFunction();

            // Check if we have a back edge.
            if ( visited.count( v ) && !visited[ v ].has_value() )
            {
                die( std::string{ "Found recursion in the given module." }
                        + "SHOOP cannot handle recursion at the moment!\n"
                        + "Offending functions: " + u->getName().str() + ", " + v->getName().str() );
            }

            if ( !visited.count( v ) )
                visit( succ );
        }

        visited[ u ] = counter++;

        if ( !u->isDeclaration() )
            _functions.push_back( u );
    };

    for ( const auto [ inst, node ] : *graph.getExternalCallingNode() )
        if ( !visited.count( node->getFunction() ) )
            visit( node );
}

void outer_only_analyser::load_function_queue()
{
    ASSERT_EQ( _functions.size(), 0 );

    auto graph = llvm::CallGraph( *_module );

    for ( const auto& [ inst, node ] : graph )
    {
        const auto f = node->getFunction();

        if ( node.get() != graph.getCallsExternalNode()
          && node.get() != graph.getExternalCallingNode()
          && !f->isDeclaration()
          && node->getNumReferences() == 1 ) // Referenced only by the ExternalCallingNode
            _functions.push_back( f );
    }
}

void module_analyser::check_unsupported_functions() const
{
    const auto is_scalar_type = []( const llvm::Type *t )
    {
        return t->isIntegerTy( 8 ) || t->isIntegerTy( 16 ) || t->isIntegerTy( 32 )
            || t->isIntegerTy( 64 ) || t->isFloatTy() || t->isDoubleTy() || t->isVoidTy();
    };

    const auto is_offending = [ & ]( const llvm::Function *fun )
    {
        const auto ft = fun->getFunctionType();

        if ( !is_scalar_type( ft->getReturnType() ) || ft->isVarArg() || ft->params().size() > 64 )
            return true;

        for ( const auto t : ft->params() )
        {
            if ( !is_scalar_type( t ) )
                return true;
        }

        return false;
    };

    auto offending = std::vector< const llvm::Function * >();
    std::copy_if( _functions.begin(), _functions.end(), std::back_inserter( offending ), is_offending );

    if ( !offending.empty() )
    {
        auto ss = std::stringstream();

        for ( const auto *fun : offending )
        {
            ss << "'" << fun->getName().str() << "' has too complex signature\n";
        }

        ss << "\nNote: at the moment, only scalar types "
              "(i.e. types representable as an integral or floating point value, and void)"
              " are supported\n";

        die( ss.str() );
    }
}

void module_analyser::analyse_queue()
{
    while ( !_functions.empty() )
    {
        auto current = _functions.front();

        std::cout << "Analysing '" << current->getName().str() << "'... \n" << std::flush;

        const auto analyser = make_analyser( *current );
        const auto pre = analyser->run();

        if ( _opts._insert_preconditions )
        {
            std::cout << "Inserting precondition... " << std::flush;

            if ( pre != formula::constant( true ) )
                insert_precondition_assert( current, pre );

            std::cout << "ok\n";
        }

        _functions.pop_front();
    }
}

std::unique_ptr< function_analyser > module_analyser::make_analyser( const llvm::Function &fun ) const
{
    if ( _opts._analyser_strategy == function_strategy::lattice_climb )
        return std::make_unique< lattice_climb_analyser >( *this, fun );
    if ( _opts._analyser_strategy == function_strategy::terms_only )
        return std::make_unique< all_terms_analyser >( *this, fun );

    UNREACHABLE( "make (function) analyser: match not exhaustive" );
}

void module_analyser::dump( const divine::mc::BitCode &bc, const std::string &file_name ) const
{
    if ( _opts._dump_dir )
    {
        const auto result_path = brq::join_path( *_opts._dump_dir, file_name );
        brq::llvm_write( bc._module.get(), result_path );
    }
}

void module_analyser::insert_precondition_assert( llvm::Function *fun, const formula &pre )
{
    auto args = std::vector< llvm::Value * >{};

    for ( auto it = fun->arg_begin(); it != fun->arg_end(); ++it )
        args.push_back( &*it );

    auto bb_builder = divine::smt::builder::LLVM( *_context, args );

    const auto asserted_expr = divine::smt::evaluate( bb_builder, pre.representation() );
    auto *block = &bb_builder.block(); // TODO: This should probably return a pointer!
    block->setName( "entry" );

    // Insert the generated block at the beginning, assert on asserted_expr, them jump to the original
    // first block of the function.

    auto fun_first_bb = &*fun->begin();
    fun_first_bb->setName( "original_entry" );

    block->insertInto( fun, fun_first_bb );

    auto irb = llvm::IRBuilder<>{ block };
    irb.CreateCall( get_assert_precondition( *_module ), std::vector< llvm::Value * >{ asserted_expr } );
    irb.CreateBr( fun_first_bb );
}

}