package com.ruoyi.generator.service.front;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import java.lang.reflect.Field;
import java.util.Arrays;
import java.util.Collections;
import org.junit.Before;
import org.junit.Test;
import org.mockito.ArgumentCaptor;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
import com.alibaba.fastjson2.JSON;
import com.ruoyi.generator.domain.front.AiGenerationTask;
import com.ruoyi.generator.domain.front.AiUsageLedger;
import com.ruoyi.generator.domain.front.dto.AiGenerationTaskCreateRequest;
import com.ruoyi.generator.domain.front.dto.AppBlueprintDesign;
import com.ruoyi.generator.domain.front.dto.GenerateAppBlueprintRequest;
import com.ruoyi.generator.mapper.front.AiGenerationTaskMapper;
import com.ruoyi.generator.mapper.front.AiUsageLedgerMapper;

public class AiGenerationTaskWorkerTest
{
    private AiGenerationTaskWorker worker;

    @Mock
    private AiGenerationTaskMapper aiGenerationTaskMapper;
    @Mock
    private AiUsageLedgerMapper aiUsageLedgerMapper;
    @Mock
    private IAiGenerateService aiGenerateService;
    @Mock
    private AiQuotaService aiQuotaService;

    @Before
    public void setUp() throws Exception
    {
        MockitoAnnotations.initMocks(this);
        worker = new AiGenerationTaskWorker();
        setField("aiGenerationTaskMapper", aiGenerationTaskMapper);
        setField("aiUsageLedgerMapper", aiUsageLedgerMapper);
        setField("aiGenerateService", aiGenerateService);
        setField("aiQuotaService", aiQuotaService);
        setField("retryPolicy", new AiTaskRetryPolicy());
        setField("costCalculator", new AiCostCalculator());
    }

    @Test
    public void processTaskMarksSuccessAndStoresResult()
    {
        AiGenerationTask task = task("app_blueprint");
        when(aiGenerationTaskMapper.selectAiGenerationTaskById(88L)).thenReturn(task);
        when(aiGenerationTaskMapper.claimTask(org.mockito.ArgumentMatchers.eq(88L), any(String.class))).thenReturn(1);
        AppBlueprintDesign blueprint = new AppBlueprintDesign();
        blueprint.setRoles(Collections.emptyList());
        when(aiGenerateService.generateAppBlueprint(any(Long.class), any(Long.class), any(GenerateAppBlueprintRequest.class)))
                .thenReturn(blueprint);

        worker.processTask(88L);

        ArgumentCaptor<AiGenerationTask> taskCaptor = ArgumentCaptor.forClass(AiGenerationTask.class);
        verify(aiGenerationTaskMapper, org.mockito.Mockito.atLeastOnce()).updateAiGenerationTask(taskCaptor.capture());
        AiGenerationTask saved = taskCaptor.getAllValues().get(taskCaptor.getAllValues().size() - 1);
        assertEquals("SUCCEEDED", saved.getStatus());
        assertEquals(Integer.valueOf(100), saved.getProgress());
        assertTrue(saved.getTotalTokens() > 0);
        assertTrue(saved.getCostCents() > 0);
        assertTrue(saved.getResultPayload().contains("\"roles\":[]"));
        verify(aiQuotaService).releaseRunning(7L);
        verify(aiUsageLedgerMapper).insertAiUsageLedger(any(AiUsageLedger.class));
    }

    @Test
    public void processTaskSchedulesRetryForRetryableFailure()
    {
        AiGenerationTask task = task("app_blueprint");
        task.setAttempts(1);
        when(aiGenerationTaskMapper.selectAiGenerationTaskById(88L)).thenReturn(task);
        when(aiGenerationTaskMapper.claimTask(org.mockito.ArgumentMatchers.eq(88L), any(String.class))).thenReturn(1);
        when(aiGenerateService.generateAppBlueprint(any(Long.class), any(Long.class), any(GenerateAppBlueprintRequest.class)))
                .thenThrow(new RuntimeException("DeepSeek request timeout"));

        worker.processTask(88L);

        ArgumentCaptor<AiGenerationTask> taskCaptor = ArgumentCaptor.forClass(AiGenerationTask.class);
        verify(aiGenerationTaskMapper, org.mockito.Mockito.atLeastOnce()).updateAiGenerationTask(taskCaptor.capture());
        AiGenerationTask saved = taskCaptor.getAllValues().get(taskCaptor.getAllValues().size() - 1);
        assertEquals("RETRY_WAITING", saved.getStatus());
        assertEquals(Integer.valueOf(2), saved.getAttempts());
        verify(aiQuotaService).releaseRunning(7L);
    }

