import gradio as gr
import os
import zipfile
import tempfile
import shutil
from pathlib import Path
import cv2
from tqdm import tqdm

def process_videos(zip_file):
    # 一時ディレクトリを作成
    with tempfile.TemporaryDirectory() as temp_dir:
        # ZIPファイルを一時ディレクトリに解凍
        with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
            zip_ref.extractall(temp_dir)
        
        # 出力用の一時ディレクトリを作成
        output_temp_dir = tempfile.mkdtemp()
        output_upscaled_dir = tempfile.mkdtemp()
        
        # 処理したファイルの数を追跡
        processed_count = 0
        
        # 全てのmp4ファイルを処理
        for root, dirs, files in os.walk(temp_dir):
            for file in files:
                if file.lower().endswith('.mp4'):
                    input_path = os.path.join(root, file)
                    
                    # 出力パスの作成（元のディレクトリ構造を維持）
                    relative_path = os.path.relpath(os.path.join(root, file), temp_dir)
                    output_path = os.path.join(output_temp_dir, relative_path)
                    output_upscaled_path = os.path.join(output_upscaled_dir, f"upscaled_{relative_path}")
                    
                    # 出力ディレクトリの作成
                    os.makedirs(os.path.dirname(output_path), exist_ok=True)
                    os.makedirs(os.path.dirname(output_upscaled_path), exist_ok=True)
                    
                    # ビデオの読み込み
                    cap = cv2.VideoCapture(input_path)
                    
                    # 入力ビデオのプロパティを取得
                    fps = int(cap.get(cv2.CAP_PROP_FPS))
                    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
                    
                    # 出力ビデオライターの設定
                    out_normal = cv2.VideoWriter(output_path, fourcc, fps, (1080, 1920))
                    out_upscaled = cv2.VideoWriter(output_upscaled_path, fourcc, fps, (2160, 3840))
                    
                    while cap.isOpened():
                        ret, frame = cap.read()
                        if not ret:
                            break
                        
                        # 通常のリサイズ (1080x1920)
                        resized_frame = cv2.resize(frame, (1080, 1920))
                        out_normal.write(resized_frame)
                        
                        # 高解像度アップスケール (2160x3840)
                        upscaled_frame = cv2.resize(resized_frame, (2160, 3840), 
                                                  interpolation=cv2.INTER_CUBIC)
                        out_upscaled.write(upscaled_frame)
                    
                    cap.release()
                    out_normal.release()
                    out_upscaled.release()
                    processed_count += 1
        
        if processed_count == 0:
            return None
        
        # 処理済みファイルをZIPにまとめる
        output_zip = tempfile.NamedTemporaryFile(delete=False, suffix='.zip')
        with zipfile.ZipFile(output_zip.name, 'w', zipfile.ZIP_DEFLATED) as zf:
            # 通常サイズの動画を追加
            for root, dirs, files in os.walk(output_temp_dir):
                for file in files:
                    file_path = os.path.join(root, file)
                    arcname = os.path.relpath(file_path, output_temp_dir)
                    zf.write(file_path, arcname)
            
            # アップスケールされた動画を追加
            for root, dirs, files in os.walk(output_upscaled_dir):
                for file in files:
                    file_path = os.path.join(root, file)
                    arcname = os.path.relpath(file_path, output_upscaled_dir)
                    zf.write(file_path, f"upscaled/{arcname}")
        
        # 一時ディレクトリの削除
        shutil.rmtree(output_temp_dir)
        shutil.rmtree(output_upscaled_dir)
        
        return output_zip.name

# Gradioインターフェースの作成
with gr.Blocks() as app:
    gr.Markdown("## 動画リサイズツール")
    gr.Markdown("1. フォルダ内の動画ファイルをZIPにまとめてアップロードしてください")
    gr.Markdown("2. 全ての動画が以下の2つの解像度で出力されます:")
    gr.Markdown("   - 標準解像度: 1080x1920")
    gr.Markdown("   - 高解像度: 2160x3840")
    
    with gr.Row():
        input_file = gr.File(label="ZIPファイルをアップロード", file_types=[".zip"])
        output_file = gr.File(label="処理済みZIPファイル")
    
    upload_button = gr.Button("処理開始")
    upload_button.click(fn=process_videos, inputs=[input_file], outputs=[output_file])
    
    gr.Markdown("### 注意事項")
    gr.Markdown("- アップロードするZIPファイルには.mp4ファイルが含まれている必要があります")
    gr.Markdown("- 元のフォルダ構造は維持されます")
    gr.Markdown("- 処理には時間がかかる場合があります")

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