package tech.easyflow.ai.documentimport.task;

import org.junit.Test;
import tech.easyflow.ai.documentimport.DocumentImportBatchDtos;
import tech.easyflow.ai.entity.DocumentImportBatch;
import tech.easyflow.ai.entity.DocumentImportBatchItem;
import tech.easyflow.ai.enums.DocumentImportBatchItemStage;
import tech.easyflow.ai.enums.DocumentImportBatchItemStatus;
import tech.easyflow.ai.enums.DocumentImportBatchStatus;
import tech.easyflow.ai.enums.DocumentImportMode;
import tech.easyflow.ai.mapper.DocumentImportBatchItemMapper;
import tech.easyflow.ai.mapper.DocumentImportBatchMapper;
import tech.easyflow.ai.service.DocumentImportBatchItemService;
import tech.easyflow.ai.service.DocumentImportBatchService;

import java.math.BigInteger;
import java.util.Date;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyBoolean;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

/**
 * 文档批量导入状态汇总测试。
 */
public class DocumentImportBatchTrackerTest {

    /**
     * 验证运行批次的完成、处理、失败与等待数量。
     */
    @Test
    public void shouldAggregateRunningBatchProgress() {
        DocumentImportBatchService batchService = mock(DocumentImportBatchService.class);
        DocumentImportBatchItemService itemService = mock(DocumentImportBatchItemService.class);
        DocumentImportBatchMapper batchMapper = mock(DocumentImportBatchMapper.class);
        DocumentImportBatchItemMapper itemMapper = mock(DocumentImportBatchItemMapper.class);
        DocumentImportBatch batch = batch(DocumentImportBatchStatus.RUNNING, 4);
        batch.setCompletedCount(1);
        batch.setProcessingCount(1);
        batch.setFailedCount(1);
        batch.setPendingCount(1);
        when(batchService.getById(batch.getId())).thenReturn(batch);

        DocumentImportBatchTracker tracker =
            new DocumentImportBatchTracker(batchService, itemService, batchMapper, itemMapper);
        DocumentImportBatchDtos.StatusResponse response = tracker.refreshBatch(batch.getId());

        assertEquals(1, response.getCompletedCount().intValue());
        assertEquals(1, response.getProcessingCount().intValue());
        assertEquals(1, response.getFailedCount().intValue());
        assertEquals(1, response.getPendingCount().intValue());
        assertEquals(50, response.getProgressPercent().intValue());
        assertEquals(DocumentImportBatchStatus.RUNNING.name(), response.getStatus());
        verify(batchService, never()).updateById(batch, false);
    }

    /**
     * 验证兼容可重试计数与人工可继续的失败总数保持一致。
     */
    @Test
    public void shouldExposeAllFailuresAsContinuable() {
        DocumentImportBatchService batchService = mock(DocumentImportBatchService.class);
        DocumentImportBatchItemService itemService = mock(DocumentImportBatchItemService.class);
        DocumentImportBatchMapper batchMapper = mock(DocumentImportBatchMapper.class);
        DocumentImportBatchItemMapper itemMapper = mock(DocumentImportBatchItemMapper.class);
        DocumentImportBatch batch = batch(DocumentImportBatchStatus.INTERRUPTED, 2);
        batch.setCompletedCount(1);
        batch.setFailedCount(1);
        batch.setRetryableFailedCount(0);
        batch.setInterruptCode("redis_unavailable");
        batch.setInterruptMessage("缓存与消息服务异常");
        batch.setInterruptedAt(new Date());
        when(batchService.getById(batch.getId())).thenReturn(batch);

        DocumentImportBatchTracker tracker =
            new DocumentImportBatchTracker(batchService, itemService, batchMapper, itemMapper);
        DocumentImportBatchDtos.StatusResponse response = tracker.refreshBatch(batch.getId());

        assertEquals(DocumentImportBatchStatus.INTERRUPTED.name(), response.getStatus());
        assertEquals(100, response.getProgressPercent().intValue());
        assertEquals(1, response.getRetryableFailedCount().intValue());
        assertEquals("redis_unavailable", response.getInterruptCode());
        assertEquals("缓存与消息服务异常", response.getInterruptMessage());
        assertEquals(batch.getInterruptedAt(), response.getInterruptedAt());
    }

