import cv2
import json
import os
import sys
from .pose import PoseEstimator

def preprocess_video(video_path, output_json_path):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"Error: Could not open video {video_path}")
        return

    estimator = PoseEstimator(model_complexity=2) # Higher complexity for reference
    all_landmarks = []
    
    frame_count = 0
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        results = estimator.process_frame(frame)
        landmarks = estimator.extract_landmarks(results)
        all_landmarks.append({
            "frame": frame_count,
            "landmarks": landmarks
        })
        
        frame_count += 1
        if frame_count % 30 == 0:
            print(f"Processed {frame_count} frames...")

    with open(output_json_path, 'w') as f:
        json.dump(all_landmarks, f)
    
    cap.release()
    estimator.close()
    print(f"Finished! Saved to {output_json_path}")

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python -m src.core.preprocess_reference <video_path> <output_json_path>")
    else:
        preprocess_video(sys.argv[1], sys.argv[2])