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

import com.kamco.cd.training.postgres.core.ModelTrainJobCoreService;
import com.kamco.cd.training.train.dto.EvalRunRequest;
import com.kamco.cd.training.train.dto.TrainRunRequest;
import com.kamco.cd.training.train.dto.TrainRunResult;
import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import lombok.RequiredArgsConstructor;
import lombok.extern.log4j.Log4j2;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

@Log4j2
@Service
@RequiredArgsConstructor
public class DockerTrainService {

  // 실행할 Docker 이미지명
  @Value("${train.docker.image}")
  private String image;

  // 학습 요청 데이터가 위치한 호스트 디렉토리
  @Value("${train.docker.request_dir}")
  private String requestDir;

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

  // 공유메모리 사이즈 설정 (대용량 학습시 필요)
  @Value("${train.docker.shm_size:16g}")
  private String shmSize;

  // data 경로 request,response 상위 폴더
  @Value("${train.docker.base_path}")
  private String basePath;

  @Value("${train.docker.symbolic_link_dir}")
  private String symbolicDir;

  // IPC host 사용 여부
  @Value("${train.docker.ipc_host:true}")
  private boolean ipcHost;

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

  @Value("${hyper.parameter.gpus}")
  private String hyperGpus;

  @Value("${hyper.parameter.gpu-ids}")
  private String hyperGpuIds;

  @Value("${hyper.parameter.batch-size}")
  private Integer batchSize;

  private final ModelTrainJobCoreService modelTrainJobCoreService;

