#include "analyze.h"
#include "builtins/builtins.h"

namespace goose::analyze
{
    template< typename I, typename F >
    optional< z3::expr > BuildZ3BinExpr( const I& instr, F&& func )
    {
        auto lhs = BuildZ3ExprFromValue( instr.lhs() );
        if( !lhs )
            return nullopt;

        auto rhs = BuildZ3ExprFromValue( instr.rhs() );
        if( !rhs )
            return nullopt;

        return func( *lhs, *rhs );
    }

    // TODO: inline cfg. We'll need those functions to take some context with a z3 solver and an optional precondition (indicating which basic block we're in)
    // so that they can push assertions directly from the inline CFG (this will have to recurse into the code that builds assertions from a cfg, which doesn't even exist yet)

    // TODO: Call. Lots of work to do there: it needs to extract the function's assertions, map args to function param ids, and replace GetVars to thos ids with the corresponding
    // expressions.

    // TODO: Create temporary. Need to keep track of which z3 expressions are assigned to what temporaries, much like in codegen.

    // TODO: GetTemporary.

    // TODO: GetVar. Retrieve the latest z3 expression stored in that var, in the same way that we track allocas in codegen.
    // TODO: SetVar. Store a new z3 expression into that var.

    // TODO: Phi. Construct an expression combining all of the incoming z3 expressions from all the possible incoming blocks.

    // TODO: LoadConstStr. Build a z3 str value.

    // TODO: bitwise operators. They need to work on bitvecs, so we'll need a param indicating that we want to
    // build bitvecs rather than ints, and use it when building the sub expression of something that needs bitvecs.

    template< typename T >
    optional< z3::expr > BuildZ3Op( const T& instr )
    {
        return nullopt;
    }

    // TODO: Not. We don't have a dedicated opcode in llr for it bc llvm didn't need it,
    // but we'll need it so add it (and lower it to xor with complement for llvm in codegen)

    optional< z3::expr > BuildZ3Op( const And& instr )
    {
        // TODO this works only for bools, have to use & for bitvecs
        return BuildZ3BinExpr( instr, []( auto&& lhs, auto&& rhs ) { return lhs && rhs; } );
    }

    optional< z3::expr > BuildZ3Op( const Or& instr )
    {
        // TODO this works only for bools, have to use & for bitvecs
        return BuildZ3BinExpr( instr, []( auto&& lhs, auto&& rhs ) { return lhs || rhs; } );
    }

    // TODO: Xor, Shl, LShr, AShr

    optional< z3::expr > BuildZ3Op( const Add& instr )
    {
        return BuildZ3BinExpr( instr, []( auto&& lhs, auto&& rhs ) { return lhs + rhs; } );
    }

    optional< z3::expr > BuildZ3Op( const Sub& instr )
    {
        return BuildZ3BinExpr( instr, []( auto&& lhs, auto&& rhs ) { return lhs - rhs; } );
    }

    optional< z3::expr > BuildZ3Op( const Mul& instr )
    {
        return BuildZ3BinExpr( instr, []( auto&& lhs, auto&& rhs ) { return lhs * rhs; } );
    }

    // TODO: UDiv (check how unsigneds need to be handled. I've seen somewhere saying to use bitvectors,
    // but the api have unsigned methods for nums? Maybe in current versions of z3 we don't need to use bitvectors for unsigned?)

    optional< z3::expr > BuildZ3Op( const SDiv& instr )
    {
        return BuildZ3BinExpr( instr, []( auto&& lhs, auto&& rhs ) { return lhs / rhs; } );
    }

    //  TODO: URem

    optional< z3::expr > BuildZ3Op( const SRem& instr )
    {
        // TODO: need to use a different function for bitvectors
        return BuildZ3BinExpr( instr, []( auto&& lhs, auto&& rhs ) { return z3::srem( lhs, rhs ); } );
    }

    optional< z3::expr > BuildZ3Op( const Eq& instr )
    {
        return BuildZ3BinExpr( instr, []( auto&& lhs, auto&& rhs ) { return lhs == rhs; } );
    }

    optional< z3::expr > BuildZ3Op( const Neq& instr )
    {
        return BuildZ3BinExpr( instr, []( auto&& lhs, auto&& rhs ) { return lhs != rhs; } );
    }

    // TODO: UGT, UGE, ULT, ULE

    optional< z3::expr > BuildZ3Op( const SGT& instr )
    {
        return BuildZ3BinExpr( instr, []( auto&& lhs, auto&& rhs ) { return lhs > rhs; } );
    }

    optional< z3::expr > BuildZ3Op( const SGE& instr )
    {
        return BuildZ3BinExpr( instr, []( auto&& lhs, auto&& rhs ) { return lhs >= rhs; } );
    }

    optional< z3::expr > BuildZ3Op( const SLT& instr )
    {
        return BuildZ3BinExpr( instr, []( auto&& lhs, auto&& rhs ) { return lhs < rhs; } );
    }

    optional< z3::expr > BuildZ3Op( const SLE& instr )
    {
        return BuildZ3BinExpr( instr, []( auto&& lhs, auto&& rhs ) { return lhs <= rhs; } );
    }

    optional< z3::expr > BuildZ3Op( const llr::Instruction& instr )
    {
        return visit( [&]( auto&& e )
        {
            return BuildZ3Op( e );
        }, instr.content() );
    }

    optional< z3::expr > BuildZ3ExprFromConstant( const Value& val )
    {
        auto& c = GetZ3Context();

        // TODO: what about very large int types? Can z3 represent them?

        if( auto b = FromValue< bool >( val ) )
            return c.bool_val( *b );
        else if( auto intVal = FromValue< APSInt >( val ) )
             return c.int_val( intVal->getExtValue() );
        else if( auto strVal = FromValue< string >( val ) )
            return c.string_val( *strVal );

        return nullopt;
    }

    optional< z3::expr > BuildZ3ExprFromValue( const Value& val )
    {
        if( val.isPoison() )
            return nullopt;

        if( val.isConstant() )
            return BuildZ3ExprFromConstant( val );

        return BuildZ3Op( *val.llr() );
    }
}
