import os
import re

import cv2

from util.ears_ai import EarsAI


def video_to_frames(video_path, output_dir):
    """動画からフレームを抽出して保存"""
    os.makedirs(output_dir, exist_ok=True)
    video = cv2.VideoCapture(video_path)
    if not video.isOpened():
        raise IOError(f"Could not open video file: {video_path}")

    frame_num = 0
    while True:
        success, frame = video.read()
        if not success:
            break
        frame_num += 1
        cv2.imwrite(os.path.join(output_dir, f"{frame_num}-frame.png"), frame)

    video.release()
    print(f"All frames saved to {output_dir}")
    return frame_num


def process_frames_with_posenet(frames_dir, output_dir):
    """PoseNetで骨格検出を行い、結果を保存"""
    ears_ai = EarsAI()
    os.makedirs(output_dir, exist_ok=True)

    # フレーム画像を番号順にソート
    png_files = sorted(
        [f for f in os.listdir(frames_dir) if f.lower().endswith(".png")],
        key=lambda x: int(re.search(r"(\d+)", x).group(1)),
    )

    for image_file_name in png_files:
        print(f"Processing image: {image_file_name}")
        # フレームを読み込み
        frame = cv2.imread(os.path.join(frames_dir, image_file_name))
        if frame is None:
            continue

        # PoseNetで骨格検出
        pose_overlay_img, *_ = ears_ai.pose_detect(frame, None)

        # 結果を保存
        cv2.imwrite(
            os.path.join(output_dir, image_file_name),
            cv2.cvtColor(pose_overlay_img, cv2.COLOR_RGB2BGR),
        )


def create_video_from_images(image_dir, output_path, fps=30):
    """画像シーケンスから動画を作成"""
    images = sorted(
        [img for img in os.listdir(image_dir) if img.endswith(".png")],
        key=lambda x: int(re.search(r"(\d+)", x).group()),
    )

    if not images:
        print(f"No images found in {image_dir}")
        return

    frame = cv2.imread(os.path.join(image_dir, images[0]))
    height, width, _ = frame.shape

    video = cv2.VideoWriter(
        output_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height)
    )

    for image in images:
        img = cv2.imread(os.path.join(image_dir, image))
        video.write(img)

    video.release()
    print(f"Created video: {output_path}")


def main():
    # 設定
    video_path = "./video/Test3-1.mp4"  # 入力動画のパス
    output_base_dir = "output-posenet"  # 出力ディレクトリ
    frames_dir = os.path.join(output_base_dir, "frames")  # フレーム保存ディレクトリ
    pose_frames_dir = os.path.join(
        output_base_dir, "pose_frames"
    )  # ポーズ検出結果保存ディレクトリ
    output_video_path = os.path.join(
        output_base_dir, "pose_detection.mp4"
    )  # 出力動画パス

    # 出力ベースディレクトリの作成
    os.makedirs(output_base_dir, exist_ok=True)

    # メイン処理
    print("Starting video processing...")
    print("1. Extracting frames from video...")
    video_to_frames(video_path, frames_dir)

    print("2. Processing frames with PoseNet...")
    process_frames_with_posenet(frames_dir, pose_frames_dir)

    print("3. Creating output video...")
    create_video_from_images(pose_frames_dir, output_video_path)

    print("Processing complete!")


if __name__ == "__main__":
    main()