package com.kamco.cd.training.postgres.repository.train;

import static com.kamco.cd.training.postgres.entity.QModelTestTrainingRunEntity.modelTestTrainingRunEntity;
import static com.kamco.cd.training.postgres.entity.QModelTrainJobEntity.modelTrainJobEntity;

import com.kamco.cd.training.common.enums.JobStatusType;
import com.kamco.cd.training.common.enums.JobType;
import com.kamco.cd.training.postgres.entity.ModelTrainJobEntity;
import com.kamco.cd.training.postgres.entity.QModelTrainJobEntity;
import com.querydsl.jpa.impl.JPAQueryFactory;
import jakarta.persistence.EntityManager;
import java.util.List;
import java.util.Optional;
import org.springframework.stereotype.Repository;

@Repository
public class ModelTrainJobRepositoryImpl implements ModelTrainJobRepositoryCustom {

  private final JPAQueryFactory queryFactory;

  public ModelTrainJobRepositoryImpl(EntityManager em) {
    this.queryFactory = new JPAQueryFactory(em);
  }

  /** modelId의 attempt_no 최대값. (없으면 0) */
  @Override
  public int findMaxAttemptNo(Long modelId) {
    QModelTrainJobEntity j = modelTrainJobEntity;

    Integer max =
        queryFactory.select(j.attemptNo.max()).from(j).where(j.modelId.eq(modelId)).fetchOne();

    return max != null ? max : 0;
  }

  /**
   * modelId의 최신 job 1건 (보통 id desc / queuedDttm desc 등) - attemptNo 기준으로도 가능하지만, 여기선 id desc가 가장
   * 단순.
   */
  @Override
  public Optional<ModelTrainJobEntity> findLatestByModelId(Long modelId) {
    QModelTrainJobEntity j = modelTrainJobEntity;

    ModelTrainJobEntity job =
        queryFactory.selectFrom(j).where(j.modelId.eq(modelId)).orderBy(j.id.desc()).fetchFirst();

    return Optional.ofNullable(job);
  }

  @Override
  public Optional<ModelTrainJobEntity> findByContainerName(String containerName) {
    QModelTrainJobEntity j = modelTrainJobEntity;

    ModelTrainJobEntity job =
        queryFactory
            .selectFrom(j)
            .where(j.containerName.eq(containerName))
            .orderBy(j.id.desc())
            .fetchFirst();

    return Optional.ofNullable(job);
  }

  @Override
  public void insertModelTestTrainingRun(Long modelId, Long jobId, int epoch) {
    Integer maxAttemptNo =
        queryFactory
            .select(modelTestTrainingRunEntity.attemptNo.max().coalesce(0))
            .from(modelTestTrainingRunEntity)
            .where(modelTestTrainingRunEntity.modelId.eq(modelId))
            .fetchOne();

    int nextAttemptNo = (maxAttemptNo == null ? 1 : maxAttemptNo + 1);

    queryFactory
        .insert(modelTestTrainingRunEntity)
        .columns(
            modelTestTrainingRunEntity.modelId,
            modelTestTrainingRunEntity.attemptNo,
            modelTestTrainingRunEntity.jobId,
            modelTestTrainingRunEntity.epoch)
        .values(modelId, nextAttemptNo, jobId, epoch)
        .execute();
  }

  @Override
  public List<ModelTrainJobEntity> findRunningJobs() {
    return queryFactory
        .select(modelTrainJobEntity)
        .from(modelTrainJobEntity)
        .where(
            modelTrainJobEntity
                .statusCd
                .eq(JobStatusType.RUNNING.getId())
                .and(modelTrainJobEntity.jobType.eq(JobType.TRAIN.getId())))
        .orderBy(modelTrainJobEntity.id.desc())
        .fetch();
  }
}
