package com.lanyuanxiaoyao.service.executor.manager.service;

import cn.hutool.core.io.IoUtil;
import cn.hutool.core.map.MapBuilder;
import cn.hutool.core.map.MapUtil;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.lanyuanxiaoyao.service.executor.Runner;
import com.lanyuanxiaoyao.service.executor.core.TaskConstants;
import com.lanyuanxiaoyao.service.executor.core.TaskContext;
import com.lanyuanxiaoyao.service.executor.manager.configuration.ExecutorConfiguration;
import com.lanyuanxiaoyao.service.executor.manager.configuration.HadoopConfiguration;
import java.io.IOException;
import java.nio.charset.Charset;
import java.time.Duration;
import java.util.ArrayList;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import org.apache.flink.client.cli.ClientOptions;
import org.apache.flink.configuration.*;
import org.apache.flink.yarn.configuration.YarnConfigOptions;
import org.apache.flink.yarn.configuration.YarnDeploymentTarget;
import org.apache.hadoop.fs.FSDataInputStream;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.yarn.api.records.ApplicationId;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.eclipse.collections.api.factory.Lists;
import org.eclipse.collections.api.factory.Maps;
import org.eclipse.collections.api.list.ImmutableList;
import org.eclipse.collections.api.list.MutableList;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cache.annotation.Cacheable;
import org.springframework.http.converter.json.Jackson2ObjectMapperBuilder;
import org.springframework.retry.annotation.Retryable;
import org.springframework.stereotype.Service;

import static com.lanyuanxiaoyao.service.common.Constants.HALF_HOUR;
import static com.lanyuanxiaoyao.service.common.Constants.MINUTE;

/**
 * @author lanyuanxiaoyao
 * @date 2024-01-10
 */
@Service
public class ExecutorTaskService {
    private static final Logger logger = LoggerFactory.getLogger(ExecutorTaskService.class);

    private final HadoopConfiguration hadoopConfiguration;
    private final ExecutorConfiguration executorConfiguration;
    private final ObjectMapper mapper;

    public ExecutorTaskService(HadoopConfiguration hadoopConfiguration, ExecutorConfiguration executorConfiguration, Jackson2ObjectMapperBuilder builder) {
        this.hadoopConfiguration = hadoopConfiguration;
        this.executorConfiguration = executorConfiguration;
        this.mapper = builder.build();
    }

    private String taskId() {
        return IdUtil.nanoId(8);
    }

    private void setEnvironment(Configuration configuration, String key, String value) {
        configuration.setString(ResourceManagerOptions.CONTAINERIZED_MASTER_ENV_PREFIX + key, value);
        configuration.setString(ResourceManagerOptions.CONTAINERIZED_TASK_MANAGER_ENV_PREFIX + key, value);
    }

    private Configuration generateConfiguration(String taskId, String name) {
        Configuration configuration = new Configuration();
        configuration.setBoolean(JobManagerOptions.JVM_DIRECT_MEMORY_LIMIT_ENABLED, true);
        configuration.setString(AkkaOptions.ASK_TIMEOUT, "10 min");
        configuration.setString(AkkaOptions.TCP_TIMEOUT, "15 min");
        configuration.setString(AkkaOptions.LOOKUP_TIMEOUT, "10 min");
        configuration.set(ClientOptions.CLIENT_TIMEOUT, Duration.ofMinutes(30));
        // Kerberos认证
        configuration.setBoolean(SecurityOptions.KERBEROS_LOGIN_USETICKETCACHE, true);
        configuration.setString(SecurityOptions.KERBEROS_LOGIN_KEYTAB, hadoopConfiguration.getKerberosKeytabPath());
        configuration.setString(SecurityOptions.KERBEROS_LOGIN_PRINCIPAL, hadoopConfiguration.getKerberosPrincipal());
        configuration.setLong(HeartbeatManagerOptions.HEARTBEAT_INTERVAL, MINUTE);
        configuration.setLong(HeartbeatManagerOptions.HEARTBEAT_TIMEOUT, HALF_HOUR);
        configuration.setString(AkkaOptions.ASK_TIMEOUT, "1 min");
        configuration.setString(AkkaOptions.TCP_TIMEOUT, "2 min");
        configuration.setBoolean(CoreOptions.CHECK_LEAKED_CLASSLOADER, false);
        configuration.setString(YarnConfigOptions.APPLICATION_ATTEMPTS, "4");
        configuration.setString(ResourceManagerOptions.CONTAINERIZED_MASTER_ENV_PREFIX + "MALLOC_ARENA_MAX", "1");
        configuration.setString(ResourceManagerOptions.CONTAINERIZED_TASK_MANAGER_ENV_PREFIX + "MALLOC_ARENA_MAX", "1");
        configuration.setInteger(RestOptions.PORT, 8081);
        configuration.setString(RestOptions.BIND_PORT, "8084-9400");
        configuration.setString(DeploymentOptions.TARGET, YarnDeploymentTarget.APPLICATION.getName());
        configuration.set(JobManagerOptions.TOTAL_PROCESS_MEMORY, MemorySize.parse("2048m"));
        configuration.set(JobManagerOptions.JVM_METASPACE, MemorySize.parse("128m"));
        configuration.set(TaskManagerOptions.TOTAL_FLINK_MEMORY, MemorySize.parse("10240m"));
        configuration.set(TaskManagerOptions.MANAGED_MEMORY_SIZE, MemorySize.parse("0m"));
        configuration.set(TaskManagerOptions.JVM_METASPACE, MemorySize.parse("128m"));
        configuration.setInteger(TaskManagerOptions.NUM_TASK_SLOTS, 10);
        configuration.setString(YarnConfigOptions.APPLICATION_NAME, StrUtil.format("Service_Task {} #{}", name, taskId));

        setEnvironment(configuration, "task_id", taskId);

        configuration.set(PipelineOptions.JARS, new ArrayList<String>() {{
            add(executorConfiguration.getTaskJarPath());
        }});
        return configuration;
    }

