# Dependencies, see also requirement.txt ;)
import gradio as gr
import cv2 
import numpy as np
import os
from datetime import datetime  # 新增时间模块
from scenedetect import open_video, SceneManager
from scenedetect.detectors import ContentDetector
from moviepy.editor import VideoFileClip

def convert_to_tuple(list):
    return tuple(list)

def clear_app():
    return None, 27, None, None, None

def find_scenes(video_path, threshold):
    # 创建时间戳目录
    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    output_dir = f"output_{timestamp}"
    os.makedirs(output_dir, exist_ok=True)

    filename = os.path.splitext(os.path.basename(video_path))[0]
    video = open_video(video_path)
    scene_manager = SceneManager()
    scene_manager.add_detector(ContentDetector(threshold=threshold))
    
    scene_manager.detect_scenes(video, show_progress=True)
    scene_list = scene_manager.get_scene_list()
    
    data_outputs = []  # 改为局部变量
    gradio_components_outputs = []  # 改为局部变量
    
    data_outputs.append(scene_list)
    gradio_components_outputs.append("json")
    
    timecodes = []
    if not scene_list:
        gr.Warning("No scenes detected in this video")
        return None, None, None

    timecodes.append({"title": filename + ".mp4", "fps": scene_list[0][0].get_framerate()})
    
    shots = []
    stills = []
    
    for i, shot in enumerate(scene_list):
        framerate = shot[0].get_framerate()
        shot_in = shot[0].get_frames() / framerate
        shot_out = shot[1].get_frames() / framerate
        
        tc_in = shot[0].get_timecode()
        tc_out = shot[1].get_timecode()
        frame_in = shot[0].get_frames()
        frame_out = shot[1].get_frames()
        
        timecode = {"tc_in": tc_in, "tc_out": tc_out, "frame_in": frame_in, "frame_out": frame_out}
        timecodes.append(timecode)
        
        # 修改输出路径到时间戳目录
        target_name = os.path.join(output_dir, f"shot_{i+1}_{filename}.mp4")
        
        with VideoFileClip(video_path) as clip:
            subclip = clip.subclip(shot_in, shot_out)
            subclip.write_videofile(
                target_name,
                codec="libx264",
                audio_codec="aac",
                threads=4,
                preset="fast",
                ffmpeg_params=["-crf", "23"]
            )
        
        shots.append(target_name)
        data_outputs.append(target_name)
        gradio_components_outputs.append("video")

        vid = cv2.VideoCapture(video_path)
        frame_id = shot[0].get_frames()
        vid.set(cv2.CAP_PROP_POS_FRAMES, frame_id)
        ret, frame = vid.read()

        # 修改截图保存路径
        img = os.path.join(output_dir, f"{frame_id}_screenshot.png")
        cv2.imwrite(img, frame)
        stills.append((img, f'shot {i+1}'))
        vid.release()
    
    data_outputs.append(shots)
    gradio_components_outputs.append("file")
    data_outputs.append(stills)
    gradio_components_outputs.append("gallery")
    
    results = convert_to_tuple(data_outputs)
    print(f"所有输出文件保存在：{os.path.abspath(output_dir)}")

    return timecodes, shots, stills

with gr.Blocks() as demo:
    with gr.Column():
        gr.Markdown("""
            # Scene Edit Detection
            Automatically find all the shots in a video.
            Accepts mp4 format. Works only with videos that have cuts in them.
            """)
                
        with gr.Row():
            with gr.Column():
                video_input = gr.Video(sources="upload", format="mp4", label="视频输入", mirror_webcam=False)
                threshold = gr.Slider(label="场景切换检测阈值（越低越敏感）", minimum=15.0, maximum=40.0, value=27.0)
                with gr.Row():
                    clear_button = gr.Button("清除")
                    run_button = gr.Button("开始处理", variant="primary")

            with gr.Column():
                json_output = gr.JSON(label="场景分析结果")
                file_output = gr.File(label="下载分割片段")
                gallery_output = gr.Gallery(label="场景缩略图", object_fit="cover", columns=3)

        run_button.click(fn=find_scenes, inputs=[video_input, threshold], outputs=[json_output, file_output, gallery_output])
        clear_button.click(fn=clear_app, inputs=None, outputs=[video_input, threshold, json_output, file_output, gallery_output])

        # 添加示例部分
        gr.Examples(
            examples=[
                ["anime_kiss.mp4", 27],
            ],
            inputs=[video_input, threshold],
            outputs=[json_output, file_output, gallery_output],
            fn=find_scenes,
            cache_examples=False,
            label="示例视频"
        )

demo.queue().launch(debug=True, share=True)