#include "builtins/builtins.h"

namespace goose::builtins
{
    Term BuildArgPatternFromTDecl( const TDecl& td )
    {
        return ValueToEIR( ValuePattern( HOLE( "_"_sid ), td.type(), HOLE( "_"_sid ) ) );
    }

    TCGen TypeCheckTDecl( const Term& lhs, const Term& rhs, TypeCheckingContext tcc )
    {
        auto tdecl = FromValue< TDecl >( *ValueFromEIR( lhs ) );
        assert( tdecl );

        auto tdeclHole = HOLE( tdecl->name() );

        auto pat = ValueToEIR( Value( tdecl->type(), HOLE( "_"_sid ) ) );

        for( auto&& [s,tcc] : Unify( pat, rhs, tcc ) )
        {
            // We need to unify the result with a hole named after the decl. However, since both sides of
            // this unification orignally appeared on the LHS, we need to setup RHS to alias the LHS namespace for this.
            auto savedRHSNamespaceIndex = tcc.RHSNamespaceIndex();
            tcc.setRHSNamespaceIndex( tcc.LHSNamespaceIndex() );

            for( auto&& [s,tcc] : TypeCheck( s, tdeclHole, tcc ) )
            {
                tcc.setRHSNamespaceIndex( savedRHSNamespaceIndex );
                co_yield { s, tcc };
            }
        }
    }

    void SetupTDeclTypeChecking( Env& e )
    {
        auto tDeclPat = ValueToEIR( Value( GetValueType< TDecl >(), VEC( ANYTERM( _ ), ANYTERM( _ ) ) ) );

        e.typeCheckingRuleSet()->addHalfUnificationRule( TCRINFOS, tDeclPat,
            []( const Term& lhs, TypeCheckingContext& c )
            {
                auto tdecl = FromValue< TDecl >( *ValueFromEIR( lhs ) );
                assert( tdecl );
                return HalfUnify( tdecl->type(), c );
            } );

        e.typeCheckingRuleSet()->addTypeCheckingRule( TCRINFOS, tDeclPat, ANYTERM( _ ), TypeCheckTDecl );
        e.typeCheckingRuleSet()->addTypeCheckingRule( TCRINFOS, tDeclPat, tDeclPat, TypeCheckTDecl );

        // tfunc tdecl param / tfunc arg
        auto tFuncTypePat = ValueToEIR( Value( TypeType(), VEC( TSID( texpr ), TSID( tfunc ),
            ANYTERM( _ ), ANYTERM( _ ), ANYTERM( _ ), ANYTERM( _ ) ) ) );

        auto tDeclTFuncPat = ParamPat( GetValueType< TDecl >(), VEC( tFuncTypePat, ANYTERM( _ ) ) );

        e.typeCheckingRuleSet()->addTypeCheckingRule( TCRINFOS,

            tDeclTFuncPat,

            ValueToEIR( ValuePattern(
                TSID( constant ),
                move( tFuncTypePat ),
                ANYTERM( _ ) ) ),

        []( const Term& lhs, const Term& rhs, TypeCheckingContext tcc ) -> TCGen
        {
            auto tdecl = FromValue< TDecl >( *ValueFromEIR( lhs ) );
            assert( tdecl );

            auto tfuncType = FromValue< TFuncType >( *ValueFromEIR( tdecl->type() ) );
            assert( tfuncType );

            auto callPat = BuildArgPatternFromTDecl( *tdecl );
            auto tdeclHole = HOLE( tdecl->name() );

            auto rhsVal = *ValueFromEIR( rhs );

            auto constraintPat = BuildTFuncSignature( tcc.context(), *tfuncType );
            assert( constraintPat );

            ConstrainedFunc cfunc( *constraintPat, GetTFuncInvocationRule(), rhsVal );
            auto cFuncTerm = ValueToEIR( ToValue( move( cfunc ) ) );

            // Create a new named hole namespace to isolate holes from the passed function from those in
            // the called function.
            auto savedRHSNamespaceIndex = tcc.RHSNamespaceIndex();
            tcc.setRHSNamespaceIndex( tcc.newNamespaceIndex() );

            auto oldValueRequired = tcc.isValueResolutionRequired();
            tcc.setValueResolutionRequired( false );

            for( auto&& [s, tcc] : TypeCheck( callPat, rhs, tcc ) )
            {
                // Restore the namespace
                tcc.setRHSNamespaceIndex( savedRHSNamespaceIndex );
                tcc.setValueResolutionRequired( oldValueRequired );

                // We need to unify the result with a hole named after the decl. However, since both sides of
                // this unification orignally appeared on the LHS, we need to setup RHS to alias the LHS namespace for this.
                tcc.setRHSNamespaceIndex( tcc.LHSNamespaceIndex() );

                for( auto&& [s,tcc] : Unify( cFuncTerm, tdeclHole, tcc ) )
                {
                    tcc.setRHSNamespaceIndex( savedRHSNamespaceIndex );
                    co_yield { s, tcc };
                }
            }
        } );

        // tfunc tdecl param / overloadset arg
        e.typeCheckingRuleSet()->addTypeCheckingRule( TCRINFOS,

            move( tDeclTFuncPat ),

            ValueToEIR( ValuePattern(
                TSID( constant ),
                GetValueType< ptr< builtins::OverloadSet > >(),
                ANYTERM( _ ) ) ),

        []( const Term& lhs, const Term& rhs, TypeCheckingContext tcc ) -> TCGen
        {
            auto tdecl = FromValue< TDecl >( *ValueFromEIR( lhs ) );
            assert( tdecl );

            auto tfuncType = FromValue< TFuncType >( *ValueFromEIR( tdecl->type() ) );
            assert( tfuncType );

            auto callPat = BuildArgPatternFromTDecl( *tdecl );
            auto tdeclHole = HOLE( tdecl->name() );

            auto rhsVal = *ValueFromEIR( rhs );

            auto constraintPat = BuildTFuncSignature( tcc.context(), *tfuncType );
            assert( constraintPat );

            ConstrainedFunc cfunc( *constraintPat, GetOverloadSetInvocationRule(), rhsVal );
            auto cFuncTerm = ValueToEIR( ToValue( move( cfunc ) ) );

            // Create a new named hole namespace to isolate holes from the passed function from those in
            // the called function.
            auto savedRHSNamespaceIndex = tcc.RHSNamespaceIndex();
            tcc.setRHSNamespaceIndex( tcc.newNamespaceIndex() );

            auto oldValueRequired = tcc.isValueResolutionRequired();
            tcc.setValueResolutionRequired( false );

            for( auto&& [s, tcc] : TypeCheck( callPat, rhs, tcc ) )
            {
                // Restore the namespace
                tcc.setRHSNamespaceIndex( savedRHSNamespaceIndex );
                tcc.setValueResolutionRequired( oldValueRequired );

                // We need to unify the result with a hole named after the decl. However, since both sides of
                // this unification orignally appeared on the LHS, we need to setup RHS to alias the LHS namespace for this.
                tcc.setRHSNamespaceIndex( tcc.LHSNamespaceIndex() );

                for( auto&& [s,tcc] : Unify( cFuncTerm, tdeclHole, tcc ) )
                {
                    tcc.setRHSNamespaceIndex( savedRHSNamespaceIndex );
                    co_yield { s, tcc };
                }
            }
        } );
    }
}
