package com.yaoai.agent.service.impl;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.yaoai.agent.dto.AgentRunCommand;
import com.yaoai.agent.service.ScriptUnderstandingService;
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.AiTaskMapper;
import com.yaoai.domain.mapper.AgentRunPlanMapper;
import com.yaoai.domain.mapper.AssemblyTaskMapper;
import com.yaoai.domain.mapper.CharacterMapper;
import com.yaoai.domain.mapper.EpisodeMapper;
import com.yaoai.domain.mapper.GraphRunMapper;
import com.yaoai.domain.mapper.GraphStepMapper;
import com.yaoai.domain.mapper.OutlineMapper;
import com.yaoai.domain.mapper.SceneMapper;
import com.yaoai.domain.mapper.StoryboardMapper;
import com.yaoai.pipeline.service.AssemblyPipelineService;
import com.yaoai.pipeline.service.AssetGenPipelineService;
import com.yaoai.pipeline.service.OutlinePipelineService;
import com.yaoai.pipeline.service.StoryboardPipelineService;
import com.yaoai.pipeline.service.VideoTaskPipelineService;
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 AiTaskMapper aiTaskMapper;
    @Mock
    private AssemblyTaskMapper assemblyTaskMapper;
    @Mock
    private AgentRunPlanMapper agentRunPlanMapper;
    @Mock
    private OutlineMapper outlineMapper;
    @Mock
    private EpisodeMapper episodeMapper;
    @Mock
    private CharacterMapper characterMapper;
    @Mock
    private SceneMapper sceneMapper;
    @Mock
    private StoryboardMapper storyboardMapper;
    @Mock
    private AgentSseManager sseManager;
    @Mock
    private LlmService llmService;
    @Mock
    private ScriptUnderstandingService scriptUnderstandingService;
    @Mock
    private OutlinePipelineService outlinePipeline;
    @Mock
    private AssetGenPipelineService assetGenPipeline;
    @Mock
    private StoryboardPipelineService storyboardPipeline;
    @Mock
    private VideoTaskPipelineService videoTaskPipeline;
    @Mock
    private AssemblyPipelineService assemblyPipeline;

    @Test
    void quickVideoPlan_shouldAppendVideoPipelineStepsAfterStoryboardGeneration() {
        AgentRunServiceImpl service = new AgentRunServiceImpl(
                graphRunMapper,
                graphStepMapper,
                aiTaskMapper,
                assemblyTaskMapper,
                agentRunPlanMapper,
                outlineMapper,
                episodeMapper,
                characterMapper,
                sceneMapper,
                storyboardMapper,
                sseManager,
                llmService,
                new ObjectMapper(),
                scriptUnderstandingService,
                outlinePipeline,
                assetGenPipeline,
                storyboardPipeline,
                videoTaskPipeline,
                assemblyPipeline
        );
        AgentRunCommand command = runCommand("goal");
        command.setRunIntent("quick_video");

        List<?> steps = ReflectionTestUtils.invokeMethod(service, "planSteps", false, command);
        List<String> stepKeys = steps.stream()
                .map(step -> {
                    Object key = ReflectionTestUtils.invokeMethod(step, "key");
                    return String.valueOf(key);
                })
                .toList();

        assertTrue(stepKeys.indexOf("image_generation") > stepKeys.indexOf("supervision_2"));
        assertTrue(stepKeys.indexOf("video_generation") > stepKeys.indexOf("image_generation"));
        assertTrue(stepKeys.indexOf("video_composition") > stepKeys.indexOf("video_generation"));
    }

    @Test
    void storyboardGen_shouldPopulateMissingVideoPromptsAfterGeneratingStoryboards() {
        AgentRunServiceImpl service = new AgentRunServiceImpl(
                graphRunMapper,
                graphStepMapper,
                aiTaskMapper,
                assemblyTaskMapper,
                agentRunPlanMapper,
                outlineMapper,
                episodeMapper,
                characterMapper,
                sceneMapper,
                storyboardMapper,
                sseManager,
                llmService,
                new ObjectMapper(),
                scriptUnderstandingService,
                outlinePipeline,
                assetGenPipeline,
                storyboardPipeline,
                videoTaskPipeline,
                assemblyPipeline
        );

        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");
        when(scriptUnderstandingService.toGenerationContext(any(), any())).thenReturn("goal");

        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(characterMapper.findByProject(projectId, tenantId)).thenReturn(List.of());
        when(sceneMapper.findByProject(projectId, tenantId)).thenReturn(List.of());
        when(storyboardMapper.findByProject(projectId, tenantId)).thenReturn(List.of());
        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,
                runCommand("goal"),
                false
        );

        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;
    }

    private static AgentRunCommand runCommand(String goal) {
        AgentRunCommand command = new AgentRunCommand();
        command.setGoal(goal);
        command.setRunIntent("auto");
        command.setForceRegenerate(false);
        command.setSelectedSteps(List.of());
        command.setContextOptions(new AgentRunCommand.ContextOptions());
        return command;
    }
}
