package tech.easyflow.agent.runtime.artifact;

import com.mybatisflex.core.query.QueryWrapper;
import com.mybatisflex.core.tenant.TenantManager;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;
import tech.easyflow.agent.entity.AgentArtifact;
import tech.easyflow.agent.mapper.AgentArtifactMapper;

import java.util.ArrayList;
import java.util.Date;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

/**
 * 有界领取并补偿清理过期、待删除和删除失败 Artifact。
 */
@Component
public class AgentArtifactCleanupScheduler {

    private static final int BATCH_SIZE = 100;
    private static final String PUBLISH_TIMEOUT_ERROR = "ARTIFACT_PUBLISH_TIMEOUT";
    private static final String SESSION_UNAVAILABLE_ERROR = "ARTIFACT_SESSION_UNAVAILABLE";

    private final AgentArtifactMapper mapper;
    private final AgentArtifactService artifactService;

    /**
     * 创建清理任务。
     *
     * @param mapper Artifact Mapper
     * @param artifactService Artifact 服务
     */
    public AgentArtifactCleanupScheduler(AgentArtifactMapper mapper,
                                         AgentArtifactService artifactService) {
        this.mapper = mapper;
        this.artifactService = artifactService;
    }

    /**
     * 周期性清理，单次最多处理一百条，避免全表扫描和长时间占用调度线程。
     */
    @Scheduled(fixedDelayString = "${easyflow.agent.workspace.cleanup-interval:30m}")
    public void cleanup() {
        TenantManager.withoutTenantCondition(() -> {
            cleanupWithoutTenantCondition();
            return null;
        });
    }

    /**
     * 在关闭 ORM 当前租户条件的作用域中执行全租户清理。
     */
    private void cleanupWithoutTenantCondition() {
        Date now = new Date();
        List<AgentArtifact> expired = mapper.selectListByQuery(QueryWrapper.create()
                .eq(AgentArtifact::getChatMode, AgentArtifactService.MODE_DRAFT)
                .eq(AgentArtifact::getStatus, AgentArtifactStatus.AVAILABLE.name())
                .le(AgentArtifact::getExpiresAt, now)
                .orderBy(AgentArtifact::getId, true)
                .limit(BATCH_SIZE));
        for (AgentArtifact artifact : expired) {
            AgentArtifact update = new AgentArtifact();
            update.setStatus(AgentArtifactStatus.DELETE_PENDING.name());
            update.setNextRetryAt(now);
            update.setModified(now);
            mapper.updateByQuery(update, QueryWrapper.create()
                    .eq(AgentArtifact::getId, artifact.getId())
                    .eq(AgentArtifact::getStatus, AgentArtifactStatus.AVAILABLE.name()));
        }

        Map<java.math.BigInteger, AgentArtifact> candidates = new LinkedHashMap<>();
        List<AgentArtifact> abandonedPublishing = mapper.selectListByQuery(QueryWrapper.create()
                .eq(AgentArtifact::getStatus, AgentArtifactStatus.PUBLISHING.name())
                .le(AgentArtifact::getNextRetryAt, now)
                .orderBy(AgentArtifact::getId, true)
                .limit(BATCH_SIZE));
        for (AgentArtifact artifact : abandonedPublishing) {
            AgentArtifact update = new AgentArtifact();
            update.setStatus(AgentArtifactStatus.DELETE_PENDING.name());
            update.setNextRetryAt(now);
            update.setLastErrorCode(PUBLISH_TIMEOUT_ERROR);
            update.setModified(now);
            int changed = mapper.updateByQuery(update, QueryWrapper.create()
                    .eq(AgentArtifact::getId, artifact.getId())
                    .eq(AgentArtifact::getStatus, AgentArtifactStatus.PUBLISHING.name())
                    .le(AgentArtifact::getNextRetryAt, now));
            if (changed == 1) {
                artifact.setStatus(AgentArtifactStatus.DELETE_PENDING.name());
                artifact.setNextRetryAt(now);
                artifact.setLastErrorCode(PUBLISH_TIMEOUT_ERROR);
                candidates.put(artifact.getId(), artifact);
            }
        }
        if (candidates.size() < BATCH_SIZE) {
            List<AgentArtifact> orphans = mapper.selectOrphanedFormalArtifacts(BATCH_SIZE - candidates.size());
            for (AgentArtifact artifact : orphans) {
                AgentArtifact update = new AgentArtifact();
                update.setStatus(AgentArtifactStatus.DELETE_PENDING.name());
                update.setNextRetryAt(now);
                update.setLastErrorCode(SESSION_UNAVAILABLE_ERROR);
                update.setModified(now);
                int changed = mapper.updateByQuery(update, QueryWrapper.create()
                        .eq(AgentArtifact::getId, artifact.getId())
                        .eq(AgentArtifact::getStatus, artifact.getStatus()));
                if (changed == 1) {
                    artifact.setStatus(AgentArtifactStatus.DELETE_PENDING.name());
                    artifact.setNextRetryAt(now);
                    artifact.setLastErrorCode(SESSION_UNAVAILABLE_ERROR);
                    candidates.put(artifact.getId(), artifact);
                }
            }
        }
        if (candidates.size() < BATCH_SIZE) {
            add(candidates, mapper.selectListByQuery(QueryWrapper.create()
                    .eq(AgentArtifact::getStatus, AgentArtifactStatus.DELETE_PENDING.name())
                    .orderBy(AgentArtifact::getId, true)
                    .limit(BATCH_SIZE - candidates.size())));
        }
        if (candidates.size() < BATCH_SIZE) {
            add(candidates, mapper.selectListByQuery(QueryWrapper.create()
                    .eq(AgentArtifact::getStatus, AgentArtifactStatus.DELETE_FAILED.name())
                    .le(AgentArtifact::getNextRetryAt, now)
                    .orderBy(AgentArtifact::getId, true)
                    .limit(BATCH_SIZE - candidates.size())));
        }
        new ArrayList<>(candidates.values()).forEach(artifactService::deleteObject);
    }

    private void add(Map<java.math.BigInteger, AgentArtifact> target, List<AgentArtifact> artifacts) {
        for (AgentArtifact artifact : artifacts) {
            target.putIfAbsent(artifact.getId(), artifact);
        }
    }
}
