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

def process_videos(zip_file):
    """ZIPファイル内の動画を2つの解像度にリサイズする"""
    if zip_file is None:
        return None

    # 一時ディレクトリを作成
    with tempfile.TemporaryDirectory() as temp_dir:
        # ZIPファイルを一時ディレクトリに解凍
        with zipfile.ZipFile(zip_file, '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))
                    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

                    # 出力ビデオライターの設定
                    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 Interface
app = gr.Interface(
    fn=process_videos,
    inputs=gr.File(label="ZIPファイルをアップロード", file_types=[".zip"]),
    outputs=gr.File(label="処理済みZIPファイル"),
    title="🎬 動画リサイズツール",
    description="""
## 使い方
1. フォルダ内の動画ファイル(.mp4)をZIPにまとめてアップロード
2. 「Submit」ボタンをクリック
3. 処理完了後、ZIPファイルをダウンロード

## 出力される解像度
- **標準解像度**: 1080x1920
- **高解像度**: 2160x3840 (upscaledフォルダ内)

## 注意事項
- 元のフォルダ構造は維持されます
- 処理には時間がかかる場合があります
    """,
    flagging_mode="never"
)

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