package cn.reghao.tnb.sb.tomcat;

import cn.reghao.tnb.sb.tomcat.model.ClusterProperties;
import cn.reghao.tnb.sb.tomcat.model.MemberConfig;
import cn.reghao.tnb.sb.tomcat.model.ReceiverConfig;
import lombok.Data;
import org.apache.catalina.Engine;
import org.apache.catalina.ha.session.DeltaManager;
import org.apache.catalina.ha.tcp.SimpleTcpCluster;
import org.apache.catalina.tribes.group.GroupChannel;
import org.apache.catalina.tribes.group.interceptors.MessageDispatchInterceptor;
import org.apache.catalina.tribes.group.interceptors.StaticMembershipInterceptor;
import org.apache.catalina.tribes.group.interceptors.TcpFailureDetector;
import org.apache.catalina.tribes.membership.StaticMember;
import org.apache.catalina.tribes.transport.nio.NioReceiver;
import org.springframework.boot.web.embedded.tomcat.TomcatServletWebServerFactory;
import org.springframework.boot.web.server.WebServerFactoryCustomizer;
import org.springframework.context.annotation.Configuration;

/**
 * @author reghao
 * @date 2024-09-04 19:44:21
 */
@Data
//@Configuration
public class TomcatClusterContextCustomizer implements WebServerFactoryCustomizer<TomcatServletWebServerFactory> {
    private ClusterProperties clusterProperties;

    public TomcatClusterContextCustomizer(ClusterProperties clusterProperties) {
        this.clusterProperties = clusterProperties;
    }

    @Override
    public void customize(TomcatServletWebServerFactory factory) {
        if (clusterProperties.getEnable()) {
            factory.addContextCustomizers(context -> {
                configureCluster((Engine) context.getParent().getParent());
                context.setManager(new DeltaManager());
                context.setDistributable(true);
            });
        }
    }

    private void configureCluster(Engine engine) {
        ReceiverConfig receiver = clusterProperties.getReceiver();
        SimpleTcpCluster cluster = new SimpleTcpCluster();
        GroupChannel groupChannel = new GroupChannel();
        cluster.setChannel(groupChannel);
        NioReceiver nioReceiver = new NioReceiver();
        nioReceiver.setPort(receiver.getPort());
        nioReceiver.setAddress(receiver.getAddress());
        groupChannel.setChannelReceiver(nioReceiver);

        StaticMembershipInterceptor staticMembershipInterceptor = getStaticMembershipService();
        groupChannel.addInterceptor(new TcpFailureDetector());
        groupChannel.addInterceptor(staticMembershipInterceptor);
        groupChannel.addInterceptor(new MessageDispatchInterceptor());
        engine.setCluster(cluster);
    }

    private StaticMembershipInterceptor getStaticMembershipService() {
        MemberConfig[] members = clusterProperties.getMember();
        StaticMembershipInterceptor staticMembershipInterceptor = new StaticMembershipInterceptor();
        for (MemberConfig memberConfig : members) {
            StaticMember staticMember = new StaticMember();
            staticMember.setUniqueId(memberConfig.getId());
            staticMember.setHost(memberConfig.getHost());
            staticMember.setPort(memberConfig.getPort());
            staticMember.setLocal(memberConfig.getLocal());
            staticMembershipInterceptor.addStaticMember(staticMember);
        }
        return staticMembershipInterceptor;
    }
}
