src/modules/video_generator.py

"""
霜云(Shimokumo) - 视频制作模块

提供分镜脚本生成、角色动画控制、字幕生成（SRT格式）、
BGM音乐管理和视频渲染流程等功能。
"""

import os
import re
import time
import uuid
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

from utils.logger import get_logger

logger = get_logger("Shimokumo.VideoGenerator")


# ==================== 分镜脚本 ====================

@dataclass
class StoryboardFrame:
    """分镜帧数据类"""
    frame_number: int = 0
    """帧编号"""
    duration: float = 5.0
    """持续时长（秒）"""
    scene_description: str = ""
    """场景描述"""
    camera_movement: str = "static"
    """镜头运动：static/pan_left/pan_right/zoom_in/zoom_out/tracking/follow"""
    camera_angle: str = "eye_level"
    """镜头角度：eye_level/high_angle/low_angle/bird_eye/worm_eye/dutch"""
    character_actions: Dict[str, str] = field(default_factory=dict)
    """角色动作，key为角色名，value为动作描述"""
    dialogue: Dict[str, str] = field(default_factory=dict)
    """对白，key为角色名，value为台词"""
    narration: str = ""
    """旁白文本"""
    bgm_mood: str = ""
    """BGM情绪/风格"""
    sfx: List[str] = field(default_factory=list)
    """音效列表"""
    transition: str = "cut"
    """转场效果：cut/fade/dissolve/wipe/zoom"""
    subtitle_text: str = ""
    """字幕文本"""
    notes: str = ""
    """备注"""

    def to_prompt(self) -> str:
        """生成分镜提示词"""
        parts: List[str] = [f"【分镜 {self.frame_number}】（{self.duration:.1f}秒）"]
        parts.append(f"场景：{self.scene_description}")
        parts.append(f"镜头：{self.camera_angle} / {self.camera_movement}")
        if self.character_actions:
            parts.append("角色动作：")
            for char, action in self.character_actions.items():
                parts.append(f"  {char}: {action}")
        if self.dialogue:
            parts.append("对白：")
            for char, line in self.dialogue.items():
                parts.append(f"  {char}: 「{line}」")
        if self.narration:
            parts.append(f"旁白：{self.narration}")
        if self.bgm_mood:
            parts.append(f"BGM：{self.bgm_mood}")
        if self.sfx:
            parts.append(f"音效：{'、'.join(self.sfx)}")
        if self.transition != "cut":
            parts.append(f"转场：{self.transition}")
        return "\n".join(parts)

    def to_dict(self) -> Dict[str, Any]:
        """转为字典"""
        return {
            "frame_number": self.frame_number,
            "duration": self.duration,
            "scene_description": self.scene_description,
            "camera_movement": self.camera_movement,
            "camera_angle": self.camera_angle,
            "character_actions": self.character_actions,
            "dialogue": self.dialogue,
            "narration": self.narration,
            "bgm_mood": self.bgm_mood,
            "sfx": self.sfx,
            "transition": self.transition,
            "subtitle_text": self.subtitle_text,
        }


# ==================== 字幕 ====================

