#!/opt/homebrew/bin/python3.11
"""
MediaPipe Holistic Processor v1.11 — Pipeline-ready body/hand/face mesh extraction

Runs at 8Hz (--sample-interval auto from --hz) and outputs:
  body 33 keypoints + hands 21x2 + face mesh 468 (optional)

Output:
{
  "frame_count": N,
  "fps": 30.0,
  "frames": [
    {
      "frame": 10,
      "timestamp": 0.33,
      "persons": [
        {
          "person_id": 0,
          "pose": {"landmarks": [[x,y,z,visibility], ...], "keypoints_33": [...]},
          "left_hand": {"landmarks": [[x,y,z], ...]},
          "right_hand": {"landmarks": [[x,y,z], ...]},
          "face_mesh": {"landmarks": [[x,y,z], ...]}  // optional
        }
      ]
    }
  ]
}
"""

import json, argparse, math, time, os, sys
import cv2
import numpy as np
import mediapipe as mp


class MediaPipeHolisticProcessor:
    def __init__(self, sample_interval: int = 0, hz: float = 8.0,
                 model_complexity: int = 1, refine_face: bool = False,
                 output_face_mesh: bool = False):
        self.sample_interval = sample_interval
        self.hz = hz
        self.mp_holistic = mp.solutions.holistic
        self.holistic = self.mp_holistic.Holistic(
            static_image_mode=False,
            model_complexity=model_complexity,
            smooth_landmarks=True,
            refine_face_landmarks=refine_face,
            min_detection_confidence=0.5,
            min_tracking_confidence=0.5,
        )
        self.output_face_mesh = output_face_mesh

    def run(self, video_path: str, output_path: str, uuid: str = ""):
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise RuntimeError(f"Cannot open: {video_path}")
        fps = cap.get(cv2.CAP_PROP_FPS)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        if self.sample_interval <= 0:
            self.sample_interval = max(1, int(math.ceil(fps / self.hz)))

        print(f"[MEDIAPIPE] Video: {width}x{height}, {fps:.1f}fps, "
              f"{total_frames}f, interval={self.sample_interval}")

        frames_out = []
        frame_num = 0
        processed = 0
        t0 = time.time()

        while True:
            ret, frame = cap.read()
            if not ret:
                break
            frame_num += 1

            if (frame_num - 1) % self.sample_interval != 0:
                continue

            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            results = self.holistic.process(rgb)

            if not results.pose_landmarks:
                continue

            persons = []
            person_data = {"person_id": 0}

            pose_lm = []
            for lm in results.pose_landmarks.landmark:
                pose_lm.append([lm.x, lm.y, lm.z, lm.visibility])
            person_data["pose"] = {
                "landmarks": pose_lm,
                "keypoints_33": [f"{lm.x:.4f},{lm.y:.4f},{lm.z:.4f},{lm.visibility:.2f}"
                                for lm in results.pose_landmarks.landmark],
            }

            if results.left_hand_landmarks:
                lh = []
                for lm in results.left_hand_landmarks.landmark:
                    lh.append([lm.x, lm.y, lm.z])
                person_data["left_hand"] = {"landmarks": lh}
            if results.right_hand_landmarks:
                rh = []
                for lm in results.right_hand_landmarks.landmark:
                    rh.append([lm.x, lm.y, lm.z])
                person_data["right_hand"] = {"landmarks": rh}

            if self.output_face_mesh and results.face_landmarks:
                fm = []
                for lm in results.face_landmarks.landmark:
                    fm.append([lm.x, lm.y, lm.z])
                person_data["face_mesh"] = {"landmarks": fm}

            persons.append(person_data)

            timestamp = frame_num / fps
            frames_out.append({
                "frame": frame_num,
                "timestamp": timestamp,
                "persons": persons,
            })
            processed += 1

            if processed % 100 == 0:
                elapsed = time.time() - t0
                pct = int(frame_num * 100 / max(total_frames, 1))
                print(f"[MEDIAPIPE] {processed} frames processed, "
                      f"{pct}% complete, {elapsed:.0f}s")

        cap.release()
        self.holistic.close()

        output = {
            "frame_count": len(frames_out),
            "fps": fps,
            "frames": frames_out,
        }
        with open(output_path, "w") as f:
            json.dump(output, f, indent=2, ensure_ascii=False)

        elapsed = time.time() - t0
        print(f"[MEDIAPIPE] Done: {len(frames_out)} frames, {elapsed:.1f}s → {output_path}")


def main():
    parser = argparse.ArgumentParser(description="MediaPipe Holistic Processor v1.11")
    parser.add_argument("video_path", help="Video file path")
    parser.add_argument("output_path", help="Output JSON path")
    parser.add_argument("--uuid", "-u", default="")
    parser.add_argument("--sample-interval", type=int, default=0,
                        help="Frame interval (0 = auto from --hz)")
    parser.add_argument("--hz", type=float, default=8.0)
    parser.add_argument("--model-complexity", type=int, default=1, choices=[0, 1, 2])
    parser.add_argument("--face-mesh", action="store_true",
                        help="Include 468-point face mesh")
    args = parser.parse_args()

    processor = MediaPipeHolisticProcessor(
        sample_interval=args.sample_interval,
        hz=args.hz,
        model_complexity=args.model_complexity,
        output_face_mesh=args.face_mesh,
    )
    processor.run(args.video_path, args.output_path, args.uuid)


if __name__ == "__main__":
    main()