  /**
   * Docker 학습 컨테이너를 동기 실행 - 요청 스레드에서 docker run 실행 - 컨테이너 종료까지 대기 - stdout/stderr 로그 수집 후 결과 반환
   *
   * @param req
   * @param containerName
   * @return
   * @throws Exception
   */
  public TrainRunResult runTrainSync(TrainRunRequest req, String containerName) throws Exception {

    List<String> cmd = buildDockerRunCommand(containerName, req);

    log.info("=== Docker Train Command ===");
    log.info("Container: {}", containerName);
    log.info("Command: {}", String.join(" ", cmd));
    log.info("================================");

    ProcessBuilder pb = new ProcessBuilder(cmd);
    pb.redirectErrorStream(true);

    // 로그는 별도 스레드에서 읽기 (메인 스레드가 readLine에 안 걸리게)
    StringBuilder logBuilder = new StringBuilder();
    Process p = pb.start();

    log.info("[TRAIN-BOOT] docker run started. container={}", containerName);

    try {
      log.info("[TRAIN-BOOT] pid={}", p.pid()); // Java 9+
    } catch (Throwable ignore) {
    }

    try {
      // 바로 죽었는지 100ms만 체크
      if (p.waitFor(100, TimeUnit.MILLISECONDS)) {
        int exit = p.exitValue();
        String earlyLogs;
        synchronized (logBuilder) {
          earlyLogs = logBuilder.toString();
        }
        log.error(
            "[TRAIN-BOOT] docker run exited immediately. container={} exit={}",
            containerName,
            exit);
        log.error("[TRAIN-BOOT] early logs:\n{}", earlyLogs);
      } else {
        log.info("[TRAIN-BOOT] docker run is still running. container={}", containerName);
      }
    } catch (Exception e) {
      log.warn("[TRAIN-BOOT] early-exit check failed: {}", e.toString(), e);
    }

    Pattern epochPattern = Pattern.compile("Epoch\\(train\\)\\s+\\[(\\d+)\\]\\[(\\d+)/(\\d+)\\]");

    // 너무 잦은 업데이트 방지용
    AtomicInteger lastEpoch = new AtomicInteger(0);
    AtomicInteger lastIter = new AtomicInteger(0);

    Thread logThread =
        new Thread(
            () -> {
              try (BufferedReader br =
                  new BufferedReader(
                      new InputStreamReader(p.getInputStream(), StandardCharsets.UTF_8))) {

                String line;
                while ((line = br.readLine()) != null) {

                  synchronized (logBuilder) {
                    logBuilder.append(line).append('\n');
                  }

                  Matcher m = epochPattern.matcher(line);
                  if (m.find()) {
                    int epoch = Integer.parseInt(m.group(1));
                    int iter = Integer.parseInt(m.group(2));
                    int totalIter = Integer.parseInt(m.group(3));

                    // (선택) maxEpochs는 req에서 알고 있으니 req.getEpochs() 같은 걸로 사용
                    int maxEpochs = req.getEpochs() != null ? req.getEpochs() : 0;

                    // 쓰로틀링: 에폭 끝 or 10 iter마다
                    boolean shouldUpdate = (iter == totalIter) || (iter % 10 == 0);

                    // 중복 방지
                    if (shouldUpdate) {
                      int prevEpoch = lastEpoch.get();
                      int prevIter = lastIter.get();
                      if (epoch != prevEpoch || iter != prevIter) {
                        lastEpoch.set(epoch);
                        lastIter.set(iter);

                        // Step 구분 (컨테이너 이름으로 판별)
                        boolean isStep1 = containerName.startsWith("train-");
                        boolean isStep2 = containerName.startsWith("eval-");

                        // 진행률 계산 (Step별 구분)
                        double progress = 0.0;
                        if (maxEpochs > 0) {
                          // Epoch + Iteration 기반 정밀 진행률 계산
                          double epochProgress =
                              ((double) (epoch - 1) + ((double) iter / totalIter)) / maxEpochs;

                          // Step1 (학습): 7.5% ~ 42.5% (0% ~ 50% 중 15% ~ 85%)
                          // Step2 (테스트): 57.5% ~ 92.5% (50% ~ 100% 중 15% ~ 85%)
                          if (isStep1) {
                            // Step1: 0% ~ 50% 범위, 그 중 15% ~ 85% = 7.5% ~ 42.5%
                            progress = 7.5 + (epochProgress * 35.0);
                          } else if (isStep2) {
                            // Step2: 50% ~ 100% 범위, 그 중 15% ~ 85% = 57.5% ~ 92.5%
                            progress = 57.5 + (epochProgress * 35.0);
                          } else {
                            // 기본값 (하위 호환)
                            progress = 15.0 + (epochProgress * 70.0);
                          }
                        }

                        String stepLabel = isStep1 ? "STEP1" : isStep2 ? "STEP2" : "UNKNOWN";

                        log.info(
                            "[TRAIN] {} container={} epoch={}/{} iter={}/{} | Progress: {}%",
                            stepLabel,
                            containerName,
                            epoch,
                            maxEpochs,
                            iter,
                            totalIter,
                            String.format("%.2f", progress));

                        modelTrainJobCoreService.updateEpoch(containerName, epoch);
                      }
                    }
                  }
                }
              } catch (Exception e) {
                log.warn("logThread error: {}", e.toString());
              }
            },
            "train-log-" + containerName);

    logThread.setDaemon(true);
    logThread.start();

    int timeoutSeconds = req.getTimeoutSeconds() != null ? req.getTimeoutSeconds() : 7200;
    boolean finished = p.waitFor(timeoutSeconds, TimeUnit.SECONDS);

    if (!finished) {
      //  docker run 프로세스도 같이 끊어야 readLine이 풀림
      p.destroy();
      if (!p.waitFor(2, TimeUnit.SECONDS)) {
        p.destroyForcibly();
      }
      killContainer(containerName);

      String logs;
      synchronized (logBuilder) {
        logs = logBuilder.toString();
      }

      return new TrainRunResult(
          null, // jobId (없으면 null)
          containerName,
          -1,
          "TIMEOUT",
          logs);
    }

    int exit = p.exitValue();

    // 로그 스레드가 마무리할 시간을 조금 줌(없어도 되지만 로그 누락 방지용)
    logThread.join(500);

    String logs;
    synchronized (logBuilder) {
      logs = logBuilder.toString();
    }

    return new TrainRunResult(null, containerName, exit, exit == 0 ? "SUCCESS" : "FAILED", logs);
  }