    @Test
    public void processTaskMarksFailedForNonRetryableFailure()
    {
        AiGenerationTask task = task("app_blueprint");
        when(aiGenerationTaskMapper.selectAiGenerationTaskById(88L)).thenReturn(task);
        when(aiGenerationTaskMapper.claimTask(org.mockito.ArgumentMatchers.eq(88L), any(String.class))).thenReturn(1);
        when(aiGenerateService.generateAppBlueprint(any(Long.class), any(Long.class), any(GenerateAppBlueprintRequest.class)))
                .thenThrow(new RuntimeException("配额不足"));

        worker.processTask(88L);

        ArgumentCaptor<AiGenerationTask> taskCaptor = ArgumentCaptor.forClass(AiGenerationTask.class);
        verify(aiGenerationTaskMapper, org.mockito.Mockito.atLeastOnce()).updateAiGenerationTask(taskCaptor.capture());
        AiGenerationTask saved = taskCaptor.getAllValues().get(taskCaptor.getAllValues().size() - 1);
        assertEquals("FAILED", saved.getStatus());
        assertEquals(Integer.valueOf(1), saved.getAttempts());
        verify(aiQuotaService).releaseRunning(7L);
    }

    @Test
    public void processDueTasksRunsClaimableTasks()
    {
        AiGenerationTask first = task("app_blueprint");
        first.setTaskId(88L);
        AiGenerationTask second = task("app_blueprint");
        second.setTaskId(89L);
        when(aiGenerationTaskMapper.selectClaimableTasks(2)).thenReturn(Arrays.asList(first, second));
        when(aiGenerationTaskMapper.selectAiGenerationTaskById(88L)).thenReturn(first);
        when(aiGenerationTaskMapper.selectAiGenerationTaskById(89L)).thenReturn(second);
        when(aiGenerationTaskMapper.claimTask(any(Long.class), any(String.class))).thenReturn(1);
        AppBlueprintDesign blueprint = new AppBlueprintDesign();
        blueprint.setRoles(Collections.emptyList());
        when(aiGenerateService.generateAppBlueprint(any(Long.class), any(Long.class), any(GenerateAppBlueprintRequest.class)))
                .thenReturn(blueprint);

        int count = worker.processDueTasks(2);

        assertEquals(2, count);
        verify(aiGenerationTaskMapper).selectClaimableTasks(2);
        verify(aiGenerationTaskMapper).selectAiGenerationTaskById(88L);
        verify(aiGenerationTaskMapper).selectAiGenerationTaskById(89L);
    }

    private AiGenerationTask task(String generateType)
    {
        AiGenerationTaskCreateRequest request = new AiGenerationTaskCreateRequest();
        request.setGenerateType(generateType);
        request.setProjectName("客户中心");
        request.setProjectDesc("管理客户资料");
        request.setIndustryTemplate("crm");
        request.setExtraRequirements("客户分层");
        AiGenerationTask task = new AiGenerationTask();
        task.setTaskId(88L);
        task.setProjectId(10L);
        task.setUserId(7L);
        task.setGenerateType(generateType);
        task.setProvider("deepseek");
        task.setModel("deepseek-chat");
        task.setStatus("QUEUED");
        task.setAttempts(0);
        task.setMaxAttempts(3);
        task.setRequestPayload(JSON.toJSONString(request));
        return task;
    }

    private void setField(String name, Object value) throws Exception
    {
        Field field = AiGenerationTaskWorker.class.getDeclaredField(name);
        field.setAccessible(true);
        field.set(worker, value);
    }
}
