package com.yaoai.pipeline.service.impl;

import com.baomidou.mybatisplus.core.conditions.update.UpdateWrapper;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.yaoai.ai.core.model.ChatMessage;
import com.yaoai.ai.core.model.ChatRequest;
import com.yaoai.ai.core.service.LlmService;
import com.yaoai.billing.dto.BillingChargeRequest;
import com.yaoai.billing.service.BillingService;
import com.yaoai.common.context.UserContext;
import com.yaoai.common.exception.BizException;
import com.yaoai.common.exception.ErrorCode;
import com.yaoai.domain.entity.Character;
import com.yaoai.domain.entity.Episode;
import com.yaoai.domain.entity.Scene;
import com.yaoai.domain.entity.Storyboard;
import com.yaoai.domain.mapper.CharacterMapper;
import com.yaoai.domain.mapper.EpisodeMapper;
import com.yaoai.domain.mapper.SceneMapper;
import com.yaoai.domain.mapper.StoryboardMapper;
import com.yaoai.pipeline.service.StoryboardPipelineService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.stream.Collectors;

@Slf4j
@Service
@RequiredArgsConstructor
public class StoryboardPipelineServiceImpl implements StoryboardPipelineService {
    private static final int STORYBOARD_PROMPT_CONCURRENCY = 3;


    private static final String SYSTEM_PROMPT_BASE = """
            You are a professional storyboard artist.
            Split the episode script into a sequence of continuous shots, where each shot lasts 4-8 seconds.
            %s

            Return only a JSON array. Each item must follow this schema:
            [
              {
                "sequence_num": 1,
                "scene_number": "001",
                "short_description": "One-sentence shot description in Chinese",
                "detailed_description": "Detailed scene description in Chinese",
                "characters": "@character or @scene refs separated by commas, e.g. @Alice,@Cafe",
                "dialogues": "Dialogue content, empty string when absent",
                "camera_direction": "Camera movement, e.g. static/push/pull/pan/tracking",
                "composition_guide": "Composition guidance, e.g. close-up/medium/wide",
                "duration_seconds": 5,
                "start_frame_prompt": "English still-frame prompt, under 30 words",
                "motion_script": "Short Chinese motion description for video generation"
              }
            ]
            Do not include any explanation outside the JSON array.
            """;

    private static final String PROMPT_GEN_SYSTEM = """
            You are a professional storyboard prompt writer.
            Based on episode context, character appearance, scene setup, and the current storyboard shot,
            turn the shot into a concise multi-segment video-generation prompt.

            Requirements:
            - Stay faithful to the provided plot and the visual setup.
            - Respect character appearance, costume, and scene details.
            - Use 2-4 time segments whose total duration matches the storyboard duration.
            - Each segment should include timing, visuals, dialogue when present, and audio mood.
            - Return only the prompt text with no extra explanation.
            """;

    private final LlmService llmService;
    private final StoryboardMapper storyboardMapper;
    private final EpisodeMapper episodeMapper;
    private final CharacterMapper characterMapper;
    private final SceneMapper sceneMapper;
    private final ObjectMapper objectMapper;
    private final BillingService billingService;

