package com.sunyard.chsm.pool;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.sunyard.chsm.auth.AppUser;
import com.sunyard.chsm.constant.SecurityConstant;
import com.sunyard.chsm.enums.DeviceTmkStatus;
import com.sunyard.chsm.mapper.CryptoServiceDeviceGroupMapper;
import com.sunyard.chsm.mapper.SpDeviceMapper;
import com.sunyard.chsm.model.dto.DeviceCheckRes;
import com.sunyard.chsm.model.entity.CryptoServiceDeviceGroup;
import com.sunyard.chsm.model.entity.Device;
import com.sunyard.chsm.sdf.adapter.SdfApiAdapter;
import com.sunyard.chsm.sdf.adapter.SdfApiAdapterFactory;
import com.sunyard.chsm.sdf.model.EccCipher;
import com.sunyard.chsm.service.TmkService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.stereotype.Component;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/**
 * @author liulu
 * @since 2024/12/10
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class DeviceManager implements InitializingBean {

    private static final Map<Long, AtomicInteger> ROUND_MAP = new HashMap<>();

    //    private Map<Long, List<DeviceContext>> deviceMap = new HashMap<>();
    private Map<Long, List<Long>> serviceDeviceMap;
    private final Map<Long, DeviceContext> usedDeviceMap = new ConcurrentHashMap<>();

    private boolean enableSoftDevice = false;
    private TMKContext softContext;

    private final TmkService tmkService;
    private final SpDeviceMapper spDeviceMapper;
    private final CryptoServiceDeviceGroupMapper cryptoServiceDeviceGroupMapper;


    public void returnContextToPool(TMKContext context) {
        Optional.ofNullable(context.getPool())
                .ifPresent(it -> {
                    context.setPool(null);
                    it.returnObject(context);
                });
    }

    public TMKContext chooseOne() {
        RequestAttributes attributes = RequestContextHolder.currentRequestAttributes();
        AppUser user = (AppUser) attributes.getAttribute(SecurityConstant.ATTRIBUTE_APP_USER, RequestAttributes.SCOPE_REQUEST);
        Assert.notNull(user, "登录用户不能为空");
        //noinspection unchecked
        List<Long> serviceIds = (List<Long>) attributes.getAttribute("used_service_ids", RequestAttributes.SCOPE_REQUEST);
        Assert.isTrue(!CollectionUtils.isEmpty(serviceIds), "应用: " + user.getName() + "没有可用服务");

        AtomicInteger atomicInteger = ROUND_MAP.computeIfAbsent(user.getAppId(), k -> new AtomicInteger(1));
        if (atomicInteger.get() > Integer.MAX_VALUE - 10000) {
            atomicInteger.set(1);
        }
        List<DeviceContext> contexts = serviceDeviceMap.entrySet().stream()
                .filter(it -> serviceIds.contains(it.getKey()))
                .map(Map.Entry::getValue)
                .flatMap(Collection::stream)
                .map(usedDeviceMap::get)
                .filter(Objects::nonNull)
                .collect(Collectors.toList());
        DeviceContext device = getNextDevice(contexts, atomicInteger.getAndIncrement());
        TMKContext tmkContext;
        if (device == null) {
            if (log.isDebugEnabled()) {
                HttpServletRequest request = ((ServletRequestAttributes) (RequestContextHolder.currentRequestAttributes())).getRequest();
                log.debug("request: {}, use soft device 1", request.getRequestURI());
            }
            tmkContext = getSoftContext();
        } else {
            if (log.isDebugEnabled()) {
                HttpServletRequest request = ((ServletRequestAttributes) (RequestContextHolder.currentRequestAttributes())).getRequest();
                log.debug("request: {}, choose device: {}", request.getRequestURI(), device.getIpPort());
            }
            try {
                GenericObjectPool<TMKContext> pool = device.getPool();
                tmkContext = pool.borrowObject(2000);
                tmkContext.setPool(pool);
            } catch (Exception e) {
                usedDeviceMap.remove(device.getId());
                if (log.isDebugEnabled()) {
                    HttpServletRequest request = ((ServletRequestAttributes) (RequestContextHolder.currentRequestAttributes())).getRequest();
                    log.debug("request: {}, use soft device 2", request.getRequestURI());
                }
                tmkContext = getSoftContext();
            }
        }
        Assert.notNull(tmkContext, "应用: " + user.getName() + "没有可用的密码设备");
        return tmkContext;
    }

    public synchronized TMKContext getSoftContext() {
        if (!enableSoftDevice) {
            return null;
        }
        if (Objects.nonNull(softContext)) {
            return softContext;
        }
        byte[] softEncTmk = tmkService.getSoftDeviceEncTmk();
        if (softEncTmk == null || softEncTmk.length == 0) {
            return null;
        }
        SdfApiAdapter bcAdapter = SdfApiAdapterFactory.getBcAdapter();
        String hk = bcAdapter.importKeyWithISKECC("", 1, EccCipher.fromBytes(softEncTmk));
        softContext = new TMKContext();
        softContext.setSdfApiAdapter(bcAdapter);
        softContext.setKeyHandle(hk);
        return softContext;
    }


    public static DeviceContext getNextDevice(List<DeviceContext> devices, int totalCalls) {
        if (devices == null || devices.isEmpty()) {
            return null;
        }

        int totalWeight = 0;
        for (DeviceContext device : devices) {
            totalWeight += device.getWeight();
        }
        int index = totalCalls % totalWeight;

        for (DeviceContext device : devices) {
            if (index < device.getWeight()) {
                return device;
            }
            index -= device.getWeight();
        }

        return null; // 理论上不会到这里
    }

    private void syncDevice() {
        log.debug(">>>>>>>>>>>>>>> start sync device <<<<<<<<<<<<<<<");
        enableSoftDevice = tmkService.isEnableSoftDevice();

        List<Device> devices = spDeviceMapper.selectList(
                new LambdaQueryWrapper<Device>()
                        .eq(Device::getTmkStatus, DeviceTmkStatus.finished)
                        .gt(Device::getGroupId, 0)
        );
        if (CollectionUtils.isEmpty(devices)) {
            log.info("no device for sync ...");
            usedDeviceMap.clear();
            return;
        }
        Map<Long, List<Long>> groupDeviceMap = devices.stream()
                .collect(Collectors.groupingBy(Device::getGroupId, Collectors.mapping(Device::getId, Collectors.toList())));
        List<CryptoServiceDeviceGroup> serviceDeviceGroups = cryptoServiceDeviceGroupMapper
                .selectList(new LambdaQueryWrapper<CryptoServiceDeviceGroup>()
                        .in(CryptoServiceDeviceGroup::getDeviceGroupId, groupDeviceMap.keySet()));
        if (CollectionUtils.isEmpty(serviceDeviceGroups)) {
            usedDeviceMap.clear();
            return;
        }
        serviceDeviceMap = serviceDeviceGroups.stream()
                .collect(Collectors.toMap(CryptoServiceDeviceGroup::getServiceId,
                        it -> groupDeviceMap.get(it.getDeviceGroupId())));

        for (Device device : devices) {
            DeviceCheckRes checkRes = tmkService.checkDevice(device);
            if (checkRes.isHasError()) {
                log.warn("设备:{}, 检测异常:{}", device.getServiceIp(), checkRes.getMessage());
                usedDeviceMap.remove(device.getId());
                continue;
            }
            DeviceContext exist = usedDeviceMap.get(device.getId());
            if (Objects.isNull(exist)) {
                usedDeviceMap.remove(device.getId());
                DeviceContext context = getDeviceContext(device, checkRes);
                if (context != null) {
                    usedDeviceMap.put(device.getId(), context);
                }
            } else {
                if (!Objects.equals(checkRes.getDeviceSerial(), exist.getDeviceSerial())
                        || !Objects.equals(checkRes.getPubKey(), exist.getPubKey())) {
                    DeviceContext context = getDeviceContext(device, checkRes);
                    if (context != null) {
                        usedDeviceMap.put(device.getId(), context);
                    }
                }
            }
        }
    }

    private static DeviceContext getDeviceContext(Device device, DeviceCheckRes checkRes) {
        try {
            DeviceContext context = new DeviceContext();
            context.setId(device.getId());
            context.setIp(device.getServiceIp());
            context.setPort(device.getServicePort());
            context.setModel(device.getManufacturerModel());
            context.setEncKeyIdx(device.getEncKeyIdx());
            context.setAccessCredentials(device.getAccessCredentials());
            context.setDeviceSerial(device.getDeviceSerial());
            context.setPubKey(device.getPubKey());
            context.setEncTmk(device.getEncTmk());
            context.setWeight(device.getWeight());

            context.setSdfApiAdapter(checkRes.getSdfApiAdapter());

            GenericObjectPoolConfig<TMKContext> config = new GenericObjectPoolConfig<>();
            config.setMaxWait(Duration.ofSeconds(5));
            config.setJmxEnabled(false);
            config.setMinIdle(2);
            config.setTimeBetweenEvictionRuns(Duration.ofMinutes(1));
            config.setTestWhileIdle(true);

            TMKContextFactory tenantTMKContextFactory = new TMKContextFactory(checkRes.getSdfApiAdapter(), context);
            GenericObjectPool<TMKContext> pool = new GenericObjectPool<>(tenantTMKContextFactory, config);
            context.setPool(pool);
            return context;
        } catch (Exception ex) {
            log.warn("device conn error: {}", device, ex);
            return null;
        }
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        Executors.newSingleThreadScheduledExecutor()
                .scheduleWithFixedDelay(this::syncDevice, 0L, 5L, TimeUnit.MINUTES);
    }
}