  /**
   * 학습 docker run command
   *
   * @param containerName
   * @param req
   * @return
   */
  private List<String> buildDockerRunCommand(String containerName, TrainRunRequest req) {

    List<String> c = new ArrayList<>();

    c.add("docker");
    c.add("run");

    // 컨테이너 이름 지정
    c.add("--name");
    c.add(containerName);

    // 실행 종료 시 자동 삭제
    c.add("--rm");

    // GPU 전체 사용
    c.add("--gpus");
    c.add("all");

    // IPC host 사용 여부
    if (ipcHost) {
      c.add("--ipc=host");
    }

    // 공유메모리 설정
    c.add("--shm-size=" + shmSize);

    // 메모리 관련 ulimit 설정
    c.add("--ulimit");
    c.add("memlock=-1");
    c.add("--ulimit");
    c.add("stack=67108864");

    // 환경변수 설정
    c.add("-e");
    c.add("OPENCV_LOG_LEVEL=ERROR");
    c.add("-e");
    c.add("NCCL_DEBUG=INFO");
    c.add("-e");
    c.add("NCCL_IB_DISABLE=1");
    c.add("-e");
    c.add("NCCL_P2P_DISABLE=0");
    c.add("-e");
    c.add("NCCL_SOCKET_IFNAME=eth0");

    // 요청/결과 디렉토리 볼륨 마운트
    c.add("-v");
    c.add(basePath + ":" + basePath); // 심볼릭 링크와 연결되는 실제 파일 경로도 마운트를 해줘야 함
    c.add("-v");
    if (req.getReqTmpYn()) {
      c.add(symbolicDir + ":/data"); // 요청할경로 : tmp 심볼릭 사용하는 것이니 symbolicDir로 호출
    } else {
      c.add(requestDir + ":/data"); // 요청할경로 : tmp 심볼릭 사용하지 않으니 request로 호출
    }
    c.add("-v");
    c.add(responseDir + ":/checkpoints"); // 저장될경로

    // 표준입력 유지 (-it 대신 -i만 사용)
    c.add("-i");

    // 사용할 이미지
    c.add(image);

    // ===== 컨테이너 내부 실행 명령 =====
    c.add("python");
    c.add("/workspace/change-detection-code/train_wrapper.py");

    // ===== 기본 파라미터 =====
    addArg(c, "--dataset-folder", req.getDatasetFolder());
    addArg(c, "--output-folder", req.getOutputFolder());
    addArg(c, "--input-size", req.getInputSize());
    addArg(c, "--crop-size", req.getCropSize());
    //    addArg(c, "--batch-size", req.getBatchSize());
    //    addArg(c, "--gpu-ids", req.getGpuIds()); // null
    if ("prod".equals(profile)) {
      addArg(c, "--batch-size", 2); // 학습서버 GPU 1개인 곳은 batch-size:2 까지만 가능

    } else {
      addArg(c, "--batch-size", batchSize); // 학습서버 GPU 1개인 곳은 batch-size:2 까지만 가능
    }
    addArg(c, "--gpus", hyperGpus); // 학습서버 GPU 1개인 곳은 1이어야 함
    addArg(c, "--gpu-ids", hyperGpuIds); // 학습서버 GPU 1개인 곳은 0이어야 함

    addArg(c, "--lr", req.getLearningRate());
    addArg(c, "--backbone", req.getBackbone());
    addArg(c, "--epochs", req.getEpochs());

    // ===== Data =====
    addArg(c, "--train-num-workers", req.getTrainNumWorkers());
    addArg(c, "--val-num-workers", req.getValNumWorkers());
    addArg(c, "--test-num-workers", req.getTestNumWorkers());
    addArg(c, "--train-shuffle", req.getTrainShuffle());
    addArg(c, "--train-persistent", req.getTrainPersistent());
    addArg(c, "--val-persistent", req.getValPersistent());

    // ===== Model Architecture =====
    addArg(c, "--drop-path-rate", req.getDropPathRate());
    addArg(c, "--frozen-stages", req.getFrozenStages());
    addArg(c, "--neck-policy", req.getNeckPolicy());
    addArg(c, "--class-weight", req.getClassWeight());
    addArg(c, "--decoder-channels", req.getDecoderChannels());

    // ===== Loss & Optimization =====
    addArg(c, "--weight-decay", req.getWeightDecay());
    addArg(c, "--layer-decay-rate", req.getLayerDecayRate());
    addArg(c, "--ignore-index", req.getIgnoreIndex());
    addArg(c, "--ddp-find-unused-params", req.getDdpFindUnusedParams());
    addArg(c, "--num-layers", req.getNumLayers());

    // ===== Evaluation =====
    addArg(c, "--metrics", req.getMetrics());
    addArg(c, "--save-best", req.getSaveBest());
    addArg(c, "--save-best-rule", req.getSaveBestRule());
    addArg(c, "--val-interval", req.getValInterval());
    addArg(c, "--log-interval", req.getLogInterval());
    addArg(c, "--vis-interval", req.getVisInterval());

    // ===== Augmentation =====
    addArg(c, "--rot-prob", req.getRotProb());
    addArg(c, "--rot-degree", req.getRotDegree());
    addArg(c, "--flip-prob", req.getFlipProb());
    addArg(c, "--exchange-prob", req.getExchangeProb());
    addArg(c, "--brightness-delta", req.getBrightnessDelta());
    addArg(c, "--contrast-range", req.getContrastRange());
    addArg(c, "--saturation-range", req.getSaturationRange());
    addArg(c, "--hue-delta", req.getHueDelta());

    // 재시작을 위한 파라미터
    if (req.getResumeFrom() != null && !req.getResumeFrom().isBlank()) {
      c.add("--resume");
      addArg(c, "--load-from", req.getResumeFrom());
    } else if ("TRANSFER".equals(req.getTrainType())
        && req.getBeforeModelUUID() != null
        && req.getBeforeBestEpoch() != null) {
      // 재시작이 아닌데 전이학습인 경우, 전이학습을 위한 파라미터
      String finalLoadFromPath;
      String bestModelFileName = "best_changed_fscore_epoch_" + req.getBeforeBestEpoch() + ".pth";
      String epochFileName = "epoch_" + req.getBeforeBestEpoch() + ".pth";

      // 이전 모델 response 경로
      Path beforeModelDir = Path.of(responseDir).resolve(req.getBeforeModelUUID().toString());
      // 이전 모델 response 의 베스트 fscore pth 모델 파일 경로
      Path bestFilePath = beforeModelDir.resolve(bestModelFileName);

      // 베스트 fscore 파일이 있으면 그 경로를 사용, 없으면 epoch_00.pth 파일 경로 사용
      if (Files.exists(bestFilePath)) {
        finalLoadFromPath = bestFilePath.toString();
      } else {
        finalLoadFromPath = beforeModelDir.resolve(epochFileName).toString();
      }
      addArg(c, "--load-from", finalLoadFromPath);
    }

    addArg(c, "--save-interval", 1);
    return c;
  }