    @Override
    public List<Storyboard> generateStoryboards(Long episodeId, Long projectId, Long tenantId) {
        Episode episode = episodeMapper.selectById(episodeId);
        if (episode == null) {
            throw new BizException(ErrorCode.NOT_FOUND, "Episode not found");
        }

        String assetContext = buildAssetContext(projectId, tenantId);
        String systemPrompt = String.format(SYSTEM_PROMPT_BASE, assetContext);
        String userMsg = String.format(
                "Episode title: %s%n%nEpisode script:%n%s",
                nullSafe(episode.getTitle()),
                episode.getScript() != null ? episode.getScript() : nullSafe(episode.getSummary())
        );

        log.info("Generating storyboards: episodeId={}, projectId={}", episodeId, projectId);

        ChatRequest request = ChatRequest.builder()
                .messages(List.of(
                        ChatMessage.system(systemPrompt),
                        ChatMessage.user(userMsg)
                ))
                .build();

        billingService.checkBalance(BillingChargeRequest.builder()
                .tenantId(tenantId)
                .userId(UserContext.require())
                .projectId(projectId)
                .operation("storyboard_generate")
                .modality("TEXT")
                .modelProvider(llmService.getProvider())
                .modelId(llmService.getModelId())
                .modelName(llmService.getModelName())
                .unitCount(1)
                .refId(String.valueOf(episodeId))
                .build());

        String raw = llmService.chat(request);
        log.debug("Storyboard LLM response length={}", raw.length());

        try {
            String json = extractJson(raw);
            List<Map<String, Object>> items = objectMapper.readValue(json, new TypeReference<>() {});
            Set<String> sceneRefs = projectSceneRefs(projectId, tenantId);

            List<Storyboard> result = new ArrayList<>();
            for (Map<String, Object> item : items) {
                StoryboardRefs refs = splitStoryboardRefs((String) item.getOrDefault("characters", ""), sceneRefs);

                Storyboard storyboard = new Storyboard();
                storyboard.setProjectId(projectId);
                storyboard.setEpisodeId(episodeId);
                storyboard.setTenantId(tenantId);
                storyboard.setSequenceNum(((Number) item.getOrDefault("sequence_num", result.size() + 1)).intValue());
                storyboard.setSceneNumber((String) item.getOrDefault("scene_number", String.format("%03d", result.size() + 1)));
                storyboard.setShortDescription((String) item.getOrDefault("short_description", ""));
                storyboard.setDetailedDescription((String) item.getOrDefault("detailed_description", ""));
                storyboard.setCharacters(refs.characters());
                storyboard.setSceneRef(refs.sceneRef());
                storyboard.setDialogues((String) item.getOrDefault("dialogues", ""));
                storyboard.setCameraDirection((String) item.getOrDefault("camera_direction", ""));
                storyboard.setCompositionGuide((String) item.getOrDefault("composition_guide", ""));
                storyboard.setDurationSeconds(((Number) item.getOrDefault("duration_seconds", 5)).intValue());
                storyboard.setStartFramePrompt((String) item.getOrDefault("start_frame_prompt", ""));
                storyboard.setMotionScript((String) item.getOrDefault("motion_script", ""));
                storyboard.setStatus("draft");
                storyboardMapper.insert(storyboard);
                result.add(storyboard);
            }

            billingService.charge(BillingChargeRequest.builder()
                    .tenantId(tenantId)
                    .userId(UserContext.require())
                    .projectId(projectId)
                    .operation("storyboard_generate")
                    .modality("TEXT")
                    .modelProvider(llmService.getProvider())
                    .modelId(llmService.getModelId())
                    .modelName(llmService.getModelName())
                    .unitCount(result.size())
                    .refId(String.valueOf(episodeId))
                    .build());
            log.info("Storyboards generated: count={}", result.size());
            return result;
        } catch (Exception e) {
            log.error("Storyboard parsing failed", e);
            throw new BizException(ErrorCode.INTERNAL_ERROR, "Storyboard parsing failed: " + e.getMessage());
        }
    }

    @Override
    public List<Storyboard> listByEpisode(Long episodeId, Long tenantId) {
        return storyboardMapper.findByEpisode(episodeId, tenantId);
    }

    @Override
    public List<Storyboard> listByProject(Long projectId, Long tenantId) {
        return storyboardMapper.findByProject(projectId, tenantId);
    }

    @Override
    public Storyboard create(Storyboard storyboard) {
        int maxSeq = storyboardMapper.maxSequenceByEpisode(storyboard.getEpisodeId());
        storyboard.setSequenceNum(maxSeq + 1);
        storyboard.setStatus("draft");
        storyboardMapper.insert(storyboard);
        return storyboard;
    }

    @Override
    @Transactional
    public Storyboard update(Long id, Long tenantId, Storyboard patch) {
        Storyboard existing = storyboardMapper.selectById(id);
        if (existing == null || !tenantId.equals(existing.getTenantId())) {
            throw new BizException(ErrorCode.NOT_FOUND, "Storyboard not found");
        }

        boolean bodyFieldsChanged = hasStoryboardBodyChanges(existing, patch);
        patch.setId(id);
        patch.setTenantId(tenantId);
        int updatedRows = storyboardMapper.updateById(patch);
        if (updatedRows != 1) {
            throw new BizException(ErrorCode.INTERNAL_ERROR, "Failed to update storyboard");
        }
        if (bodyFieldsChanged) {
            clearVideoPrompt(id, tenantId);
        }
        return storyboardMapper.selectById(id);
    }