    public String scanAvro(
            String key,
            String hdfs,
            String pulsar,
            String pulsarTopic,
            Boolean scanSource,
            Boolean scanQueue,
            Boolean scanLog,
            Boolean scanBase,
            Boolean scanTarget
    ) throws Exception {
        String taskId = taskId();
        Configuration configuration = generateConfiguration(taskId, "scan " + key);
        MapBuilder<String, Object> builder = MapUtil.builder();

        setEnvironment(configuration, "key", key);
        builder.put("key", key);

        if (scanLog || scanBase) {
            setEnvironment(configuration, "hdfs", hdfs);
            builder.put("scan_log", scanLog);
            builder.put("scan_base", scanBase);
            builder.put("hdfs", hdfs);
        }

        if (scanQueue) {
            setEnvironment(configuration, "pulsar", pulsar);
            setEnvironment(configuration, "pulsar_topic", pulsarTopic);
            builder.put("scan_queue", true);
            builder.put("pulsar", pulsar);
            builder.put("pulsar_topic", pulsarTopic);
        }
        ApplicationId applicationId = Runner.run(
                configuration,
                "com.lanyuanxiaoyao.service.executor.task.DataScanner",
                new String[]{
                        TaskConstants.TASK_CONTEXT_OPTION,
                        mapper.writeValueAsString(
                                new TaskContext(
                                        taskId,
                                        executorConfiguration.getTaskResultPath(),
                                        Maps.mutable.ofMap(builder.build())
                                )
                        )
                }
        );
        return applicationId.toString();
    }

    public String scanLatestOpTs(String hdfs) throws Exception {
        String taskId = taskId();
        Configuration configuration = generateConfiguration(taskId, StrUtil.format("latest_op_ts {}", hdfs));
        configuration.set(TaskManagerOptions.MANAGED_MEMORY_SIZE, MemorySize.parse("1024m"));
        MapBuilder<String, Object> builder = MapUtil.builder();

        builder.put("hdfs", hdfs);

        ApplicationId applicationId = Runner.run(
                configuration,
                "com.lanyuanxiaoyao.service.executor.task.LatestOperationTimeScan",
                new String[]{
                        TaskConstants.TASK_CONTEXT_OPTION,
                        mapper.writeValueAsString(
                                new TaskContext(
                                        taskId,
                                        executorConfiguration.getTaskResultPath(),
                                        Maps.mutable.ofMap(builder.build())
                                )
                        )
                }
        );
        return applicationId.toString();
    }

    @Cacheable(value = "results", sync = true)
    @Retryable(Throwable.class)
    public ImmutableList<String> taskResult(String taskId, Integer limit) throws IOException {
        Path resultPath = new Path(executorConfiguration.getTaskResultPath(), taskId);
        MutableList<String> results = Lists.mutable.empty();
        try (FileSystem fileSystem = FileSystem.get(new org.apache.hadoop.conf.Configuration())) {
            if (!fileSystem.exists(resultPath)) {
                throw new RuntimeException(StrUtil.format("Task {} result not found", taskId));
            }
            for (FileStatus status : fileSystem.listStatus(resultPath)) {
                if (status.isFile() && results.size() < limit) {
                    try (FSDataInputStream file = fileSystem.open(status.getPath())) {
                        IoUtil.readLines(file, Charset.defaultCharset(), results);
                    }
                }
            }
        }
        Pattern dateRegex = Pattern.compile("(\\w+) (\\d{17}) (.+)");
        return results
                .reject(StrUtil::isBlank)
                .collect(StrUtil::trim)
                .sortThisBy(line -> {
                    Matcher matcher = dateRegex.matcher(line);
                    if (matcher.matches()) {
                        return Long.valueOf(matcher.group(2));
                    }
                    return 0L;
                })
                .toImmutable();
    }
}
