import gradio as gr
import subprocess
import os
import tempfile
import shutil
from video_merger import merge_videos

def extract_last_frame(video_file):
    if video_file is None:
        return None, "請上傳影片檔案"
    
    try:
        # 創建固定名稱的輸出檔案
        temp_output = tempfile.mktemp(suffix='.jpg')
        
        # 獲取影片總幀數
        frame_count_cmd = [
            'ffprobe',
            '-v', 'error',
            '-select_streams', 'v:0',
            '-count_frames',
            '-show_entries', 'stream=nb_read_frames',
            '-of', 'default=nokey=1:noprint_wrappers=1',
            video_file
        ]
        
        frame_result = subprocess.run(frame_count_cmd, capture_output=True, text=True)
        
        # 如果無法獲取幀數，使用反向讀取方法
        if not frame_result.stdout.strip() or frame_result.returncode != 0:
            # 方法2: 使用 select 過濾器選擇最後一幀
            extract_cmd = [
                'ffmpeg',
                '-i', video_file,
                '-vf', 'select=eq(n\\,0)',
                '-vsync', '0',
                '-q:v', '2',
                '-frames:v', '1',
                temp_output,
                '-y'
            ]
            
            # 先嘗試從尾部讀取
            extract_cmd_reverse = [
                'ffmpeg',
                '-sseof', '-3',
                '-i', video_file,
                '-vframes', '1',
                '-vf', 'select=eq(pict_type\\,I)',
                '-vsync', '0',
                '-q:v', '2',
                temp_output,
                '-y'
            ]
            
            result = subprocess.run(extract_cmd_reverse, capture_output=True, text=True)
            
            # 如果失敗，使用標準方法
            if result.returncode != 0 or not os.path.exists(temp_output):
                subprocess.run(extract_cmd, capture_output=True, text=True)
        else:
            # 方法1: 使用幀數選擇最後一幀
            total_frames = int(frame_result.stdout.strip())
            last_frame = total_frames - 1
            
            extract_cmd = [
                'ffmpeg',
                '-i', video_file,
                '-vf', f'select=eq(n\\,{last_frame})',
                '-vsync', '0',
                '-q:v', '2',
                '-frames:v', '1',
                temp_output,
                '-y'
            ]
            
            subprocess.run(extract_cmd, check=True, capture_output=True)
        
        if os.path.exists(temp_output) and os.path.getsize(temp_output) > 0:
            # 重新命名為 lastframe.jpg
            output_dir = tempfile.gettempdir()
            final_output = os.path.join(output_dir, 'lastframe.jpg')
            shutil.copy2(temp_output, final_output)
            os.remove(temp_output)
            return final_output, "成功提取最後一幀！"
        else:
            return None, "提取失敗，請確認影片格式正確"
            
    except subprocess.CalledProcessError as e:
        return None, f"處理錯誤: {str(e)}"
    except Exception as e:
        return None, f"發生錯誤: {str(e)}"

def merge_videos_wrapper(video_files):
    """包裝合併函數以處理檔名"""
    if not video_files:
        return None, "請上傳影片檔案"
    
    temp_output, message = merge_videos(video_files)
    
    if temp_output and os.path.exists(temp_output):
        # 重新命名為 final.mp4
        output_dir = tempfile.gettempdir()
        final_output = os.path.join(output_dir, 'final.mp4')
        shutil.copy2(temp_output, final_output)
        os.remove(temp_output)
        return final_output, message
    else:
        return None, message

# 創建 Gradio 界面
with gr.Blocks(title="Grok影片工具箱") as demo:
    gr.Markdown("# Grok影片工具箱")
    gr.Markdown("提供影片最後一幀提取和多段影片合併功能")
    
    with gr.Tabs():
        # Tab 1: 提取最後一幀
        with gr.Tab("提取最後一幀"):
            gr.Markdown("上傳影片（最大 100MB），自動提取最後一幀為 JPEG 圖片")
            with gr.Row():
                with gr.Column():
                    video_input = gr.Video(
                        label="上傳影片",
                        max_length=None,
                        height=400
                    )
                    extract_btn = gr.Button("提取最後一幀", variant="primary")
                
                with gr.Column():
                    image_output = gr.Image(
                        label="最後一幀 (lastframe.jpg)",
                        type="filepath",
                        height=400
                    )
                    status_output = gr.Textbox(label="狀態", interactive=False)
            
            extract_btn.click(
                fn=extract_last_frame,
                inputs=[video_input],
                outputs=[image_output, status_output]
            )
        
        # Tab 2: 合併影片
        with gr.Tab("合併影片"):
            gr.Markdown("上傳多個 MP4 影片，按照上傳順序合併為一個影片")
            with gr.Row():
                with gr.Column():
                    videos_input = gr.File(
                        label="上傳影片檔案（可多選）",
                        file_count="multiple",
                        file_types=["video"],
                        height=300
                    )
                    merge_btn = gr.Button("合併影片", variant="primary")
                
                with gr.Column():
                    merged_video_output = gr.Video(
                        label="合併後的影片 (final.mp4)",
                        height=400
                    )
                    merge_status_output = gr.Textbox(label="狀態", interactive=False)
            
            merge_btn.click(
                fn=merge_videos_wrapper,
                inputs=[videos_input],
                outputs=[merged_video_output, merge_status_output]
            )

if __name__ == "__main__":
    demo.launch()