package com.lanyuanxiaoyao.service.launcher.synchronizer.service;

import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.lanyuanxiaoyao.service.common.Constants;
import com.lanyuanxiaoyao.service.common.entity.RunMeta;
import com.lanyuanxiaoyao.service.common.utils.NameHelper;
import com.lanyuanxiaoyao.service.launcher.ExecutorService;
import com.lanyuanxiaoyao.service.launcher.configuration.ClusterConfiguration;
import com.lanyuanxiaoyao.service.launcher.configuration.HadoopConfiguration;
import com.lanyuanxiaoyao.service.launcher.configuration.ZookeeperConfiguration;
import dev.failsafe.Failsafe;
import dev.failsafe.RetryPolicy;
import java.io.IOException;
import java.time.Duration;
import java.util.concurrent.TimeUnit;
import javax.annotation.PreDestroy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.locks.InterProcessLock;
import org.apache.curator.framework.recipes.locks.InterProcessMutex;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.yarn.api.records.ApplicationId;
import org.apache.hadoop.yarn.client.api.YarnClient;
import org.apache.zookeeper.data.Stat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.converter.json.Jackson2ObjectMapperBuilder;
import org.springframework.stereotype.Service;

/**
 * Sync 服务
 *
 * @author ZhangJiacheng
 * @date 2023-12-07
 */
@Service("sync-service")
public class SynchronizerService {
    private static final Logger logger = LoggerFactory.getLogger(SynchronizerService.class);
    private static final RetryPolicy<ApplicationId> RETRY_POLICY = RetryPolicy.<ApplicationId>builder()
            .handle(Exception.class)
            .withDelay(Duration.ofSeconds(1))
            .withMaxAttempts(3)
            .build();

    private final HadoopConfiguration hadoopConfiguration;
    private final ClusterConfiguration clusterConfiguration;
    private final CuratorFramework zookeeperClient;
    private final ExecutorService executorService;
    private final YarnClient yarnClient;
    private final ObjectMapper mapper;

    public SynchronizerService(HadoopConfiguration hadoopConfiguration, ClusterConfiguration clusterConfiguration, ExecutorService executorService, ZookeeperConfiguration zookeeperConfiguration, Jackson2ObjectMapperBuilder builder) {
        this.hadoopConfiguration = hadoopConfiguration;
        this.clusterConfiguration = clusterConfiguration;
        this.executorService = executorService;

        mapper = builder.build();

        yarnClient = YarnClient.createYarnClient();
        yarnClient.init(new Configuration());
        yarnClient.start();

        zookeeperClient = CuratorFrameworkFactory.builder()
                .connectString(zookeeperConfiguration.getConnectUrl())
                .retryPolicy(new ExponentialBackoffRetry((int) Constants.SECOND, 3))
                .sessionTimeoutMs((int) Constants.SECOND)
                .connectionTimeoutMs((int) (30 * Constants.SECOND))
                .build();
        zookeeperClient.start();
    }

    @PreDestroy
    public void destroy() throws IOException {
        if (ObjectUtil.isNotNull(yarnClient)) {
            yarnClient.stop();
            yarnClient.close();
        }

        if (ObjectUtil.isNotNull(this.zookeeperClient)) {
            this.zookeeperClient.close();
        }
    }

    public String start(Long flinkJobId) throws Exception {
        String lockPath = NameHelper.syncLauncherLockPath(flinkJobId);
        InterProcessLock lock = new InterProcessMutex(zookeeperClient, lockPath);
        try {
            if (lock.acquire(2, TimeUnit.SECONDS)) {
                Stat stat = zookeeperClient.checkExists().forPath(NameHelper.syncRunningLockPath(flinkJobId, null));
                if (ObjectUtil.isNotNull(stat)) {
                    String message = StrUtil.format("Job {} is running", flinkJobId);
                    logger.warn(message);
                    throw new RuntimeException(message);
                }
                ApplicationId applicationId = Failsafe
                        .with(RETRY_POLICY)
                        .get(() -> executorService.runSync(
                                flinkJobId,
                                hadoopConfiguration.getKerberosKeytabPath(),
                                hadoopConfiguration.getKerberosPrincipal(),
                                clusterConfiguration.getCluster()
                        ));
                return applicationId.toString();
            } else {
                String message = StrUtil.format("Un acquire lock for {}", flinkJobId);
                logger.warn(message);
                throw new RuntimeException(message);
            }
        } finally {
            if (lock.isAcquiredInThisProcess()) {
                lock.release();
            }
        }
    }

    public void stop(Long flinkJobId) throws Exception {
        String syncLockPath = NameHelper.syncRunningLockPath(flinkJobId, null);
        Stat stat = zookeeperClient.checkExists().forPath(syncLockPath);
        if (ObjectUtil.isNull(stat)) {
            return;
        }
        byte[] bytes = zookeeperClient.getData().forPath(syncLockPath);
        RunMeta meta = mapper.readValue(new String(bytes), RunMeta.class);
        logger.info("Run meta: {}", meta);
        String applicationIdText = meta.getApplicationId();
        if (StrUtil.isBlank(applicationIdText)) {
            throw new RuntimeException(StrUtil.format("Cannot parse application id from zookeeper: {}", syncLockPath));
        }
        ApplicationId applicationId = ApplicationId.fromString(applicationIdText);
        Failsafe.with(RETRY_POLICY)
                .run(() -> yarnClient.killApplication(applicationId));
        logger.info("{} have be killed", applicationId);
    }
}