    @Override
    public void delete(Long id, Long tenantId) {
        Storyboard existing = storyboardMapper.selectById(id);
        if (existing == null || !tenantId.equals(existing.getTenantId())) {
            throw new BizException(ErrorCode.NOT_FOUND, "Storyboard not found");
        }
        storyboardMapper.deleteById(id);
    }

    @Override
    public void reorder(List<Long> orderedIds, Long tenantId) {
        for (int i = 0; i < orderedIds.size(); i++) {
            storyboardMapper.updateSequence(orderedIds.get(i), i + 1, tenantId);
        }
    }

    @Override
    @Transactional
    public String generatePrompt(Long storyboardId, Long tenantId) {
        Storyboard storyboard = storyboardMapper.selectById(storyboardId);
        if (storyboard == null || !tenantId.equals(storyboard.getTenantId())) {
            throw new BizException(ErrorCode.NOT_FOUND, "Storyboard not found");
        }

        StringBuilder userMsg = new StringBuilder();
        if (storyboard.getEpisodeId() != null) {
            Episode episode = episodeMapper.selectById(storyboard.getEpisodeId());
            if (episode != null) {
                String episodeContent = episode.getScript() != null && !episode.getScript().isBlank()
                        ? episode.getScript()
                        : nullSafe(episode.getSummary());
                userMsg.append(String.format(
                        "Episode context:%nEpisode %d - %s%n%s%n%n",
                        episode.getEpisodeNumber() != null ? episode.getEpisodeNumber() : 0,
                        nullSafe(episode.getTitle()),
                        episodeContent
                ));
            }
        }

        List<Character> characters = characterMapper.findByProject(storyboard.getProjectId(), tenantId);
        if (!characters.isEmpty()) {
            userMsg.append("Character visual setup:\n");
            for (Character character : characters) {
                userMsg.append(String.format(
                        "- @%s (%s): personality=%s, costume=%s, visual=%s%n",
                        nullSafe(character.getName()),
                        nullSafe(character.getRoleType()),
                        nullSafe(character.getPersonality()),
                        nullSafe(character.getCostume()),
                        nullSafe(character.getVisualHint())
                ));
            }
            userMsg.append("\n");
        }

        List<Scene> scenes = sceneMapper.findByProject(storyboard.getProjectId(), tenantId);
        if (!scenes.isEmpty()) {
            userMsg.append("Scene visual setup:\n");
            for (Scene scene : scenes) {
                userMsg.append(String.format(
                        "- @%s (%s): %s%n",
                        nullSafe(scene.getName()),
                        sceneTypeLabel(scene.getSceneType()),
                        nullSafe(scene.getDescription())
                ));
            }
            userMsg.append("\n");
        }

        StoryboardRefs promptRefs = promptRefsForCurrentShot(storyboard, scenes);
        userMsg.append(String.format(
                "Current storyboard shot:%n" +
                        "- Sequence: %d%n" +
                        "- Short description: %s%n" +
                        "- Detailed description: %s%n" +
                        "- Character refs: %s%n" +
                        "- Scene refs: %s%n" +
                        "- Dialogues: %s%n" +
                        "- Camera direction: %s%n" +
                        "- Composition guide: %s%n" +
                        "- Duration seconds: %d%n",
                storyboard.getSequenceNum() != null ? storyboard.getSequenceNum() : 0,
                nullSafe(storyboard.getShortDescription()),
                nullSafe(storyboard.getDetailedDescription()),
                nullSafe(promptRefs.characters()),
                nullSafe(promptRefs.sceneRef()),
                nullSafe(storyboard.getDialogues()),
                nullSafe(storyboard.getCameraDirection()),
                nullSafe(storyboard.getCompositionGuide()),
                storyboard.getDurationSeconds() != null ? storyboard.getDurationSeconds() : 5
        ));

        ChatRequest request = ChatRequest.builder()
                .messages(List.of(
                        ChatMessage.system(PROMPT_GEN_SYSTEM),
                        ChatMessage.user(userMsg.toString())
                ))
                .build();

        log.info(
                "Generating storyboard prompt: storyboardId={}, episodeId={}, characters={}, scenes={}",
                storyboardId,
                storyboard.getEpisodeId(),
                characters.size(),
                scenes.size()
        );

        billingService.checkBalance(BillingChargeRequest.builder()
                .tenantId(tenantId)
                .userId(UserContext.require())
                .projectId(storyboard.getProjectId())
                .operation("storyboard_prompt_generate")
                .modality("TEXT")
                .modelProvider(llmService.getProvider())
                .modelId(llmService.getModelId())
                .modelName(llmService.getModelName())
                .unitCount(1)
                .refId(String.valueOf(storyboardId))
                .build());

        String prompt = llmService.chat(request);

        billingService.charge(BillingChargeRequest.builder()
                .tenantId(tenantId)
                .userId(UserContext.require())
                .projectId(storyboard.getProjectId())
                .operation("storyboard_prompt_generate")
                .modality("TEXT")
                .modelProvider(llmService.getProvider())
                .modelId(llmService.getModelId())
                .modelName(llmService.getModelName())
                .unitCount(1)
                .refId(String.valueOf(storyboardId))
                .build());

        Storyboard promptUpdate = new Storyboard();
        promptUpdate.setId(storyboard.getId());
        promptUpdate.setVideoPrompt(prompt);
        int updatedRows = storyboardMapper.updateById(promptUpdate);
        if (updatedRows != 1) {
            throw new BizException(ErrorCode.INTERNAL_ERROR, "Failed to persist storyboard prompt");
        }
        return prompt;
    }

