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

using namespace goose::diagnostics;

namespace goose::verify
{
    optional< Z3Val > BuildZ3ValFromConstant( const sema::Context& c, const Value& val )
    {
        auto tinfo = TypeCache::GetInstance()->getTypeInfo( c, val.type() );
        if( !tinfo )
            return nullopt;

        auto zexpr = tinfo->ctor( val );
        if( !zexpr )
            return nullopt;

        return Z3Val { *zexpr, *ValueFromIRExpr( val.type() ) };
    }

    optional< Z3Val > BuildZ3ConstantFromType( const sema::Context& c, const Value& type, const string& name )
    {
        auto tinfo = TypeCache::GetInstance()->getTypeInfo( c, ValueToIRExpr( type ) );
        if( !tinfo )
            return nullopt;

        assert( tinfo->sort );
        return Z3Val { GetZ3Context().constant( name.c_str(), *tinfo->sort ), type };
    }

    optional< Z3Val > BuildZ3ConstantFromType( const sema::Context& c, const Term& type, const string& name )
    {
        auto tinfo = TypeCache::GetInstance()->getTypeInfo( c, type );
        if( !tinfo )
            return nullopt;

        assert( tinfo->sort );
        return Z3Val { GetZ3Context().constant( name.c_str(), *tinfo->sort ), *ValueFromIRExpr( type ) };
    }

    z3::expr GetAsBitVec( const z3::expr& expr, const Value& type )
    {
        if( expr.is_bv() )
            return expr;

        assert( expr.is_int() );

        auto intType = FromValue< builtins::IntegerType >( type );
        assert( intType );

        return z3::int2bv( intType->m_numBits, expr );
    }

    z3::expr GetAsBitVec( const Z3Val& zv )
    {
        return GetAsBitVec( zv.expr, zv.type );
    }

    z3::expr GetAsInt( const z3::expr& expr, const Value& type )
    {
        if( expr.is_int() )
            return expr;

        assert( expr.is_bv() );

        auto intType = FromValue< builtins::IntegerType >( type );
        assert( intType );

        return z3::bv2int( expr, intType->m_signed );
    }

    z3::expr GetAsInt( const Z3Val& zv )
    {
        return GetAsInt( zv.expr, zv.type );
    }

    template< typename I, typename F >
    optional< Z3Val > BuildZ3UnaryExpr( Builder& b, const I& instr, F&& func )
    {
        auto operand = BuildZ3ExprFromValue( b, instr.operand() );
        if( !operand )
            return nullopt;

        return Z3Val{ func( operand->expr ), operand->type };
    }

    template< typename I, typename F >
    optional< Z3Val > BuildZ3BinExpr( Builder& b, const I& instr, F&& func )
    {
        auto lhs = BuildZ3ExprFromValue( b, instr.lhs() );
        if( !lhs )
            return nullopt;

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

        if( lhs->expr.get_sort().sort_kind() == rhs->expr.get_sort().sort_kind() )
            return Z3Val{ func( lhs->expr, rhs->expr, lhs->type ), lhs->type };

        // If we are trying to do an operation on a mix of bitvec and int,
        // convert the int to a bitvec first.
        if( lhs->expr.is_bv() )
        {
            assert( rhs->expr.is_int() );
            return Z3Val{ func( lhs->expr, GetAsBitVec( *rhs ), lhs->type ), lhs->type };
        }
        else
        {
            assert( lhs->expr.is_int() );
            return Z3Val{ func( GetAsBitVec( *lhs ), rhs->expr, lhs->type ), lhs->type };
        }

        return nullopt;
    }

    template< typename T >
    optional< Z3Val > BuildZ3Op( Builder& b, const T& instr )
    {
        return nullopt;
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const Load& instr )
    {
        return LoadFromAddress( b, instr.addr(), instr.type() );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const Store& instr )
    {
        StoreToAddress( b, instr.addr(), instr.val() );
        return nullopt;
    }

    // Implemented in call.cpp
    extern optional< Z3Val > BuildZ3Op( Builder& b, const Call& instr );

    optional< Z3Val > BuildZ3Op( Builder& b, const CreateTemporary& instr )
    {
        b.setVar( instr.index(), instr.value() );
        return nullopt;
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const GetTemporary& instr )
    {
        const auto* expr = b.retrieveVar( instr.index(), instr.type() );
        if( expr )
            return Z3Val{ *expr, *ValueFromIRExpr( instr.type() ) };

        return BuildZ3ConstantFromType( b.context(), instr.type(), format( "v{}", instr.index() ) );
    }

