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

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.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import lombok.extern.log4j.Log4j2;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

@Log4j2
@Service
public class DockerTrainService {

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

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

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

  // 컨테이너 이름 prefix
  @Value("${train.docker.containerPrefix}")
  private String containerPrefix;

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

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

  /** Docker 학습 컨테이너를 동기 실행 - 요청 스레드에서 docker run 실행 - 컨테이너 종료까지 대기 - stdout/stderr 로그 수집 후 결과 반환 */
  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);

    Process p = pb.start();

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

    Pattern epochPattern = Pattern.compile("(?i)\\bepoch\\s*\\[?(\\d+)\\s*/\\s*(\\d+)\\]?\\b");

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

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

                  // 1) 로그 누적
                  synchronized (logBuilder) {
                    logBuilder.append(line).append('\n');
                  }

                  // 2) epoch 감지 + DB 업데이트
                  Matcher m = epochPattern.matcher(line);
                  if (m.find()) {
                    int currentEpoch = Integer.parseInt(m.group(1));
                    int totalEpoch = Integer.parseInt(m.group(2));

                    log.info("[EPOCH] container={} {}/{}", containerName, currentEpoch, totalEpoch);

                    // TODO 실행중인 에폭 저장 필요하면 만들어야함
                    // TODO 하지만 여기서 트랜젝션 걸리는 db 작업하면 안좋다고하는데..?
                    //                modelTrainMngCoreService.updateCurrentEpoch(modelId,
                    // currentEpoch, totalEpoch);
                  }
                }
              } catch (Exception e) {
                log.warn("logThread error: {}", e.toString());
              }
            },
            "train-log-" + containerName);
    //        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) {
    //              }
    //            },
    //            "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 + "-" + req.getUuid().substring(0, 8));

    // 실행 종료 시 자동 삭제
    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(requestDir + ":/data");
    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", "4BDBBDF99D04477A927CC9EBA760B845" /*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());
    //    addArg(c, "--gpus", req.getGpus());
    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());

    addArg(c, "--resume-from", req.getResumeFrom());
    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);
    c.add(s);
  }

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

  public TrainRunResult runEvalSync(EvalRunRequest req, String containerName) 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");

    String modelFile = "best_changed_fscore_epoch_" + epoch + ".pth";

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

    c.add("docker");
    c.add("run");
    c.add("--name");
    c.add(containerName + "=" + req.getUuid().substring(0, 8));
    c.add("--rm");

    c.add("--gpus");
    c.add("all");
    if (ipcHost) c.add("--ipc=host");
    c.add("--shm-size=" + shmSize);

    c.add("-v");
    c.add(requestDir + ":/data");
    c.add("-v");
    c.add(responseDir + ":/checkpoints");

    c.add(image);

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

    c.add("--dataset_dir");
    c.add("/data/" + uuid);

    c.add("--model");
    c.add("/checkpoints/" + uuid + "/" + modelFile);

    return c;
  }
}