@dataclass
class SubtitleEntry:
    """字幕条目"""
    index: int = 0
    """序号"""
    start_time: float = 0.0
    """开始时间（秒）"""
    end_time: float = 0.0
    """结束时间（秒）"""
    text: str = ""
    """字幕文本"""
    style: str = ""
    """字幕样式"""

    def to_srt_time(self, seconds: float) -> str:
        """将秒数转为SRT时间格式 HH:MM:SS,mmm"""
        hours = int(seconds // 3600)
        minutes = int((seconds % 3600) // 60)
        secs = int(seconds % 60)
        millis = int((seconds % 1) * 1000)
        return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"

    def to_srt(self) -> str:
        """生成SRT格式字符串"""
        return (
            f"{self.index}\n"
            f"{self.to_srt_time(self.start_time)} --> {self.to_srt_time(self.end_time)}\n"
            f"{self.text}\n"
        )


# ==================== BGM管理 ====================

@dataclass
class BGMAudio:
    """BGM音频数据"""
    name: str = ""
    """音乐名称"""
    file_path: str = ""
    """文件路径或URL"""
    mood: str = ""
    """情绪标签"""
    duration: float = 0.0
    """时长"""
    tempo: int = 0
    """BPM"""
    volume: float = 0.5
    """音量 (0.0-1.0)"""
    loop: bool = True
    """是否循环"""
    fade_in: float = 1.0
    """淡入时长"""
    fade_out: float = 1.0
    """淡出时长"""


class BGMLibrary:
    """BGM音乐库"""

    def __init__(self):
        self.tracks: Dict[str, BGMAudio] = {}

    def add_track(self, bgm: BGMAudio) -> None:
        """添加BGM"""
        self.tracks[bgm.name] = bgm

    def get_by_mood(self, mood: str) -> List[BGMAudio]:
        """根据情绪标签获取BGM"""
        return [t for t in self.tracks.values() if mood.lower() in t.mood.lower()]

    def search(self, keyword: str) -> List[BGMAudio]:
        """搜索BGM"""
        keyword = keyword.lower()
        return [t for t in self.tracks.values() if keyword in t.name.lower() or keyword in t.mood.lower()]


# ==================== 视频项目 ====================

@dataclass
class VideoProject:
    """视频项目"""
    id: str = field(default_factory=lambda: str(uuid.uuid4())[:8])
    title: str = ""
    """视频标题"""
    resolution: Tuple[int, int] = (1920, 1080)
    """分辨率"""
    fps: int = 24
    """帧率"""
    frames: List[StoryboardFrame] = field(default_factory=list)
    """分镜列表"""
    subtitles: List[SubtitleEntry] = field(default_factory=list)
    """字幕列表"""
    bgm_library: BGMLibrary = field(default_factory=BGMLibrary)
    """BGM库"""
    total_duration: float = 0.0
    """总时长"""
    created_at: float = field(default_factory=time.time)

    def calculate_duration(self) -> float:
        """计算总时长"""
        self.total_duration = sum(f.duration for f in self.frames)
        return self.total_duration

    def add_frame(self, frame: StoryboardFrame) -> None:
        """添加分镜"""
        frame.frame_number = len(self.frames) + 1
        self.frames.append(frame)
        self.calculate_duration()


# ==================== 视频生成模块 ====================

class VideoGeneratorModule:
    """视频制作模块

    提供完整的视频制作工作流：
    1. 创建视频项目
    2. 根据剧本生成分镜脚本
    3. 生成SRT字幕
    4. BGM音乐管理
    5. 输出渲染指令/脚本

    用法：
        generator = VideoGeneratorModule(inference_engine)
        project = generator.create_project("我的动画", 1920, 1080, 24)
        frames = generator.generate_storyboard(project, script_text)
        subtitles = generator.generate_subtitles(project)
        srt_content = generator.export_srt(subtitles)
    """

    # 预定义镜头运动列表
    CAMERA_MOVEMENTS = [
        "static", "pan_left", "pan_right", "pan_up", "pan_down",
        "zoom_in", "zoom_out", "tracking", "follow", "orbit",
        "dolly_in", "dolly_out", "crane_up", "crane_down",
    ]

    # 预定义转场效果
    TRANSITIONS = [
        "cut", "fade", "dissolve", "wipe_left", "wipe_right",
        "wipe_up", "wipe_down", "zoom_in", "zoom_out",
        "iris_in", "iris_out", "slide_left", "slide_right",
    ]

    # 预定义BGM情绪
    BGM_MOODS = [
        "happy", "sad", "tense", "epic", "romantic", "mysterious",
        "calm", "energetic", "dark", "heroic", "nostalgic", "comical",
        "cheerful", "melancholy", "suspenseful", "peaceful",
    ]

    def __init__(self, inference_engine=None):
        """
        初始化视频制作模块。

        Args:
            inference_engine: 霜云推理引擎实例
        """
        self.inference = inference_engine

    def create_project(
        self,
        title: str = "",
        width: int = 1920,
        height: int = 1080,
        fps: int = 24,
    ) -> VideoProject:
        """
        创建视频项目。

        Args:
            title: 视频标题
            width: 视频宽度
            height: 视频高度
            fps: 帧率

        Returns:
            VideoProject实例
        """
        project = VideoProject(
            title=title,
            resolution=(width, height),
            fps=fps,
        )
        logger.info(f"创建视频项目: 「{title}」 ({width}x{height} @ {fps}fps)")
        return project

    def generate_storyboard(
        self,
        project: VideoProject,
        script: str,
        characters: Optional[List[str]] = None,
        style: str = "anime",
    ) -> List[StoryboardFrame]:
        """
        根据剧本生成分镜脚本。

        Args:
            project: 视频项目
            script: 剧本文本
            characters: 出场角色列表
            style: 风格（anime/realistic/cg）

        Returns:
            分镜帧列表
        """
        if self.inference:
            prompt = (
                f"请根据以下剧本，生成详细的分镜脚本。\n"
                f"风格：{style}\n"
            )
            if characters:
                prompt += f"出场角色：{'、'.join(characters)}\n"
            prompt += (
                f"\n剧本内容：\n{script}\n\n"
                f"请为每个分镜指定：\n"
                f"- 场景描述\n"
                f"- 镜头运动（{', '.join(self.CAMERA_MOVEMENTS)}）\n"
                f"- 镜头角度\n"
                f"- 角色动作\n"
                f"- 对白/旁白\n"
                f"- BGM情绪\n"
                f"- 转场效果\n"
                f"- 持续时长（秒）\n\n"
                f"格式：用「分镜N」标记每个分镜。"
            )

            response = self.inference.generate(prompt, max_new_tokens=3000, temperature=0.8)
            frames = self._parse_storyboard_response(response)
        else:
            frames = self._generate_template_storyboard(script, characters)

        # 添加到项目
        for frame in frames:
            project.add_frame(frame)

        logger.info(f"生成 {len(frames)} 个分镜")
        return frames

    def _generate_template_storyboard(
        self,
        script: str,
        characters: Optional[List[str]] = None,
    ) -> List[StoryboardFrame]:
        """生成模板分镜（回退方案）"""
        frames: List[StoryboardFrame] = []

        # 将剧本按段落分割
        paragraphs = [p.strip() for p in script.split("\n") if p.strip()]
        chars = characters or []

        for i, para in enumerate(paragraphs, 1):
            frame = StoryboardFrame(
                frame_number=i,
                duration=5.0,
                scene_description=para,
                camera_movement="static" if i <= 1 else (
                    "pan_right" if i % 3 == 0 else "zoom_in"
                ),
                camera_angle="eye_level",
                transition="cut" if i == 1 else "dissolve",
                bgm_mood="calm",
            )

            # 分配角色动作
            if chars:
                char = chars[i % len(chars)]
                frame.character_actions[char] = "站立，面朝前方"
                frame.dialogue[char] = para if len(para) < 100 else para[:100]

            # 提取字幕
            frame.subtitle_text = para[:80] if len(para) > 80 else para

            frames.append(frame)

        return frames

    def _parse_storyboard_response(self, response: str) -> List[StoryboardFrame]:
        """解析AI生成的分镜响应"""
        frames: List[StoryboardFrame] = []
        blocks = re.split(r"[【\[]分镜\s*(\d+)[】\]]", response)

        # blocks格式：[前缀, 编号1, 内容1, 编号2, 内容2, ...]
        for i in range(1, len(blocks) - 1, 2):
            try:
                num = int(blocks[i])
                content = blocks[i + 1].strip()
            except (ValueError, IndexError):
                continue

            frame = StoryboardFrame(
                frame_number=num,
                scene_description=content[:500],
                subtitle_text=content[:100],
            )

            # 尝试提取时长
            duration_match = re.search(r"(\d+\.?\d*)\s*秒", content)
            if duration_match:
                frame.duration = float(duration_match.group(1))

            # 尝试提取镜头运动
            for movement in self.CAMERA_MOVEMENTS:
                if movement in content.lower():
                    frame.camera_movement = movement
                    break

            frames.append(frame)

        return frames

    def generate_subtitles(self, project: VideoProject) -> List[SubtitleEntry]:
        """
        从分镜生成SRT字幕列表。

        Args:
            project: 视频项目

        Returns:
            字幕条目列表
        """
        subtitles: List[SubtitleEntry] = []
        current_time = 0.0

        for frame in project.frames:
            if not frame.subtitle_text:
                # 尝试从对白和旁白中提取
                texts: List[str] = []
                for char, line in frame.dialogue.items():
                    texts.append(f"{char}: {line}")
                if frame.narration:
                    texts.append(frame.narration)

                if texts:
                    frame.subtitle_text = " / ".join(texts)

            if frame.subtitle_text:
                entry = SubtitleEntry(
                    index=len(subtitles) + 1,
                    start_time=current_time,
                    end_time=current_time + frame.duration,
                    text=frame.subtitle_text,
                )
                subtitles.append(entry)

            current_time += frame.duration

        project.subtitles = subtitles
        logger.info(f"生成 {len(subtitles)} 条字幕")
        return subtitles

    def export_srt(self, subtitles: List[SubtitleEntry]) -> str:
        """
        导出SRT格式字幕文件内容。

        Args:
            subtitles: 字幕条目列表

        Returns:
            SRT格式文本
        """
        return "\n".join(entry.to_srt() for entry in subtitles)

    def save_srt(self, subtitles: List[SubtitleEntry], file_path: str) -> bool:
        """
        保存SRT字幕文件。

        Args:
            subtitles: 字幕条目列表
            file_path: 文件保存路径

        Returns:
            是否成功
        """
        try:
            os.makedirs(os.path.dirname(file_path) or ".", exist_ok=True)
            srt_content = self.export_srt(subtitles)
            with open(file_path, "w", encoding="utf-8") as f:
                f.write(srt_content)
            logger.info(f"SRT字幕已保存: {file_path}")
            return True
        except Exception as e:
            logger.error(f"保存SRT文件失败: {e}")
            return False

    def generate_render_script(
        self,
        project: VideoProject,
        output_path: str = "output.mp4",
    ) -> str:
        """
        生成FFmpeg渲染脚本。

        Args:
            project: 视频项目
            output_path: 输出文件路径

        Returns:
            FFmpeg命令字符串
        """
        w, h = project.resolution
        fps = project.fps
        duration = project.calculate_duration()

        script_parts: List[str] = [
            "#!/bin/bash",
            "# 霜云视频渲染脚本",
            f"# 项目: {project.title}",
            f"# 分辨率: {w}x{h} @ {fps}fps",
            f"# 总时长: {duration:.1f}秒",
            f"# 分镜数: {len(project.frames)}",
            "",
            f"OUTPUT=\"{output_path}\"",
            f"RESOLUTION=\"{w}x{h}\"",
            f"FPS={fps}",
            "",
            "# 检查FFmpeg",
            "if ! command -v ffmpeg &> /dev/null; then",
            "    echo \"错误: 未找到FFmpeg，请先安装\"",
            "    exit 1",
            "fi",
            "",
        ]

        # 生成分镜帧列表
        script_parts.append("# 分镜帧列表")
        for frame in project.frames:
            script_parts.append(
                f"# 分镜{frame.frame_number}: {frame.scene_description[:50]} "
                f"({frame.duration:.1f}s)"
            )

        script_parts.extend([
            "",
            "# FFmpeg渲染命令模板",
            "# 需要根据实际的图片/视频素材替换输入文件",
            f"ffmpeg -y \\",
            f"  -r $FPS \\",
            f"  -s $RESOLUTION \\",
            f"  -i \"frames/%04d.png\" \\",
        ])

        # 添加字幕（如果有）
        if project.subtitles:
            srt_path = os.path.splitext(output_path)[0] + ".srt"
            script_parts.append(f"  -vf \"subtitles={srt_path}\" \\")

        # 添加BGM（如果有）
        if project.bgm_library.tracks:
            bgm = list(project.bgm_library.tracks.values())[0]
            script_parts.append(f"  -i \"{bgm.file_path}\" \\")
            script_parts.append(f"  -map 0:v -map 1:a \\")

        script_parts.extend([
            f"  -c:v libx264 -preset medium -crf 23 \\",
            f"  -c:a aac -b:a 128k \\",
            f"  -shortest \\",
            f"  \"$OUTPUT\"",
            "",
            'echo "渲染完成: $OUTPUT"',
        ])

        return "\n".join(script_parts)

    def get_project_summary(self, project: VideoProject) -> str:
        """
        获取视频项目摘要信息。

        Args:
            project: 视频项目

        Returns:
            格式化的摘要文本
        """
        w, h = project.resolution
        total_duration = project.calculate_duration()

        parts: List[str] = [
            f"视频项目「{project.title}」",
            f"分辨率: {w}x{h} | 帧率: {fps}fps" if (fps := project.fps) else "",
            f"总时长: {total_duration:.1f}秒 | 分镜数: {len(project.frames)}",
            f"字幕数: {len(project.subtitles)}",
            "",
        ]

        if project.frames:
            parts.append("分镜列表：")
            for frame in project.frames:
                parts.append(
                    f"  {frame.frame_number}. [{frame.duration:.1f}s] "
                    f"{frame.scene_description[:40]}... "
                    f"({frame.camera_movement}/{frame.transition})"
                )

        return "\n".join(parts)