package com.kamco.cd.training.train.service;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializationFeature;
import com.kamco.cd.training.common.enums.TrainStatusType;
import com.kamco.cd.training.postgres.core.ModelTestMetricsJobCoreService;
import com.kamco.cd.training.train.dto.ModelTrainMetricsDto.ModelMetricJsonDto;
import com.kamco.cd.training.train.dto.ModelTrainMetricsDto.ModelTestFileName;
import com.kamco.cd.training.train.dto.ModelTrainMetricsDto.ResponsePathDto;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.nio.file.StandardOpenOption;
import java.time.ZonedDateTime;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.csv.CSVFormat;
import org.apache.commons.csv.CSVParser;
import org.apache.commons.csv.CSVRecord;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

@Slf4j
@Service
@RequiredArgsConstructor
public class ModelTestMetricsJobService {

  private final ModelTestMetricsJobCoreService modelTestMetricsJobCoreService;

  @Value("${spring.profiles.active}")
  private String profile;

  // 학습 결과가 저장될 호스트 디렉토리
  @Value("${train.docker.responseDir}")
  private String responseDir;

  @Value("${file.pt-path}")
  private String ptPathDir;

  /** 결과 csv 파일 정보 등록 */
  public void findTestValidMetricCsvFiles() {

    List<ResponsePathDto> modelIds =
        modelTestMetricsJobCoreService.getTestMetricSaveNotYetModelIds();

    if (modelIds.isEmpty()) {
      return;
    }

    for (ResponsePathDto modelInfo : modelIds) {

      String testPath = responseDir + "/" + modelInfo.getUuid() + "/metrics/test.csv";
      try (BufferedReader reader =
          Files.newBufferedReader(Paths.get(testPath), StandardCharsets.UTF_8); ) {

        CSVParser parser = CSVFormat.DEFAULT.withFirstRecordAsHeader().parse(reader);

        List<Object[]> batchArgs = new ArrayList<>();

        for (CSVRecord record : parser) {

          String model = record.get("model");
          long TP = Long.parseLong(record.get("TP"));
          long FP = Long.parseLong(record.get("FP"));
          long FN = Long.parseLong(record.get("FN"));
          float precision = Float.parseFloat(record.get("precision"));
          float recall = Float.parseFloat(record.get("recall"));
          float f1_score = Float.parseFloat(record.get("f1_score"));
          float accuracy = Float.parseFloat(record.get("accuracy"));
          float iou = Float.parseFloat(record.get("iou"));
          long detection_count = Long.parseLong(record.get("detection_count"));
          long gt_count = Long.parseLong(record.get("gt_count"));

          batchArgs.add(
              new Object[] {
                modelInfo.getModelId(),
                model,
                TP,
                FP,
                FN,
                precision,
                recall,
                f1_score,
                accuracy,
                iou,
                detection_count,
                gt_count
              });
        }

        modelTestMetricsJobCoreService.insertModelMetricsTest(batchArgs);

        // test.csv 파일 읽어서 저장한 여부로만 사용하기
        modelTestMetricsJobCoreService.updateModelMetricsTrainSaveYn(
            modelInfo.getModelId(), "step2");

      } catch (IOException e) {
        throw new RuntimeException(e);
      }

      // 패키징할 파일 만들기
      modelTestMetricsJobCoreService.updatePackingStart(
          modelInfo.getModelId(), ZonedDateTime.now());

      ModelMetricJsonDto jsonDto =
          modelTestMetricsJobCoreService.getTestMetricPackingInfo(modelInfo.getModelId());
      try {
        writeJsonFile(
            jsonDto,
            Paths.get(
                responseDir
                    + "/"
                    + modelInfo.getUuid()
                    + "/"
                    + jsonDto.getModelVersion()
                    + ".json"));
      } catch (IOException e) {
        throw new RuntimeException(e);
      }

      Path responsePath = Paths.get(responseDir + "/" + modelInfo.getUuid());

      ModelTestFileName fileInfo =
          modelTestMetricsJobCoreService.findModelTestFileNames(modelInfo.getModelId());

      Path zipPath =
          Paths.get(
              responseDir + "/" + modelInfo.getUuid() + "/" + fileInfo.getModelVersion() + ".zip");
      Set<String> targetNames =
          Set.of(
              "model_config.py",
              fileInfo.getBestEpochFileName() + ".pth",
              fileInfo.getModelVersion() + ".json");

      List<Path> files = new ArrayList<>();
      try (Stream<Path> s = Files.list(responsePath)) {
        files.addAll(
            s.filter(Files::isRegularFile)
                .filter(p -> targetNames.contains(p.getFileName().toString()))
                .collect(Collectors.toList()));
      } catch (IOException e) {
        throw new RuntimeException(e);
      }

      try (Stream<Path> s = Files.list(Path.of(ptPathDir))) {
        files.addAll(
            s.filter(Files::isRegularFile)
                .limit(1) // yolov8_6th-6m.pt 파일 1개만
                .collect(Collectors.toList()));
      } catch (IOException e) {
        throw new RuntimeException(e);
      }

      try {
        zipFiles(files, zipPath);

        modelTestMetricsJobCoreService.updatePackingEnd(
            modelInfo.getModelId(), ZonedDateTime.now(), TrainStatusType.COMPLETED.getId());
      } catch (IOException e) {
        modelTestMetricsJobCoreService.updatePackingEnd(
            modelInfo.getModelId(), ZonedDateTime.now(), TrainStatusType.ERROR.getId());
        throw new RuntimeException(e);
      }
    }
  }

  private void writeJsonFile(Object data, Path outputPath) throws IOException {

    Path parent = outputPath.getParent();

    if (parent != null) {
      Files.createDirectories(parent);
    }

    ObjectMapper objectMapper = new ObjectMapper();
    objectMapper.enable(SerializationFeature.INDENT_OUTPUT);

    try (OutputStream os =
        Files.newOutputStream(
            outputPath, StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING)) {
      objectMapper.writeValue(os, data);
    }
  }

  private void zipFiles(List<Path> files, Path zipPath) throws IOException {

    Path parent = zipPath.getParent();
    if (parent != null) {
      Files.createDirectories(parent);
    }

    try (ZipOutputStream zos = new ZipOutputStream(Files.newOutputStream(zipPath))) {

      for (Path file : files) {

        ZipEntry entry = new ZipEntry(file.getFileName().toString());
        zos.putNextEntry(entry);

        Files.copy(file, zos);

        zos.closeEntry();
      }
    }
  }
}
