#include "builtins/builtins.h"

using namespace goose;
using namespace goose::ir;
using namespace goose::llr;

namespace goose::builtins
{
    void SetupLocalVarUnification( Env& e )
    {
        auto localVarPattern = GetValueType< LocalVar >( ANYTERM( _ ) );

        auto constRefTypePattern = ValueToIRExpr(
            Value( GetValueType< ReferenceType >(), TVEC( TSID( reference ), TSID( const ), ANYTERM( _ ) ) ) );

        auto mutRefTypePattern = ValueToIRExpr(
            Value( GetValueType< ReferenceType >(), TVEC( TSID( reference ), TSID( mut ), ANYTERM( _ ) ) ) );

        // LocalVar unification against a param:
        // Build a const ref value.
        e.unificationRuleSet()->addSymRule(

            ParamPat( ANYTERM( _ ) ),

            ValueToIRExpr( ValuePattern(
                ANYTERM( _ ),
                localVarPattern,
                ANYTERM( _ ) ) ),

            []( const Term& lhs, const Term& rhs, UnificationContext& c ) -> UniGen
            {
                auto lvval = *ValueFromIRExpr( rhs );
                auto locvar = FromValue< LocalVar >( lvval );
                if( !locvar )
                    co_return;

                ReferenceType rt( locvar->type(), ReferenceType::Behavior::Constant );
                auto ref = ValueToIRExpr( BuildComputedValue( ValueToIRExpr( ToValue( rt ) ),
                    Load( VarAddress( locvar->index() ), rt.type() ) )
                    .setLocationId( lvval.locationId() ) );

                // Unify the param with the ref
                co_yield Unify( lhs, ref, c );
            } );

        // LocalVar unification against another LocalVar: unify their types.
        e.unificationRuleSet()->addSymRule(

            ParamPat( localVarPattern ),

            ValueToIRExpr( ValuePattern(
                ANYTERM( _ ),
                localVarPattern,
                ANYTERM( _ ) ) ),

            []( const Term& lhs, const Term& rhs, UnificationContext& uc ) -> UniGen
            {
                auto lvarType = *FromValue< LocalVarType >( *ValueFromIRExpr( ValuePatternFromIRExpr( lhs )->type() ) );

                auto rhsVal = *ValuePatternFromIRExpr( rhs );
                auto rvarType = *FromValue< LocalVarType >( *ValueFromIRExpr( rhsVal.type() ) );

                for( auto&& [s, uc] : Unify( lvarType.type(), rvarType.type(), uc ) )
                {
                    co_yield { ValueToIRExpr( Value( ValueToIRExpr( ToValue( LocalVarType( s ) ) ),
                        rhsVal.val() ) ), uc };
                }
            } );
    }
}
