"""
example_keypoint_adder.py
Manages example videos and their precomputed keypoint JSON files and rendered keypoint videos.
Ensures that each example entry has:
1. video_path: Path to the input ASL video.
2. keypoints_path: Path to the extracted keypoints JSON.
3. keypoints_video_path: Path to the video with keypoints rendered on top.

Can be run as a standalone script to batch-process all examples or add new ones.
"""

import os
import sys
import json
import argparse
import numpy as np

# Ensure UTF-8 stdout
if hasattr(sys.stdout, 'reconfigure'):
    try:
        sys.stdout.reconfigure(encoding='utf-8', errors='replace')
    except Exception:
        pass

from predict_pose import create_mediapipe_models, predict_pose, load_video_cv
from visualize_pose import render_keypoints_video


def to_json_serializable(obj):
    """Recursively converts NumPy arrays, numbers, and tensors to JSON-serializable Python types."""
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    if isinstance(obj, (np.float32, np.float64, np.float16)):
        return float(obj)
    if isinstance(obj, (np.int32, np.int64, np.int16, np.int8, np.uint8)):
        return int(obj)
    if isinstance(obj, dict):
        return {str(k): to_json_serializable(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return [to_json_serializable(v) for v in obj]
    return obj


def load_examples_json(json_path="examples.json"):
    """Loads examples configuration from a JSON file."""
    if not os.path.exists(json_path):
        return []
    try:
        with open(json_path, "r", encoding="utf-8") as f:
            data = json.load(f)
            if isinstance(data, list):
                return data
            elif isinstance(data, dict) and "examples" in data:
                return data["examples"]
            return []
    except Exception as e:
        print(f"Error loading {json_path}: {e}")
        return []


def save_examples_json(examples, json_path="examples.json"):
    """Saves examples configuration to a JSON file."""
    try:
        with open(json_path, "w", encoding="utf-8") as f:
            json.dump(examples, f, indent=2, ensure_ascii=False)
        print(f"Saved {len(examples)} examples to {json_path}")
    except Exception as e:
        print(f"Error saving {json_path}: {e}")


def process_example(
    video_path,
    title=None,
    keypoints_path=None,
    keypoints_video_path=None,
    pose_models=None,
    force=False,
    checkpoint_folder="checkpoints/pose/"
):
    """
    Checks if keypoints JSON and rendered keypoints video exist for a video.
    If not (or if force=True), extracts keypoints and renders the visualization video.
    Returns the updated example dictionary.
    """
    video_path = os.path.normpath(video_path)
    if not os.path.exists(video_path):
        print(f"[Warning] Video file not found: {video_path}")
        return None

    video_dir = os.path.dirname(video_path)
    video_stem = os.path.splitext(os.path.basename(video_path))[0]
    
    if not title:
        title = video_stem.replace("_", " ").capitalize()

    if not keypoints_path:
        keypoints_path = os.path.join(video_dir, f"{video_stem}_keypoints.json")
    keypoints_path = os.path.normpath(keypoints_path)

    if not keypoints_video_path:
        keypoints_video_path = os.path.join(video_dir, f"{video_stem}_keypoints.mp4")
    keypoints_video_path = os.path.normpath(keypoints_video_path)

    kp_exists = os.path.exists(keypoints_path) and os.path.getsize(keypoints_path) > 0
    vid_exists = os.path.exists(keypoints_video_path) and os.path.getsize(keypoints_video_path) > 0

    if kp_exists and vid_exists and not force:
        print(f"[OK] Keypoints and video already exist for '{title}' ({video_path})")
        return {
            "title": title,
            "video_path": video_path.replace("\\", "/"),
            "keypoints_path": keypoints_path.replace("\\", "/"),
            "keypoints_video_path": keypoints_video_path.replace("\\", "/"),
        }

    print(f"\n[Processing] Generating keypoints for '{title}' ({video_path})...")

    # Load pose models if not yet loaded
    if pose_models is None:
        print("  Initializing MediaPipe & YOLO models...")
        pose_models = create_mediapipe_models(checkpoint_folder)

    # Load video frames
    video_frames, fps = load_video_cv(video_path)
    if fps is None or fps <= 0 or np.isnan(fps):
        fps = 25.0
    print(f"  Loaded {len(video_frames)} frames @ {fps:.1f} FPS")

    # Predict pose & keypoints (includes 512x512 scaling for WLASL)
    pose_results = predict_pose(video_frames, pose_models)

    # Attach bounding boxes to keypoints dict for each frame
    kps_list = pose_results.get("keypoints", [])
    bbox_face_list = pose_results.get("bbox_face", [])
    bbox_lh_list = pose_results.get("bbox_left_hand", [])
    bbox_rh_list = pose_results.get("bbox_right_hand", [])

    for i, kp in enumerate(kps_list):
        if i < len(bbox_face_list):
            kp['bbox_face'] = bbox_face_list[i]
        if i < len(bbox_lh_list):
            kp['bbox_left_hand'] = bbox_lh_list[i]
        if i < len(bbox_rh_list):
            kp['bbox_right_hand'] = bbox_rh_list[i]

    # Save keypoints to JSON (strip heavy raw image arrays to keep JSON file lightweight ~50KB and ultra-fast)
    heavy_image_keys = {'images', 'cropped_images', 'cropped_left_hand', 'cropped_right_hand', 'cropped_face'}
    clean_pose_results = {k: v for k, v in pose_results.items() if k not in heavy_image_keys}

    os.makedirs(os.path.dirname(os.path.abspath(keypoints_path)), exist_ok=True)
    serializable_results = to_json_serializable(clean_pose_results)
    with open(keypoints_path, "w", encoding="utf-8") as f:
        json.dump(serializable_results, f)
    print(f"  [Saved] Lightweight Keypoints JSON -> {keypoints_path}")

    # Render keypoint overlay video
    rendered_path = render_keypoints_video(
        video_frames=video_frames,
        keypoints_list=kps_list,
        fps=fps,
        output_path=keypoints_video_path,
    )
    print(f"  [Saved] Rendered Keypoints Video -> {rendered_path}")

    return {
        "title": title,
        "video_path": video_path.replace("\\", "/"),
        "keypoints_path": keypoints_path.replace("\\", "/"),
        "keypoints_video_path": keypoints_video_path.replace("\\", "/"),
    }


def ensure_all_examples(
    json_path="examples.json",
    default_dir="examples/displayed",
    force=False,
    checkpoint_folder="checkpoints/pose/"
):
    """
    Ensures all examples in `examples.json` (or found in default_dir)
    have precomputed keypoints JSON and rendered MP4 video.
    """
    existing_examples = load_examples_json(json_path)

    # If examples.json doesn't exist or is empty, discover .mp4 files in default_dir
    if not existing_examples and os.path.exists(default_dir):
        print(f"Discovering example videos in {default_dir}...")
        for fname in sorted(os.listdir(default_dir)):
            if fname.lower().endswith(".mp4") and not fname.endswith("_keypoints.mp4"):
                vpath = os.path.join(default_dir, fname)
                stem = os.path.splitext(fname)[0]
                existing_examples.append({
                    "title": stem.capitalize(),
                    "video_path": vpath.replace("\\", "/"),
                    "keypoints_path": os.path.join(default_dir, f"{stem}_keypoints.json").replace("\\", "/"),
                    "keypoints_video_path": os.path.join(default_dir, f"{stem}_keypoints.mp4").replace("\\", "/"),
                })

    if not existing_examples:
        print("No examples found to process.")
        return []

    pose_models = None
    # Check if any model processing is needed before initializing heavy models
    needs_processing = force or any(
        not (
            os.path.exists(ex.get("keypoints_path", "")) and 
            os.path.exists(ex.get("keypoints_video_path", "")) and
            os.path.getsize(ex.get("keypoints_path", "")) > 0 and
            os.path.getsize(ex.get("keypoints_video_path", "")) > 0
        )
        for ex in existing_examples
    )

    if needs_processing:
        print("Initializing pose estimation models for processing...")
        pose_models = create_mediapipe_models(checkpoint_folder)

    updated_examples = []
    for ex in existing_examples:
        vpath = ex.get("video_path", "")
        title = ex.get("title", "")
        kpath = ex.get("keypoints_path", "")
        kvpath = ex.get("keypoints_video_path", "")

        res = process_example(
            video_path=vpath,
            title=title,
            keypoints_path=kpath,
            keypoints_video_path=kvpath,
            pose_models=pose_models,
            force=force,
            checkpoint_folder=checkpoint_folder,
        )
        if res:
            updated_examples.append(res)
        else:
            updated_examples.append(ex)

    save_examples_json(updated_examples, json_path)
    return updated_examples


def main():
    parser = argparse.ArgumentParser(
        description="Ensure all examples have precomputed keypoints JSON and rendered pose videos."
    )
    parser.add_argument(
        "--json",
        default="examples.json",
        help="Path to examples.json (default: examples.json)"
    )
    parser.add_argument(
        "--dir",
        default="examples/displayed",
        help="Directory to scan for example videos if JSON is empty (default: examples/displayed)"
    )
    parser.add_argument(
        "--add",
        default=None,
        help="Add a specific video path to examples.json and generate its keypoints"
    )
    parser.add_argument(
        "--title",
        default=None,
        help="Optional title/label when adding a single video"
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="Force re-generation of keypoints and videos even if they exist"
    )
    args = parser.parse_args()

    if args.add:
        print(f"Adding single video: {args.add}")
        examples = load_examples_json(args.json)
        pose_models = create_mediapipe_models("checkpoints/pose/")
        res = process_example(
            video_path=args.add,
            title=args.title,
            pose_models=pose_models,
            force=args.force
        )
        if res:
            # Check if entry with same video_path already exists
            norm_new = os.path.normpath(res["video_path"])
            exists_idx = -1
            for idx, ex in enumerate(examples):
                if os.path.normpath(ex.get("video_path", "")) == norm_new:
                    exists_idx = idx
                    break
            if exists_idx >= 0:
                examples[exists_idx] = res
            else:
                examples.append(res)
            save_examples_json(examples, args.json)
    else:
        ensure_all_examples(
            json_path=args.json,
            default_dir=args.dir,
            force=args.force
        )


if __name__ == "__main__":
    main()