    @Override
    public List<Storyboard> populateMissingVideoPromptsByEpisode(Long episodeId, Long tenantId) {
        List<Storyboard> pending = storyboardMapper.findMissingVideoPromptByEpisode(episodeId, tenantId);
        if (pending.isEmpty()) {
            return storyboardMapper.findByEpisode(episodeId, tenantId);
        }

        Long userId = UserContext.get();
        int poolSize = Math.min(STORYBOARD_PROMPT_CONCURRENCY, pending.size());
        ExecutorService executor = Executors.newFixedThreadPool(poolSize);
        try {
            List<CompletableFuture<Void>> futures = pending.stream()
                    .map(storyboard -> CompletableFuture.runAsync(() -> {
                        if (userId != null) {
                            UserContext.set(userId);
                        }
                        try {
                            generatePrompt(storyboard.getId(), tenantId);
                        } finally {
                            UserContext.clear();
                        }
                    }, executor))
                    .toList();

            for (CompletableFuture<Void> future : futures) {
                future.join();
            }
        } catch (CompletionException e) {
            if (e.getCause() instanceof RuntimeException runtimeException) {
                throw runtimeException;
            }
            throw e;
        } finally {
            executor.shutdown();
        }
        return storyboardMapper.findByEpisode(episodeId, tenantId);
    }

    private String buildAssetContext(Long projectId, Long tenantId) {
        StringBuilder context = new StringBuilder();
        List<Character> characters = characterMapper.findByProject(projectId, tenantId);
        List<Scene> scenes = sceneMapper.findByProject(projectId, tenantId);

        if (characters.isEmpty() && scenes.isEmpty()) {
            return "No project character or scene setup exists yet. Infer reasonably from the script.";
        }

        if (!characters.isEmpty()) {
            context.append("Project characters (refer to them as @Name in the characters field):\n");
            for (Character character : characters) {
                context.append(String.format(
                        "- @%s: personality=%s, visual=%s, costume=%s%n",
                        nullSafe(character.getName()),
                        nullSafe(character.getPersonality()),
                        nullSafe(character.getVisualHint()),
                        nullSafe(character.getCostume())
                ));
            }
        }

        if (!scenes.isEmpty()) {
            context.append("\nProject scenes (they may also appear as @Name refs in the characters field):\n");
            for (Scene scene : scenes) {
                context.append(String.format(
                        "- @%s (%s): %s%n",
                        nullSafe(scene.getName()),
                        sceneTypeLabel(scene.getSceneType()),
                        nullSafe(scene.getDescription())
                ));
            }
        }

        context.append("\nThe characters field must use @Name refs, separated by commas.");
        return context.toString();
    }

    private Set<String> projectSceneRefs(Long projectId, Long tenantId) {
        return sceneMapper.findByProject(projectId, tenantId)
                .stream()
                .map(Scene::getName)
                .filter(Objects::nonNull)
                .map(String::trim)
                .filter(name -> !name.isEmpty())
                .map(name -> "@" + name)
                .collect(Collectors.toCollection(LinkedHashSet::new));
    }

