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.postgres.core.ModelTrainJobCoreService;
import com.kamco.cd.training.postgres.core.ModelTrainMngCoreService;
import com.kamco.cd.training.train.dto.EvalRunRequest;
import com.kamco.cd.training.train.dto.ModelTrainJobDto;
import com.kamco.cd.training.train.dto.ModelTrainJobQueuedEvent;
import com.kamco.cd.training.train.dto.TrainRunRequest;
import com.kamco.cd.training.train.dto.TrainRunResult;
import java.util.Map;
import lombok.RequiredArgsConstructor;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Component;
import org.springframework.transaction.event.TransactionPhase;
import org.springframework.transaction.event.TransactionalEventListener;

@Component
@RequiredArgsConstructor
public class TrainJobWorker {

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

  @Async
  @TransactionalEventListener(phase = TransactionPhase.AFTER_COMMIT)
  public void handle(ModelTrainJobQueuedEvent event) {

    Long jobId = event.getJobId();

    ModelTrainJobDto job =
        modelTrainJobCoreService
            .findById(jobId)
            .orElseThrow(() -> new IllegalArgumentException("Job not found: " + jobId));

    if (TrainStatusType.STOPPED.getId().equals(job.getStatusCd())) {
      return;
    }

    Long modelId = job.getModelId();
    Map<String, Object> params = job.getParamsJson();

    String jobType = params != null ? String.valueOf(params.get("jobType")) : null;

    boolean isEval = "EVAL".equals(jobType);

    String containerName =
        (isEval ? "eval-" : "train-") + jobId + "-" + params.get("uuid").toString().substring(0, 8);

    Integer totalEpoch = null;
    if (params.containsKey("totalEpoch")) {
      if (params.get("totalEpoch") != null) {
        totalEpoch = Integer.parseInt(params.get("totalEpoch").toString());
      }
    }

    modelTrainJobCoreService.markRunning(jobId, containerName, null, "TRAIN_WORKER", totalEpoch);

    try {
      TrainRunResult result;

      if (isEval) {
        modelTrainMngCoreService.markStep2InProgress(modelId, jobId);
        String uuid = String.valueOf(params.get("uuid"));
        int epoch = (int) params.get("epoch");

        EvalRunRequest evalReq = new EvalRunRequest(uuid, epoch, null);
        result = dockerTrainService.runEvalSync(evalReq, containerName);

      } else {
        modelTrainMngCoreService.markStep1InProgress(modelId, jobId);
        TrainRunRequest trainReq = toTrainRunRequest(params);
        result = dockerTrainService.runTrainSync(trainReq, containerName);
      }

      ModelTrainJobDto latest =
          modelTrainJobCoreService
              .findById(jobId)
              .orElseThrow(() -> new IllegalArgumentException("Job not found: " + jobId));

      if (TrainStatusType.STOPPED.getId().equals(latest.getStatusCd())) {
        return;
      }

      if (result.getExitCode() == 0) {
        modelTrainJobCoreService.markSuccess(jobId, result.getExitCode());

        if (isEval) {
          modelTrainMngCoreService.markStep2Success(modelId);
        } else {
          modelTrainMngCoreService.markStep1Success(modelId);
        }

      } else {
        modelTrainJobCoreService.markFailed(
            jobId, result.getExitCode(), result.getStatus() + "\n" + result.getLogs());

        if (isEval) {
          modelTrainMngCoreService.markStep2Error(modelId, "exit=" + result.getExitCode());
        } else {
          modelTrainMngCoreService.markError(modelId, "exit=" + result.getExitCode());
        }
      }

    } catch (Exception e) {
      modelTrainJobCoreService.markFailed(jobId, null, e.toString());

      if ("EVAL".equals(params.get("jobType"))) {
        modelTrainMngCoreService.markStep2Error(modelId, e.getMessage());
      } else {
        modelTrainMngCoreService.markError(modelId, e.getMessage());
      }
    }
  }

  private TrainRunRequest toTrainRunRequest(Map<String, Object> paramsJson) {
    if (paramsJson == null || paramsJson.isEmpty()) {
      return null;
    }

    return objectMapper.convertValue(paramsJson, TrainRunRequest.class);
  }
}