  /** 인자 추가(키 + 값) - null / blank면 아예 추가 안 함 */
  private void addArg(List<String> c, String key, Object value) {
    if (value == null) return;
    String s = String.valueOf(value).trim();
    if (s.isEmpty()) return;

    c.add(key + "=" + s);
  }

  /** 컨테이너 강제 종료 및 제거 */
  public void killContainer(String containerName) {
    try {
      new ProcessBuilder("docker", "rm", "-f", containerName)
          .redirectErrorStream(true)
          .start()
          .waitFor(10, TimeUnit.SECONDS);
    } catch (Exception ignored) {
    }
  }

  /**
   * Docker 학습 컨테이너를 동기 실행 - 요청 스레드에서 docker run 실행 - 컨테이너 종료까지 대기 - stdout/stderr 로그 수집 후 결과 반환
   *
   * @param containerName
   * @param req
   * @return
   * @throws Exception
   */
  public TrainRunResult runEvalSync(String containerName, EvalRunRequest req) throws Exception {

    List<String> cmd = buildDockerEvalCommand(containerName, req);

    log.info("=== Docker Test Command ===");
    log.info("Container: {}", containerName);
    log.info("Command: {}", String.join(" ", cmd));
    log.info("================================");

    ProcessBuilder pb = new ProcessBuilder(cmd);
    pb.redirectErrorStream(true);

    Process p = pb.start();

    StringBuilder log = new StringBuilder();
    Thread logThread =
        new Thread(
            () -> {
              try (BufferedReader br =
                  new BufferedReader(
                      new InputStreamReader(p.getInputStream(), StandardCharsets.UTF_8))) {
                String line;
                while ((line = br.readLine()) != null) {
                  synchronized (log) {
                    log.append(line).append('\n');
                  }
                }
              } catch (Exception ignored) {
              }
            });

    logThread.setDaemon(true);
    logThread.start();

    int timeout = req.getTimeoutSeconds() != null ? req.getTimeoutSeconds() : 7200;
    boolean finished = p.waitFor(timeout, TimeUnit.SECONDS);

    if (!finished) {
      p.destroyForcibly();
      killContainer(containerName);

      String logs;
      synchronized (log) {
        logs = log.toString();
      }
      return new TrainRunResult(null, containerName, -1, "TIMEOUT", logs);
    }

    int exit = p.exitValue();
    logThread.join(500);

    String logs;
    synchronized (log) {
      logs = log.toString();
    }

    return new TrainRunResult(null, containerName, exit, exit == 0 ? "SUCCESS" : "FAILED", logs);
  }

