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

import static com.kamco.cd.training.postgres.entity.QModelMasterEntity.modelMasterEntity;

import com.kamco.cd.training.common.enums.TrainStatusType;
import com.kamco.cd.training.postgres.entity.ModelMetricsTrainEntity;
import com.kamco.cd.training.train.dto.ModelTrainMetricsDto.ResponsePathDto;
import com.querydsl.core.types.Projections;
import com.querydsl.jpa.impl.JPAQueryFactory;
import java.util.List;
import org.springframework.data.jpa.repository.support.QuerydslRepositorySupport;
import org.springframework.jdbc.core.JdbcTemplate;

public class ModelTrainMetricsJobRepositoryImpl extends QuerydslRepositorySupport
    implements ModelTrainMetricsJobRepositoryCustom {

  private final JPAQueryFactory queryFactory;
  private final JdbcTemplate jdbcTemplate;

  public ModelTrainMetricsJobRepositoryImpl(
      JPAQueryFactory queryFactory, JdbcTemplate jdbcTemplate) {
    super(ModelMetricsTrainEntity.class);
    this.queryFactory = queryFactory;
    this.jdbcTemplate = jdbcTemplate;
  }

  @Override
  public List<ResponsePathDto> getTrainMetricSaveNotYetModelIds() {
    return queryFactory
        .select(
            Projections.constructor(
                ResponsePathDto.class,
                modelMasterEntity.id,
                modelMasterEntity.responsePath,
                modelMasterEntity.uuid))
        .from(modelMasterEntity)
        .where(
            modelMasterEntity.step1EndDttm.isNotNull(),
            modelMasterEntity.step1State.eq(TrainStatusType.COMPLETED.getId()),
            modelMasterEntity
                .step1MetricSaveYn
                .isNull()
                .or(modelMasterEntity.step1MetricSaveYn.isFalse()))
        .fetch();
  }

  @Override
  public ResponsePathDto getTrainMetricSaveNotYetModelId(Long modelId) {
    return queryFactory
        .select(
            Projections.constructor(
                ResponsePathDto.class,
                modelMasterEntity.id,
                modelMasterEntity.responsePath,
                modelMasterEntity.uuid))
        .from(modelMasterEntity)
        .where(
            modelMasterEntity.id.eq(modelId),
            modelMasterEntity
                .step1MetricSaveYn
                .isNull()
                .or(modelMasterEntity.step1MetricSaveYn.isFalse()))
        .fetchOne();
  }

  @Override
  public void insertModelMetricsTrain(List<Object[]> batchArgs) {
    String sql =
        """
            insert into tb_model_metrics_train
            (model_id, epoch, iteration, loss, lr, duration_time)
            values (?, ?, ?, ?, ?, ?)
        """;

    jdbcTemplate.batchUpdate(sql, batchArgs);
  }

  @Override
  public void updateModelMetricsTrainSaveYn(Long modelId, String stepNo) {
    queryFactory
        .update(modelMasterEntity)
        .set(
            stepNo.equals("step1")
                ? modelMasterEntity.step1MetricSaveYn
                : modelMasterEntity.step2MetricSaveYn,
            true)
        .where(modelMasterEntity.id.eq(modelId))
        .execute();
  }

  @Override
  public void insertModelMetricsValidation(List<Object[]> batchArgs) {
    String sql =
        """
            insert into tb_model_metrics_validation
            (model_id, epoch, a_acc, m_fscore, m_precision, m_recall, m_iou, m_acc, changed_fscore, changed_precision, changed_recall,
             unchanged_fscore, unchanged_precision, unchanged_recall
             )
            values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """;

    jdbcTemplate.batchUpdate(sql, batchArgs);
  }

  @Override
  public void updateModelSelectedBestEpoch(Long modelId, Integer epoch) {
    queryFactory
        .update(modelMasterEntity)
        .set(modelMasterEntity.bestEpoch, epoch)
        .where(modelMasterEntity.id.eq(modelId))
        .execute();
  }
}
