package cn.novalon.manage.gateway.loadbalancer;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 自定义负载均衡器
 * 
 * 文件定义：实现多种负载均衡策略
 * 涉及业务：请求分发、服务实例选择、负载均衡策略
 * 
 * 负载均衡策略：
 * 1. 轮询
 * 2. 随机
 * 3. 加权轮询
 * 4. 最少连接
 * 
 * @author 张翔
 * @date 2026-03-26
 */
@Component
public class CustomLoadBalancer {

    private static final Logger logger = LoggerFactory.getLogger(CustomLoadBalancer.class);

    private final AtomicInteger position = new AtomicInteger(new Random().nextInt(1000));
    private final Map<String, AtomicInteger> connectionCounts = new ConcurrentHashMap<>();
    private final Map<String, Integer> weights = new ConcurrentHashMap<>();

    public ServiceInstance selectInstance(
            List<ServiceInstance> instances, 
            LoadBalanceStrategy strategy) {
        
        if (instances == null || instances.isEmpty()) {
            logger.warn("No instances available");
            return null;
        }

        ServiceInstance selectedInstance;

        switch (strategy) {
            case ROUND_ROBIN:
                selectedInstance = selectByRoundRobin(instances);
                break;
            case RANDOM:
                selectedInstance = selectByRandom(instances);
                break;
            case WEIGHTED_ROUND_ROBIN:
                selectedInstance = selectByWeightedRoundRobin(instances);
                break;
            case LEAST_CONNECTIONS:
                selectedInstance = selectByLeastConnections(instances);
                break;
            default:
                selectedInstance = selectByRoundRobin(instances);
        }

        if (selectedInstance != null) {
            logger.debug("Selected instance {}:{} using {} strategy", 
                    selectedInstance.getHost(), 
                    selectedInstance.getPort(), 
                    strategy);
        }

        return selectedInstance;
    }

    private ServiceInstance selectByRoundRobin(List<ServiceInstance> instances) {
        int pos = Math.abs(position.incrementAndGet());
        return instances.get(pos % instances.size());
    }

    private ServiceInstance selectByRandom(List<ServiceInstance> instances) {
        int index = new Random().nextInt(instances.size());
        return instances.get(index);
    }

    private ServiceInstance selectByWeightedRoundRobin(List<ServiceInstance> instances) {
        int totalWeight = instances.stream()
                .mapToInt(this::getWeight)
                .sum();

        if (totalWeight == 0) {
            return selectByRoundRobin(instances);
        }

        int randomWeight = new Random().nextInt(totalWeight);
        int currentWeight = 0;

        for (ServiceInstance instance : instances) {
            currentWeight += getWeight(instance);
            if (randomWeight < currentWeight) {
                return instance;
            }
        }

        return instances.get(0);
    }

    private ServiceInstance selectByLeastConnections(List<ServiceInstance> instances) {
        ServiceInstance selectedInstance = null;
        int minConnections = Integer.MAX_VALUE;

        for (ServiceInstance instance : instances) {
            int connections = getConnectionCount(instance);
            if (connections < minConnections) {
                minConnections = connections;
                selectedInstance = instance;
            }
        }

        return selectedInstance != null ? selectedInstance : instances.get(0);
    }

    private int getWeight(ServiceInstance instance) {
        String instanceKey = getInstanceKey(instance);
        return weights.getOrDefault(instanceKey, 1);
    }

    public void setWeight(ServiceInstance instance, int weight) {
        String instanceKey = getInstanceKey(instance);
        weights.put(instanceKey, weight);
        logger.debug("Set weight {} for instance {}", weight, instanceKey);
    }

    private int getConnectionCount(ServiceInstance instance) {
        String instanceKey = getInstanceKey(instance);
        AtomicInteger count = connectionCounts.get(instanceKey);
        return count != null ? count.get() : 0;
    }

    public void incrementConnection(ServiceInstance instance) {
        String instanceKey = getInstanceKey(instance);
        connectionCounts.computeIfAbsent(instanceKey, k -> new AtomicInteger(0)).incrementAndGet();
        logger.debug("Incremented connection count for instance {}", instanceKey);
    }

    public void decrementConnection(ServiceInstance instance) {
        String instanceKey = getInstanceKey(instance);
        AtomicInteger count = connectionCounts.get(instanceKey);
        if (count != null && count.get() > 0) {
            count.decrementAndGet();
            logger.debug("Decremented connection count for instance {}", instanceKey);
        }
    }

    private String getInstanceKey(ServiceInstance instance) {
        return instance.getHost() + ":" + instance.getPort();
    }

    public enum LoadBalanceStrategy {
        ROUND_ROBIN,
        RANDOM,
        WEIGHTED_ROUND_ROBIN,
        LEAST_CONNECTIONS
    }
}