    private StoryboardRefs splitStoryboardRefs(String rawRefs, Set<String> sceneRefs) {
        List<String> refs = parseRefs(rawRefs);
        List<String> characterRefs = new ArrayList<>();
        String sceneRef = null;

        for (String ref : refs) {
            if (sceneRefs.contains(ref)) {
                if (sceneRef == null) {
                    sceneRef = ref;
                }
            } else {
                characterRefs.add(ref);
            }
        }

        return new StoryboardRefs(joinRefs(characterRefs), sceneRef);
    }

    private StoryboardRefs promptRefsForCurrentShot(Storyboard storyboard, List<Scene> projectScenes) {
        Set<String> sceneRefs = projectScenes.stream()
                .map(Scene::getName)
                .filter(Objects::nonNull)
                .map(String::trim)
                .filter(name -> !name.isEmpty())
                .map(name -> "@" + name)
                .collect(Collectors.toCollection(LinkedHashSet::new));
        StoryboardRefs fallbackRefs = splitStoryboardRefs(storyboard.getCharacters(), sceneRefs);
        String sceneRef = storyboard.getSceneRef() != null && !storyboard.getSceneRef().isBlank()
                ? storyboard.getSceneRef()
                : fallbackRefs.sceneRef();
        return new StoryboardRefs(fallbackRefs.characters(), sceneRef);
    }

    private List<String> parseRefs(String rawRefs) {
        if (rawRefs == null || rawRefs.isBlank()) {
            return List.of();
        }
        return Arrays.stream(rawRefs.split(","))
                .map(String::trim)
                .filter(token -> !token.isBlank())
                .distinct()
                .toList();
    }

    private String joinRefs(List<String> refs) {
        if (refs.isEmpty()) {
            return null;
        }
        return String.join(",", refs);
    }

    private void clearVideoPrompt(Long storyboardId, Long tenantId) {
        int updatedRows = storyboardMapper.update(
                null,
                new UpdateWrapper<Storyboard>()
                        .eq("id", storyboardId)
                        .eq("tenant_id", tenantId)
                        .set("video_prompt", null)
        );
        if (updatedRows != 1) {
            throw new BizException(ErrorCode.INTERNAL_ERROR, "Failed to clear storyboard prompt");
        }
    }

    private boolean hasStoryboardBodyChanges(Storyboard existing, Storyboard patch) {
        return hasChanged(patch.getSceneNumber(), existing.getSceneNumber())
                || hasChanged(patch.getShortDescription(), existing.getShortDescription())
                || hasChanged(patch.getDetailedDescription(), existing.getDetailedDescription())
                || hasChanged(patch.getCharacters(), existing.getCharacters())
                || hasChanged(patch.getSceneRef(), existing.getSceneRef())
                || hasChanged(patch.getDialogues(), existing.getDialogues())
                || hasChanged(patch.getCameraDirection(), existing.getCameraDirection())
                || hasChanged(patch.getCompositionGuide(), existing.getCompositionGuide())
                || hasChanged(patch.getDurationSeconds(), existing.getDurationSeconds())
                || hasChanged(patch.getStartFramePrompt(), existing.getStartFramePrompt())
                || hasChanged(patch.getEndFramePrompt(), existing.getEndFramePrompt())
                || hasChanged(patch.getMotionScript(), existing.getMotionScript())
                || hasChanged(patch.getNotes(), existing.getNotes());
    }

    private boolean hasChanged(String patchValue, String existingValue) {
        return patchValue != null && !Objects.equals(patchValue, existingValue);
    }

    private boolean hasChanged(Integer patchValue, Integer existingValue) {
        return patchValue != null && !Objects.equals(patchValue, existingValue);
    }

    private String extractJson(String raw) {
        String text = raw.strip();
        if (text.startsWith("```")) {
            int start = text.indexOf('\n') + 1;
            int end = text.lastIndexOf("```");
            if (end > start) {
                text = text.substring(start, end).strip();
            }
        }
        return text;
    }

    private static String nullSafe(String value) {
        return value != null ? value : "";
    }

    private String sceneTypeLabel(String sceneType) {
        if ("indoor".equals(sceneType)) {
            return "indoor";
        }
        if ("outdoor".equals(sceneType)) {
            return "outdoor";
        }
        return "";
    }

    private record StoryboardRefs(String characters, String sceneRef) {
    }
}