  /**
   * 테스트 docker run command
   *
   * @param containerName
   * @param req
   * @return
   */
  private List<String> buildDockerEvalCommand(String containerName, EvalRunRequest req) {

    String uuid = req.getUuid();
    Integer epoch = req.getEpoch();
    if (uuid == null || uuid.isBlank()) throw new IllegalArgumentException("uuid is required");
    if (epoch == null || epoch <= 0) throw new IllegalArgumentException("epoch must be > 0");

    Path epochPath = Paths.get(responseDir, req.getOutputFolder());
    // 결과 폴더에 파라미터로 받은 베스트 epoch이 best_changed_fscore_epoch_ 로 시작하는 파일이 있는지 확인 후 pth 파일명 반환
    String modelFile = findCheckpoint(epochPath, epoch);

    List<String> c = new ArrayList<>();

    c.add("docker");
    c.add("run");
    c.add("--rm");

    c.add("--user"); // root 권한 문제를 해결하기 위해 임시 추가
    c.add("0:0");

    c.add("--gpus");
    c.add("all");

    c.add("--ipc=host");
    c.add("--shm-size=" + shmSize);

    c.add("-v");
    c.add(basePath + ":" + basePath); // 심볼릭 링크와 연결되는 실제 파일 경로도 마운트를 해줘야 함

    c.add("-v");
    if (req.getReqTmpYn()) {
      c.add(symbolicDir + ":/data"); // tmp 사용하는 모델은 심볼릭 링크
    } else {
      c.add(requestDir + ":/data"); // tmp 사용하지 않는 모델은 request 경로
    }

    c.add("-v");
    // c.add(responseDir + ":/checkpoints");
    c.add(responseDir + "/" + req.getOutputFolder() + ":/checkpoints");

    c.add("kamco-cd-train:latest");

    c.add("python");
    c.add("/workspace/change-detection-code/run_evaluation_pipeline.py");

    addArg(c, "--dataset-folder", req.getDatasetFolder());
    // addArg(c, "--output-folder", req.getOutputFolder());

    c.add("--epoch");
    c.add(modelFile);

    return c;
  }

  public String findCheckpoint(Path dir, int epoch) {

    String bestFileName = String.format("best_changed_fscore_epoch_%d.pth", epoch);
    String normalFileName = String.format("epoch_%d.pth", epoch);

    Path bestPath = dir.resolve(bestFileName);
    Path normalPath = dir.resolve(normalFileName);

    // 1. best 파일이 존재하면 그거 사용
    if (Files.isRegularFile(bestPath)) {
      return bestFileName;
    }

    // 2. 없으면 일반 epoch 파일 사용
    if (Files.isRegularFile(normalPath)) {
      return normalFileName;
    }

    throw new IllegalStateException("Checkpoint 파일이 없습니다. epoch=" + epoch);
  }
}
