#include "builtins/builtins.h"

namespace goose::builtins
{
    Term BuildArgPatternFromTDecl( const Context& c, Term&& sig )
    {
        return ValueToEIR( ValuePattern( HOLE( "_"_sid ), move( sig ), HOLE( "_"_sid ) ) );
    }

    void SetupTDeclTypeChecking( Env& e )
    {
        auto tDeclPat = TDeclSigPattern( ANYTERM( _ ), ANYTERM( _ ), ANYTERM( _ ) );

        e.typeCheckingRuleSet()->addHalfUnificationRule( TCRINFOS, tDeclPat,
            []( const Term& lhs, TypeCheckingContext& c )
            {
                auto [name,rptDepth,sig] = DecomposeTDSig( lhs );
                return HalfUnify( sig, c );
            } );

        e.typeCheckingRuleSet()->addTypeCheckingRule( TCRINFOS, tDeclPat, ANYTERM( _ ),
            []( const Term& lhs, const Term& rhs, TypeCheckingContext tcc ) -> TCGen
            {
                auto [name,rptDepth,typeSig] = DecomposeTDSig( lhs );
                auto tdeclHole = HOLE( name, ""_sid, rptDepth );

                auto pat = ValueToEIR( Value( move( typeSig ), HOLE( "_"_sid ) ) );
                for( auto&& [s,tcc] : TypeCheck( pat, rhs, tcc ) )
                {
                    // We need to typecheck 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 savedRHSSubContext = tcc.RHSSubContext();
                    tcc.RHSSubContext() = tcc.LHSSubContext();

                    for( auto&& [s,tcc] : TypeCheck( s, tdeclHole, tcc ) )
                    {
                        tcc.RHSSubContext() = savedRHSSubContext;
                        co_yield { s, tcc };
                    }
                }
            } );

        e.typeCheckingRuleSet()->addUnificationRule( TCRINFOS, tDeclPat, ANYTERM( _ ),
            []( const Term& lhs, const Term& rhs, TypeCheckingContext tcc ) -> TCGen
            {
                auto [name,rptDepth,typeSig] = DecomposeTDSig( lhs );
                auto tdeclHole = HOLE( name, ""_sid, rptDepth );

                auto pat = ValueToEIR( Value( move( typeSig ), 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 savedRHSSubContext = tcc.RHSSubContext();
                    tcc.RHSSubContext() = tcc.LHSSubContext();

                    for( auto&& [s,tcc] : Unify( s, tdeclHole, tcc ) )
                    {
                        tcc.RHSSubContext() = savedRHSSubContext;
                        co_yield { s, tcc };
                    }
                }
            } );

        // tfunc tdecl param / tfunc arg
        auto tDeclTFuncPat = TDeclSigPattern( ANYTERM( _ ), ANYTERM( _ ), TFuncTypeSigPattern( ANYTERM( _ ), ANYTERM( _ ) ) );

        e.typeCheckingRuleSet()->addTypeCheckingRule( TCRINFOS,

            tDeclTFuncPat,

            ValueToEIR( ValuePattern(
                TSID( constant ),
                TFuncTypePattern(),
                ANYTERM( _ ) ) ),

        []( const Term& lhs, const Term& rhs, TypeCheckingContext tcc ) -> TCGen
        {
            auto [name,rptDepth,typeSig] = DecomposeTDSig( lhs );
            auto [_,tfSig] = DecomposeTFTSig( typeSig );

            auto callPat = BuildArgPatternFromTDecl( tcc.context(), move( typeSig ) );
            auto tdeclHole = HOLE( name, ""_sid, rptDepth );

            auto rhsVal = *EIRToValue( rhs );

            ConstrainedFunc cfunc( tfSig, 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 savedRHSSubContext = tcc.RHSSubContext();
            tcc.RHSSubContext().namespaceIndex = tcc.newNamespaceIndex();

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

            for( auto&& [s, tcc] : TypeCheck( callPat, rhs, tcc ) )
            {
                // Restore the namespace
                tcc.RHSSubContext() = savedRHSSubContext;
                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.RHSSubContext() = tcc.LHSSubContext();

                for( auto&& [s,tcc] : Unify( cFuncTerm, tdeclHole, tcc ) )
                {
                    tcc.RHSSubContext() = savedRHSSubContext;
                    co_yield { s, tcc };
                }
            }
        } );

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

            move( tDeclTFuncPat ),

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

        []( const Term& lhs, const Term& rhs, TypeCheckingContext tcc ) -> TCGen
        {
            auto [name,rptDepth,typeSig] = DecomposeTDSig( lhs );
            auto [_,tfSig] = DecomposeTFTSig( typeSig );

            auto callPat = BuildArgPatternFromTDecl( tcc.context(), move( typeSig ) );
            auto tdeclHole = HOLE( name, ""_sid, rptDepth );

            auto rhsVal = *EIRToValue( rhs );

            ConstrainedFunc cfunc( tfSig, 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 savedRHSSubContext = tcc.RHSSubContext();
            tcc.RHSSubContext().namespaceIndex = tcc.newNamespaceIndex();

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

            for( auto&& [s, tcc] : TypeCheck( callPat, rhs, tcc ) )
            {
                // Restore the namespace
                tcc.RHSSubContext() = savedRHSSubContext;
                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.RHSSubContext() = tcc.LHSSubContext();

                for( auto&& [s,tcc] : Unify( cFuncTerm, tdeclHole, tcc ) )
                {
                    tcc.RHSSubContext() = savedRHSSubContext;
                    co_yield { s, tcc };
                }
            }
        } );
    }
}
