#include "builtins/builtins.h"

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

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

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

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

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

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

        // Reference type checking rule.
        e.typeCheckingRuleSet()->addTypeCheckingRule( TCRINFOS,

            ParamPat( refTypePattern ),

            ValueToIRExpr( ValuePattern(
                ANYTERM( _ ),
                refTypePattern,
                ANYTERM( _ ) ) ),

            []( const Term& lhs, const Term& rhs, const TypeCheckingContext& tcc ) -> TCGen
            {
                auto lRefType = *FromValue< ReferenceType >( *ValueFromIRExpr( ValuePatternFromIRExpr( lhs )->type() ) );

                auto rhsVal = *ValueFromIRExpr( rhs );
                auto rRefType = *FromValue< ReferenceType >( *ValueFromIRExpr( rhsVal.type() ) );

                // Unify the behaviors
                for( auto&& [b, tcc] : Unify( lRefType.behavior(), rRefType.behavior(), tcc ) )
                {
                    // Unify the types
                    for( auto&& [t, tcc] : Unify( lRefType.type(), rRefType.type(), tcc ) )
                    {
                        co_yield { ValueToIRExpr( Value( ValueToIRExpr( ToValue( ReferenceType( t, b ) ) ),
                            rhsVal.llr() ) ), tcc };
                    }
                }
            }
        );

        // mut -> const reference unification rule.
        e.typeCheckingRuleSet()->addUnificationRule( TCRINFOS,
            TSID( const ),
            TSID( mut ),

            []( const Term& lhs, const Term& rhs, const TypeCheckingContext& tcc ) -> TCGen
            {
                co_yield { TSID( const ), tcc };
            }
        );

        // temp -> const reference unification rule.
        e.typeCheckingRuleSet()->addUnificationRule( TCRINFOS,
            TSID( const ),
            TSID( temp ),

            []( const Term& lhs, const Term& rhs, const TypeCheckingContext& tcc ) -> TCGen
            {
                co_yield { TSID( const ), tcc };
            }
        );

        // temp -> mut reference unification rule.
        e.typeCheckingRuleSet()->addUnificationRule( TCRINFOS,
            TSID( mut ),
            TSID( temp ),

            []( const Term& lhs, const Term& rhs, const TypeCheckingContext& tcc ) -> TCGen
            {
                co_yield { TSID( mut ), tcc };
            }
        );

        // Reference type checking against a param (implicit dereferencing):
        // Unify the referenced value with the param.
        e.typeCheckingRuleSet()->addTypeCheckingRule( TCRINFOS,

            ANYTERM( _ ),

            ValueToIRExpr( ValuePattern(
                ANYTERM( _ ),
                refTypePattern,
                ANYTERM( _ ) ) ),

            []( const Term& lhs, const Term& rhs, const TypeCheckingContext& tcc ) -> TCGen
            {
                auto refval = *ValueFromIRExpr( rhs );
                auto ref = FromValue< Reference >( refval );
                if( !ref )
                    co_return;

                auto content = ValueToIRExpr( BuildComputedValue( ref->type().type(),
                    llr::Load( ref->address(), ref->type().type() ) )
                    .setLocationId( refval.locationId() ) );

                // TypeCheck the param with the ref's content
                co_yield TypeCheck( lhs, content, tcc );
            } );

        // LocalVar type checking against a param (implicit referencing):
        // Build a mutable ref value.
        e.typeCheckingRuleSet()->addTypeCheckingRule( TCRINFOS,

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

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

        []( const Term& lhs, const Term& rhs, const TypeCheckingContext& tcc ) -> TCGen
        {
            auto ltype = ValuePatternFromIRExpr( lhs )->type();

            auto lvval = *ValueFromIRExpr( rhs );
            auto locvar = FromValue< LocalVar >( lvval );
            if( !locvar )
                co_return;

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

            co_yield TypeCheck( lhs, ref, tcc );
        } );

        // Implicit referencing of non-variables: build a tempref
        e.typeCheckingRuleSet()->addTypeCheckingRule( TCRINFOS,

            ValueToIRExpr( ValuePattern(
                ANYTERM( _ ),
                refTypePattern,
                ANYTERM( _ ) ) ),

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

        [refTypePattern]( const Term& lhs, const Term& rhs, const TypeCheckingContext& tcc ) -> TCGen
        {
            auto rval = *ValueFromIRExpr( rhs );

            if( !tcc.context().codeBuilder() )
                co_return;

            if( !tcc.context().codeBuilder()->cfg() )
                co_return;

            ReferenceType rt( rval.type(), TSID( temp ) );
            auto tempIndex = tcc.context().codeBuilder()->cfg()->getNewTemporaryIndex();
            auto ref = ValueToIRExpr( BuildComputedValue( ValueToIRExpr( ToValue( rt ) ),
                Load( TemporaryAddress( tempIndex, rval ), rt.type() ) )
                .setLocationId( rval.locationId() ) );

            // TypeCheck the param with the ref
            for( auto&& [s,tcc] : TypeCheck( lhs, ref, tcc ) )
            {
                assert( tcc.complexity() >= GetComplexity( ParamPat( refTypePattern ) ) );
                tcc.subComplexity( GetComplexity( ParamPat( refTypePattern ) ) );
                co_yield { s, tcc };
            }
        } );
    }
}