    // Implemented in phi.cpp
    optional< Z3Val > BuildZ3Op( Builder& b, const Phi& instr );

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

    optional< Z3Val > BuildZ3Op( Builder& b, const Not& instr )
    {
        return BuildZ3UnaryExpr( b, instr, []( auto&& operand )
        {
            if( operand.is_bool() )
                return !operand;

            assert( operand.is_bv() );
            return ~operand;
        } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const And& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type )
        {
            if( lhs.is_bool() && rhs.is_bool() )
                return lhs && rhs;

            auto lhsBV = GetAsBitVec( lhs, type );
            auto rhsBV = GetAsBitVec( rhs, type );
            return lhsBV & rhsBV;
        } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const Or& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type )
        {
            if( lhs.is_bool() && rhs.is_bool() )
                return lhs || rhs;

            auto lhsBV = GetAsBitVec( lhs, type );
            auto rhsBV = GetAsBitVec( rhs, type );
            return lhsBV | rhsBV;
        } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const Xor& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type )
        {
            if( lhs.is_bool() && rhs.is_bool() )
                return mkZ3BoolXor( lhs, rhs );

            auto lhsBV = GetAsBitVec( lhs, type );
            auto rhsBV = GetAsBitVec( rhs, type );
            return lhsBV ^ rhsBV;
        } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const Shl& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type )
        {
            auto lhsBV = GetAsBitVec( lhs, type );
            auto rhsBV = GetAsBitVec( rhs, type );
            return z3::shl( lhsBV, rhsBV );
        } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const LShr& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type )
        {
            auto lhsBV = GetAsBitVec( lhs, type );
            auto rhsBV = GetAsBitVec( rhs, type );
            return z3::lshr( lhsBV, rhsBV );
        } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const AShr& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type )
        {
            auto lhsBV = GetAsBitVec( lhs, type );
            auto rhsBV = GetAsBitVec( rhs, type );
            return z3::ashr( lhsBV, rhsBV );
        } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const Add& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type ) { return lhs + rhs; } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const Sub& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type ) { return lhs - rhs; } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const Mul& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type ) { return lhs * rhs; } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const UDiv& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type ) { return z3::udiv( lhs, rhs ); } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const SDiv& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type ) { return lhs / rhs; } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const URem& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type ) { return z3::urem( lhs, rhs ); } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const SRem& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type )
        {
            if( lhs.is_bv() )
                return z3::srem( lhs, rhs );

            return z3::rem( lhs, rhs );
        } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const Eq& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type ) { return lhs == rhs; } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const Neq& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type ) { return lhs != rhs; } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const UGT& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type ) { return z3::ugt( lhs, rhs ); } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const UGE& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type ) { return z3::uge( lhs, rhs ); } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const ULT& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type ) { return z3::ult( lhs, rhs ); } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const ULE& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type ) { return z3::ule( lhs, rhs ); } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const SGT& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type ) { return lhs > rhs; } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const SGE& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type ) { return lhs >= rhs; } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const SLT& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type ) { return lhs < rhs; } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const SLE& instr )
    {
        return BuildZ3BinExpr( b, instr, []( auto&& lhs, auto&& rhs, auto&& type ) { return lhs <= rhs; } );
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const Assert& instr )
    {
        // TODO: refactor the return value here to be able to indicate failure versus "nothing to do"
        auto cond = BuildZ3ExprFromValue( b, instr.cond() );
        if( cond )
            b.checkAssertion( cond->expr, instr.cond().locationId() );
        return nullopt;
    }

    optional< Z3Val > BuildZ3Op( Builder& b, const Placeholder& instr )
    {
        const auto* expr = b.retrievePlaceholder( instr.name() );
        if( expr )
            return Z3Val{ *expr, *ValueFromIRExpr( instr.type() ) };

        return BuildZ3ConstantFromType( b.context(), instr.type(), format( "p{}", instr.name() ) );
    }

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

    optional< Z3Val > BuildZ3ExprFromValue( Builder& b, const Value& val )
    {
        if( val.isPoison() )
            return nullopt;

        if( val.isConstant() )
            return BuildZ3ValFromConstant( b.context(), val );

        if( auto expr = BuildZ3Op( b, *val.llr() )  )
            return expr;

        return BuildZ3ConstantFromType( b.context(), val.type(), format( "val{}", b.newUniqueId() ) );
    }
}
