import os
import cv2
import av
import logging
import json

import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision

logging.basicConfig(format='%(asctime)s - %(levelname)s : %(message)s', level=logging.DEBUG)
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

DEFAULT_MODEL_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'face_landmarker.task'))

def process_video(video_path, model_path=DEFAULT_MODEL_PATH, output_json=None, method='av'):
    if method == 'cv2':
        return process_video_cv2(video_path, model_path, output_json)
    else:
        return process_video_av(video_path, model_path, output_json)
    
def process_video_cv2(video_path, model_path, output_json=None):
    """Core MediaPipe processing logic using OpenCV"""
    # Initialize MediaPipe
    cache = {}
    base_options = python.BaseOptions(model_asset_path=model_path)

    options = vision.FaceLandmarkerOptions(
        base_options=base_options,
        output_face_blendshapes=True,
        running_mode=vision.RunningMode.VIDEO
    )

    with vision.FaceLandmarker.create_from_options(options) as detector:
        cap = cv2.VideoCapture(video_path)
        fps = cap.get(cv2.CAP_PROP_FPS)
        frame_count = 0
        
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret: break
            
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb_frame)
            ts = int((frame_count / fps) * 1000)
            result = detector.detect_for_video(mp_image, ts)

            if result.face_blendshapes:
                # Capture first face from blendshapes
                shapes = {c.category_name: c.score for c in result.face_blendshapes[0]}
                cache[frame_count] = shapes
            
            frame_count += 1
        cap.release()

    return cache

def process_video_av(video_path, model_path, output_json=None):
    # 這裡是處理影片的邏輯
    """Core MediaPipe processing logic using PyAV"""
    cache = {}

    if not os.path.exists(video_path) or not os.path.exists(model_path):
        logger.error(f"File not found: {video_path} or {model_path}")
        return

    logger.info('Start Analyzing Video: {video_path}')

    # Initialize MediaPipe
    base_options = python.BaseOptions(model_asset_path=model_path)
    options = vision.FaceLandmarkerOptions(
        base_options=base_options,
        output_face_blendshapes=True,
        running_mode=vision.RunningMode.VIDEO
    )

    with vision.FaceLandmarker.create_from_options(options) as detector:
        with av.open(video_path) as container:
            video_stream = container.streams.video[0]
            video_stream.thread_type = "AUTO"
            frame_count = 0
            for frame in container.decode(video_stream):
                rgb_array = frame.to_ndarray(format='rgb24')
                mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb_array)
                # ts = int((frame_count / fps) * 1000)
                ts_ms = int(frame.time * 1000)
                
                result = detector.detect_for_video(mp_image, ts_ms)
                
                if result.face_blendshapes:
                    # 抓取第一張臉的 blendshapes
                    shapes = {c.category_name: c.score for c in result.face_blendshapes[0]}
                    cache[frame.index] = shapes
            
            frame_count += 1
        
    if output_json:
        if output_json.endswith('.json'):
            output_json_dir = os.path.dirname(output_json)
            if not os.path.exists(output_json_dir):
                os.makedirs(output_json_dir)

            with open(output_json, 'w') as f:
                json.dump(cache, f, indent=4)
            logger.info(f"Output JSON saved to {output_json}")

    return cache

if __name__ == "__main__":
    # 測試用主程式
    test_video_path = "test_videos/sample_video.mp4"
    test_model_path = "models/face_landmarker.task"
    test_output_json = "output/blendshapes_output.json"

    process_video(test_video_path, test_model_path, test_output_json)