package com.yaoai.pipeline.service.impl;

import cn.dev33.satoken.stp.StpUtil;
import com.yaoai.billing.dto.BillingChargeRequest;
import com.yaoai.billing.service.BillingService;
import com.yaoai.common.context.UserContext;
import com.yaoai.domain.entity.AiTask;
import com.yaoai.domain.entity.AssemblyTask;
import com.yaoai.domain.entity.ShotAsset;
import com.yaoai.domain.entity.Storyboard;
import com.yaoai.domain.mapper.AiTaskMapper;
import com.yaoai.domain.mapper.AssemblyTaskMapper;
import com.yaoai.domain.mapper.ShotAssetMapper;
import com.yaoai.domain.mapper.StoryboardMapper;
import com.yaoai.media.service.FfmpegService;
import com.yaoai.pipeline.service.AssemblyPipelineService;
import com.yaoai.storage.service.TosService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.stream.Collectors;

@Slf4j
@Service
public class AssemblyPipelineServiceImpl implements AssemblyPipelineService {

    private final AssemblyTaskMapper assemblyTaskMapper;
    private final StoryboardMapper storyboardMapper;
    private final ShotAssetMapper shotAssetMapper;
    private final AiTaskMapper aiTaskMapper;
    private final FfmpegService ffmpegService;
    private final TosService tosService;
    private final BillingService billingService;
    private final Executor assemblyExecutor;

    public AssemblyPipelineServiceImpl(
            AssemblyTaskMapper assemblyTaskMapper,
            StoryboardMapper storyboardMapper,
            ShotAssetMapper shotAssetMapper,
            AiTaskMapper aiTaskMapper,
            FfmpegService ffmpegService,
            TosService tosService,
            BillingService billingService,
            @Qualifier("assemblyExecutor") Executor assemblyExecutor) {
        this.assemblyTaskMapper = assemblyTaskMapper;
        this.storyboardMapper = storyboardMapper;
        this.shotAssetMapper = shotAssetMapper;
        this.aiTaskMapper = aiTaskMapper;
        this.ffmpegService = ffmpegService;
        this.tosService = tosService;
        this.billingService = billingService;
        this.assemblyExecutor = assemblyExecutor;
    }

    @Override
    public AssemblyTask startAssembly(Long projectId, Long episodeId, Long tenantId, List<Long> videoTaskIds) {
        Long userId = resolveUserId();
        AssemblyTask task = new AssemblyTask();
        task.setTenantId(tenantId);
        task.setProjectId(projectId);
        task.setEpisodeId(episodeId);
        task.setStatus("pending");
        task.setShotCount(0);
        assemblyTaskMapper.insert(task);

        Long taskId = task.getId();
        CompletableFuture.runAsync(() -> {
            UserContext.set(userId);
            try {
                doAssemble(taskId, episodeId, userId, videoTaskIds);
            } finally {
                UserContext.clear();
            }
        }, assemblyExecutor);
        return task;
    }

    @Override
    public AssemblyTask getLatest(Long episodeId, Long tenantId) {
        return assemblyTaskMapper.findLatestByEpisode(episodeId, tenantId);
    }

    @Override
    public String signedDownloadUrl(AssemblyTask task) {
        if (task == null || task.getResultTosKey() == null) return null;
        return tosService.publicUrl(task.getResultTosKey());
    }

