package com.yaoai.agent.service.impl;

import com.yaoai.agent.sse.AgentSseManager;
import com.yaoai.ai.core.service.LlmService;
import com.yaoai.domain.entity.Episode;
import com.yaoai.domain.entity.GraphRun;
import com.yaoai.domain.entity.GraphStep;
import com.yaoai.domain.entity.Outline;
import com.yaoai.domain.entity.Storyboard;
import com.yaoai.domain.mapper.EpisodeMapper;
import com.yaoai.domain.mapper.GraphRunMapper;
import com.yaoai.domain.mapper.GraphStepMapper;
import com.yaoai.pipeline.service.AssetGenPipelineService;
import com.yaoai.pipeline.service.OutlinePipelineService;
import com.yaoai.pipeline.service.StoryboardPipelineService;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.ArgumentCaptor;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.test.util.ReflectionTestUtils;

import java.util.List;

import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.atLeastOnce;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

@ExtendWith(MockitoExtension.class)
class AgentRunServiceImplTest {

    @Mock
    private GraphRunMapper graphRunMapper;
    @Mock
    private GraphStepMapper graphStepMapper;
    @Mock
    private EpisodeMapper episodeMapper;
    @Mock
    private AgentSseManager sseManager;
    @Mock
    private LlmService llmService;
    @Mock
    private OutlinePipelineService outlinePipeline;
    @Mock
    private AssetGenPipelineService assetGenPipeline;
    @Mock
    private StoryboardPipelineService storyboardPipeline;

    @Test
    void storyboardGen_shouldPopulateMissingVideoPromptsAfterGeneratingStoryboards() {
        AgentRunServiceImpl service = new AgentRunServiceImpl(
                graphRunMapper,
                graphStepMapper,
                episodeMapper,
                sseManager,
                llmService,
                outlinePipeline,
                assetGenPipeline,
                storyboardPipeline
        );

        Long runId = 88L;
        Long projectId = 22L;
        Long tenantId = 33L;
        GraphRun runningRun = new GraphRun();
        runningRun.setId(runId);
        runningRun.setStatus("RUNNING");
        when(graphRunMapper.selectById(runId)).thenReturn(runningRun);
        when(graphRunMapper.updateById(any(GraphRun.class))).thenReturn(1);

        List<GraphStep> steps = List.of(
                step(runId, "decision", 0),
                step(runId, "outline_gen", 1),
                step(runId, "episode_gen", 2),
                step(runId, "supervision_1", 3),
                step(runId, "asset_gen", 4),
                step(runId, "storyboard_gen", 5),
                step(runId, "supervision_2", 6)
        );
        when(graphStepMapper.findByRunId(runId)).thenReturn(steps);
        when(graphStepMapper.updateById(any(GraphStep.class))).thenReturn(1);
        when(llmService.chat(any())).thenReturn("SCORE: A");

        Outline outline = new Outline();
        outline.setId(77L);
        outline.setTitle("outline");
        outline.setEpisodeCount(1);
        when(outlinePipeline.generateOutline(projectId, tenantId, "goal")).thenReturn(outline);

        Episode episode = new Episode();
        episode.setId(101L);
        episode.setEpisodeNumber(1);
        when(outlinePipeline.generateEpisodes(77L, projectId, tenantId)).thenReturn(List.of(episode));
        when(assetGenPipeline.extractCharacters(projectId, tenantId)).thenReturn(List.of());
        when(assetGenPipeline.extractScenes(projectId, tenantId)).thenReturn(List.of());
        when(storyboardPipeline.generateStoryboards(101L, projectId, tenantId))
                .thenReturn(List.of(new Storyboard(), new Storyboard()));
        when(storyboardPipeline.populateMissingVideoPromptsByEpisode(101L, tenantId))
                .thenReturn(List.of(storyboardWithPrompt(), storyboardWithPrompt()));

        ReflectionTestUtils.invokeMethod(
                service,
                "executeRun",
                runId,
                projectId,
                tenantId,
                1001L,
                "goal"
        );

        verify(storyboardPipeline).populateMissingVideoPromptsByEpisode(101L, tenantId);

        ArgumentCaptor<GraphStep> stepCaptor = ArgumentCaptor.forClass(GraphStep.class);
        verify(graphStepMapper, atLeastOnce()).updateById(stepCaptor.capture());
        boolean storyboardOutputMentionedPrompts = stepCaptor.getAllValues().stream()
                .anyMatch(step -> "storyboard_gen".equals(step.getStepKey())
                        && "DONE".equals(step.getStatus())
                        && step.getOutput() != null
                        && step.getOutput().contains("视频提示词"));
        assertTrue(storyboardOutputMentionedPrompts);
    }

    private static GraphStep step(Long runId, String stepKey, int seq) {
        GraphStep step = new GraphStep();
        step.setRunId(runId);
        step.setStepKey(stepKey);
        step.setStatus("PENDING");
        step.setSeq(seq);
        return step;
    }

    private static Storyboard storyboardWithPrompt() {
        Storyboard storyboard = new Storyboard();
        storyboard.setVideoPrompt("prompt");
        return storyboard;
    }
}
