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

import com.fasterxml.jackson.databind.ObjectMapper;
import com.kamco.cd.training.common.enums.TrainStatusType;
import com.kamco.cd.training.model.dto.ModelTrainMngDto;
import com.kamco.cd.training.postgres.core.ModelTrainJobCoreService;
import com.kamco.cd.training.postgres.core.ModelTrainMngCoreService;
import com.kamco.cd.training.train.dto.ModelTrainJobQueuedEvent;
import com.kamco.cd.training.train.dto.TrainRunRequest;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.time.ZonedDateTime;
import java.util.Map;
import java.util.UUID;
import lombok.RequiredArgsConstructor;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.ApplicationEventPublisher;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

@Service
@RequiredArgsConstructor
@Transactional(readOnly = true)
public class TrainJobService {

  private final ModelTrainJobCoreService modelTrainJobCoreService;
  private final ModelTrainMngCoreService modelTrainMngCoreService;
  private final DockerTrainService dockerTrainService;
  private final ObjectMapper objectMapper;
  private final ApplicationEventPublisher eventPublisher;

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

  public Long getModelIdByUuid(UUID uuid) {
    return modelTrainMngCoreService.findModelIdByUuid(uuid);
  }

  /** 실행 예약 (QUEUE 등록) */
  @Transactional
  public Long enqueue(Long modelId) {

    // 마스터 존재 확인(없으면 예외)
    modelTrainMngCoreService.findModelById(modelId);

    // 파라미터 조회
    TrainRunRequest trainRunRequest = modelTrainMngCoreService.findTrainRunRequest(modelId);

    if (trainRunRequest == null) {
      throw new IllegalArgumentException("Model not found: " + modelId);
    }

    int nextAttemptNo = modelTrainJobCoreService.findMaxAttemptNo(modelId) + 1;

    @SuppressWarnings("unchecked")
    Map<String, Object> paramsMap = objectMapper.convertValue(trainRunRequest, Map.class);
    paramsMap.put("jobType", "TRAIN");
    paramsMap.put("uuid", trainRunRequest.getUuid());
    paramsMap.put("totalEpoch", trainRunRequest.getEpochs());

    Long jobId =
        modelTrainJobCoreService.createQueuedJob(
            modelId, nextAttemptNo, paramsMap, ZonedDateTime.now());

    modelTrainMngCoreService.clearLastError(modelId);
    modelTrainMngCoreService.markInProgress(modelId, jobId);

    // 커밋 이후 Worker 실행 트리거(리스너에서 AFTER_COMMIT로 받아야 함)
    eventPublisher.publishEvent(new ModelTrainJobQueuedEvent(jobId));
    return jobId;
  }

  /**
   * 재시작
   *
   * <p>- STOPPED / ERROR 상태에서만 가능 (IN_PROGRESS면 예외) - 이전 params_json 재사용 - 새 attempt 생성
   */
  @Transactional
  public Long restart(Long modelId) {
    return createNextAttempt(modelId, ResumeMode.NONE);
  }

  /**
   * 이어하기
   *
   * @param modelId
   * @return
   */
  @Transactional
  public Long resume(Long modelId) {
    return createNextAttempt(modelId, ResumeMode.REQUIRE);
  }

  /**
   * 중단
   *
   * <p>- job 상태 CANCELED - master 상태 STOPPED
   *
   * <p>※ 실제 docker stop은 Worker/Runner가 수행(운영 안정)
   */
  @Transactional
  public void cancel(Long modelId) {

    ModelTrainMngDto.Basic master = modelTrainMngCoreService.findModelById(modelId);

    Long jobId = master.getCurrentAttemptId();
    if (jobId == null) {
      throw new IllegalStateException("실행중인 작업이 없습니다.");
    }

    var job =
        modelTrainJobCoreService
            .findById(jobId)
            .orElseThrow(() -> new IllegalStateException("Job not found"));

    String containerName = job.getContainerName();

    // 1) 컨테이너 강제 종료 + 제거 (없거나 이미 죽었어도 괜찮게)
    if (containerName != null && !containerName.isBlank()) {
      dockerTrainService.killContainer(containerName);
    }

    // 2) 상태 업데이트 (항상 수행)
    modelTrainJobCoreService.markCanceled(jobId);
    modelTrainMngCoreService.markStopped(modelId);
  }

  private Long createNextAttempt(Long modelId, ResumeMode mode) {

    ModelTrainMngDto.Basic master = modelTrainMngCoreService.findModelById(modelId);

    if (TrainStatusType.IN_PROGRESS.getId().equals(master.getStatusCd())) {
      throw new IllegalStateException("이미 진행중입니다.");
    }

    var lastJob =
        modelTrainJobCoreService
            .findLatestByModelId(modelId)
            .orElseThrow(() -> new IllegalStateException("이전 실행 이력이 없습니다."));

    int nextAttemptNo = modelTrainJobCoreService.findMaxAttemptNo(modelId) + 1;

    // 이전 params_json 재사용 (재현성)
    Map<String, Object> params = lastJob.getParamsJson();
    if (params == null || params.isEmpty()) {
      throw new IllegalStateException("이전 실행 params_json이 없습니다.");
    }

    // mode에 따라 resume 옵션 주입/제거
    Map<String, Object> nextParams = new java.util.LinkedHashMap<>(params);

    if (mode == ResumeMode.NONE) {
      // 이어하기 관련 키가 있다면 제거 (완전 새로 시작 보장)
      nextParams.remove("resumeFrom");
      nextParams.remove("resume");
    } else if (mode == ResumeMode.REQUIRE) {
      // 체크포인트 탐지해서 resumeFrom 세팅
      String resumeFrom = findResumeFromOrNull(nextParams);
      if (resumeFrom == null) {
        throw new IllegalStateException("이어하기 체크포인트가 없습니다.");
      }
      nextParams.put("resumeFrom", resumeFrom);
      nextParams.put("resume", true);
    }

    Long jobId =
        modelTrainJobCoreService.createQueuedJob(
            modelId, nextAttemptNo, nextParams, ZonedDateTime.now());

    modelTrainMngCoreService.clearLastError(modelId);
    modelTrainMngCoreService.markInProgress(modelId, jobId);

    eventPublisher.publishEvent(new ModelTrainJobQueuedEvent(jobId));
    return jobId;
  }

  private enum ResumeMode {
    NONE, // 새로 시작
    REQUIRE // 이어하기
  }

  public String findResumeFromOrNull(Map<String, Object> paramsJson) {
    if (paramsJson == null) return null;

    Object out = paramsJson.get("outputFolder");
    if (out == null) return null;

    String outputFolder = String.valueOf(out).trim(); // uuid
    if (outputFolder.isEmpty()) return null;

    // 호스트 기준 경로
    Path outDir = Paths.get(responseDir, outputFolder);

    Path last = outDir.resolve("last_checkpoint");
    if (!Files.isRegularFile(last)) return null;

    try {
      String ckptFile = Files.readString(last).trim(); // epoch_10.pth
      if (ckptFile.isEmpty()) return null;

      Path ckptHost = outDir.resolve(ckptFile);
      if (!Files.isRegularFile(ckptHost)) return null;

      // 컨테이너 경로 반환
      return "/checkpoints/" + outputFolder + "/" + ckptFile;

    } catch (Exception e) {
      return null;
    }
  }
}