    private void doAssemble(Long assemblyTaskId, Long episodeId, Long userId, List<Long> videoTaskIds) {
        AssemblyTask task = assemblyTaskMapper.selectById(assemblyTaskId);
        task.setStatus("running");
        assemblyTaskMapper.updateById(task);

        Path outputFile = null;
        try {
            List<String> urls = videoTaskIds != null && !videoTaskIds.isEmpty()
                    ? buildAssemblyUrlsFromTaskIds(videoTaskIds, episodeId, task.getTenantId())
                    : buildAssemblyUrls(episodeId, task.getTenantId());
            if (urls.isEmpty()) {
                fail(task, "没有可用的视频片段（需先生成每个镜头的视频）");
                return;
            }

            task.setShotCount(urls.size());
            assemblyTaskMapper.updateById(task);

            outputFile = ffmpegService.concatVideos(urls);

            byte[] bytes = Files.readAllBytes(outputFile);
            String tosKey = "assembly/" + task.getTenantId() + "/" + episodeId + "/" + assemblyTaskId + ".mp4";
            try (var is = new java.io.ByteArrayInputStream(bytes)) {
                tosService.upload(tosKey, is, bytes.length, "video/mp4");
            }

            task.setResultTosKey(tosKey);
            task.setStatus("succeeded");
            assemblyTaskMapper.updateById(task);
            billingService.charge(BillingChargeRequest.builder()
                    .tenantId(task.getTenantId())
                    .userId(userId)
                    .projectId(task.getProjectId())
                    .operation("assembly")
                    .modality("COMPOSE")
                    .modelProvider("ffmpeg")
                    .modelId("ffmpeg-compose")
                    .modelName("FFmpeg 视频合成")
                    .unitCount(1)
                    .refId(String.valueOf(assemblyTaskId))
                    .build());
            log.info("Assembly task {} succeeded, tosKey={}", assemblyTaskId, tosKey);

        } catch (IOException | InterruptedException e) {
            log.error("Assembly task {} failed", assemblyTaskId, e);
            fail(task, e.getMessage());
        } finally {
            if (outputFile != null) ffmpegService.cleanup(outputFile);
        }
    }

    private List<String> buildAssemblyUrlsFromTaskIds(List<Long> videoTaskIds, Long episodeId, Long tenantId) {
        List<AiTask> tasks = aiTaskMapper.selectBatchIds(videoTaskIds);
        Map<Long, AiTask> taskById = tasks.stream()
                .collect(Collectors.toMap(AiTask::getId, task -> task, (left, right) -> left, LinkedHashMap::new));
        List<String> urls = new ArrayList<>();
        for (Long taskId : videoTaskIds) {
            AiTask task = taskById.get(taskId);
            if (task == null
                    || !tenantId.equals(task.getTenantId())
                    || !episodeId.equals(task.getEpisodeId())
                    || !"succeeded".equals(task.getStatus())
                    || task.getResultVideoUrl() == null
                    || task.getResultVideoUrl().isBlank()) {
                continue;
            }
            urls.add(task.getResultVideoUrl());
        }
        if (urls.size() != videoTaskIds.size()) {
            throw new IllegalStateException("合成视频任务数量不一致：请求 " + videoTaskIds.size() + " 个，实际可用 " + urls.size() + " 个");
        }
        return urls;
    }

    private List<String> buildAssemblyUrls(Long episodeId, Long tenantId) {
        List<Storyboard> storyboards = storyboardMapper.findByEpisode(episodeId, tenantId);
        List<String> urls = new ArrayList<>();
        for (Storyboard storyboard : storyboards) {
            AiTask latestTask = aiTaskMapper.findLatestSucceededVideoByStoryboard(storyboard.getId(), tenantId);
            if (latestTask != null && latestTask.getResultVideoUrl() != null && !latestTask.getResultVideoUrl().isBlank()) {
                urls.add(latestTask.getResultVideoUrl());
                continue;
            }

            ShotAsset activeVideo = shotAssetMapper.findActiveReadyVideoByStoryboard(storyboard.getId(), tenantId);
            if (activeVideo != null && activeVideo.getExternalUrl() != null && !activeVideo.getExternalUrl().isBlank()) {
                urls.add(activeVideo.getExternalUrl());
            }
        }
        return urls;
    }

    private void fail(AssemblyTask task, String msg) {
        task.setStatus("failed");
        task.setErrorMessage(msg);
        assemblyTaskMapper.updateById(task);
    }

    private Long resolveUserId() {
        Long userId = UserContext.get();
        if (userId != null) {
            return userId;
        }
        return StpUtil.getLoginIdAsLong();
    }
}