    /**
     * 验证失败项重试通过状态 CAS 增量迁移计数。
     */
    @Test
    public void shouldMoveRetryableFailureBackToPendingAtomically() {
        DocumentImportBatchService batchService = mock(DocumentImportBatchService.class);
        DocumentImportBatchItemService itemService = mock(DocumentImportBatchItemService.class);
        DocumentImportBatchMapper batchMapper = mock(DocumentImportBatchMapper.class);
        DocumentImportBatchItemMapper itemMapper = mock(DocumentImportBatchItemMapper.class);
        DocumentImportBatch batch = batch(DocumentImportBatchStatus.RUNNING, 1);
        batch.setFailedCount(1);
        batch.setRetryableFailedCount(1);
        DocumentImportBatchItem item = new DocumentImportBatchItem();
        item.setId(BigInteger.TEN);
        item.setBatchId(batch.getId());
        item.setStage(DocumentImportBatchItemStage.INDEX.name());
        item.setStatus(DocumentImportBatchItemStatus.FAILED.name());
        item.setRetryable(true);
        when(itemService.getById(item.getId())).thenReturn(item);
        when(itemMapper.transitionStatus(
            eq(item.getId()),
            eq(DocumentImportBatchItemStatus.FAILED.name()),
            eq(DocumentImportBatchItemStage.INDEX.name()),
            eq(DocumentImportBatchItemStatus.PENDING.name()),
            eq(null),
            eq(null),
            eq(false),
            eq(1),
            any(Date.class)
        )).thenReturn(1);
        when(batchService.getById(batch.getId())).thenReturn(batch);

        DocumentImportBatchTracker tracker =
            new DocumentImportBatchTracker(batchService, itemService, batchMapper, itemMapper);

        assertTrue(tracker.updateItem(item.getId(),
            DocumentImportBatchItemStage.INDEX,
            DocumentImportBatchItemStatus.PENDING,
            null));
        verify(batchMapper).adjustCounters(
            eq(batch.getId()),
            eq(0),
            eq(0),
            eq(-1),
            eq(1),
            eq(0),
            eq(0),
            eq(0),
            eq(-1),
            any(Date.class)
        );
    }

    /**
     * 验证无文档失败项绑定后按真实文件状态刷新批次计数。
     */
    @Test
    public void shouldBindRecoveredDocumentAndRefreshCounters() {
        DocumentImportBatchService batchService =
            mock(DocumentImportBatchService.class);
        DocumentImportBatchItemService itemService =
            mock(DocumentImportBatchItemService.class);
        DocumentImportBatchMapper batchMapper =
            mock(DocumentImportBatchMapper.class);
        DocumentImportBatchItemMapper itemMapper =
            mock(DocumentImportBatchItemMapper.class);
        DocumentImportBatch batch = batch(DocumentImportBatchStatus.RUNNING, 1);
        batch.setPendingCount(1);
        DocumentImportBatchItem item = new DocumentImportBatchItem();
        item.setId(BigInteger.TEN);
        item.setBatchId(batch.getId());
        item.setStatus(DocumentImportBatchItemStatus.FAILED.name());
        item.setRetryable(true);
        BigInteger documentId = BigInteger.valueOf(99);
        when(itemService.getById(item.getId())).thenReturn(item);
        when(itemMapper.bindFailedDocument(
            eq(item.getId()), eq(documentId), any(Date.class)
        )).thenReturn(1);
        when(batchService.getById(batch.getId())).thenReturn(batch);
        DocumentImportBatchTracker tracker =
            new DocumentImportBatchTracker(
                batchService, itemService, batchMapper, itemMapper);

        tracker.bindDocument(item.getId(), documentId);

        verify(batchMapper).refreshCountersFromItems(
            eq(batch.getId()), any(Date.class));
        verify(batchMapper, never()).adjustCounters(
            any(), anyInt(), anyInt(), anyInt(), anyInt(), anyInt(),
            anyInt(), anyInt(), anyInt(), any(Date.class));
    }

    /**
     * 验证迟到任务不能把已完成文件重新改为处理中。
     */
    @Test
    public void shouldRejectLateTransitionFromCompletedToRunning() {
        DocumentImportBatchService batchService = mock(DocumentImportBatchService.class);
        DocumentImportBatchItemService itemService = mock(DocumentImportBatchItemService.class);
        DocumentImportBatchMapper batchMapper = mock(DocumentImportBatchMapper.class);
        DocumentImportBatchItemMapper itemMapper = mock(DocumentImportBatchItemMapper.class);
        DocumentImportBatchItem item = new DocumentImportBatchItem();
        item.setId(BigInteger.TEN);
        item.setBatchId(BigInteger.ONE);
        item.setStatus(DocumentImportBatchItemStatus.COMPLETED.name());
        when(itemService.getById(item.getId())).thenReturn(item);

        DocumentImportBatchTracker tracker =
            new DocumentImportBatchTracker(batchService, itemService, batchMapper, itemMapper);

        assertFalse(tracker.updateItem(item.getId(),
            DocumentImportBatchItemStage.INDEX,
            DocumentImportBatchItemStatus.RUNNING,
            null));
        verify(itemMapper, never()).transitionStatus(
            any(), any(), any(), any(), any(), any(), anyBoolean(), anyInt(), any(Date.class));
    }

    private DocumentImportBatch batch(DocumentImportBatchStatus status, int totalCount) {
        DocumentImportBatch batch = new DocumentImportBatch();
        batch.setId(BigInteger.ONE);
        batch.setKnowledgeId(BigInteger.TWO);
        batch.setImportMode(DocumentImportMode.AUTO.name());
        batch.setStatus(status.name());
        batch.setTotalCount(totalCount);
        batch.setTotalBytes(100L);
        batch.setCompletedCount(0);
        batch.setProcessingCount(0);
        batch.setFailedCount(0);
        batch.setPendingCount(0);
        batch.setUploadedCount(0);
        batch.setSkippedCount(0);
        batch.setCancelledCount(0);
        batch.setRetryableFailedCount(0);
        return batch;
    }
}
