#!/opt/homebrew/bin/python3.11
"""
Pose Processor - AI-Driven Processor Contract Version 1.0

Compliant with AI-Driven Processor Contract v1.0
Effective Date: 2026-03-27

Features:
1. Standardized command-line interface
2. Redis progress reporting
3. Signal handling (SIGTERM, SIGINT)
4. Health check mode
5. Resource monitoring
6. Contract-compliant JSON output
7. Unified configuration
8. Pose estimation using YOLOv8 Pose
"""

import sys
import json
import os
import argparse
import signal
import time
import subprocess
import traceback
from datetime import datetime
from typing import Dict, Any, Optional, Tuple
import atexit

# Redis Publisher for progress reporting
try:
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from redis_publisher import RedisPublisher

    REDIS_AVAILABLE = True
except ImportError:
    REDIS_AVAILABLE = False
    print(
        "WARNING: RedisPublisher not available, progress reporting disabled",
        file=sys.stderr,
    )

# Contract version
CONTRACT_VERSION = "1.0"
PROCESSOR_NAME = (
    "/Users/accusys/momentry_core_0.1/scripts/pose_processor_contract_v1.py"
)
PROCESSOR_VERSION = "1.0.0"
MODEL_NAME = "yolov8n-pose.pt"
MODEL_VERSION = "8.0"

# YOLO Pose keypoint names (COCO dataset 17 keypoints)
POSE_KEYPOINT_NAMES = [
    "nose",
    "left_eye",
    "right_eye",
    "left_ear",
    "right_ear",
    "left_shoulder",
    "right_shoulder",
    "left_elbow",
    "right_elbow",
    "left_wrist",
    "right_wrist",
    "left_hip",
    "right_hip",
    "left_knee",
    "right_knee",
    "left_ankle",
    "right_ankle",
]


class PoseProcessor:
    """Pose Estimation Processor"""

    def __init__(
        self,
        video_path: str,
        output_path: str,
        uuid: Optional[str] = None,
        check_health: bool = False,
    ):
        self.video_path = video_path
        self.output_path = output_path
        self.uuid = uuid
        self.check_health = check_health

        # Configuration from environment variables with defaults
        self.timeout = int(os.environ.get("MOMENTRY_POSE_TIMEOUT", "7200"))
        self.model_size = os.environ.get("MOMENTRY_POSE_MODEL_SIZE", "yolov8n-pose.pt")
        self.confidence = float(os.environ.get("MOMENTRY_POSE_CONFIDENCE", "0.25"))
        self.iou = float(os.environ.get("MOMENTRY_POSE_IOU", "0.45"))
        self.gpu_enabled = (
            os.environ.get("MOMENTRY_POSE_GPU", "false").lower() == "true"
        )
        self.keypoint_confidence = float(
            os.environ.get("MOMENTRY_POSE_KEYPOINT_CONFIDENCE", "0.5")
        )
        self.max_persons = int(os.environ.get("MOMENTRY_POSE_MAX_PERSONS", "10"))

        # Initialize Redis publisher if available
        self.publisher = None
        if REDIS_AVAILABLE and uuid:
            self.publisher = RedisPublisher(uuid)

        # State tracking
        self.start_time = None
        self.is_interrupted = False

        # Set up signal handlers
        signal.signal(signal.SIGTERM, self._signal_handler)
        signal.signal(signal.SIGINT, self._signal_handler)

        # Register cleanup
        atexit.register(self._cleanup)

    def _signal_handler(self, signum, frame):
        """Handle termination signals gracefully"""
        self.is_interrupted = True
        self.publish(
            "warning", f"Received signal {signum}, saving progress and exiting..."
        )
        sys.exit(130 if signum == signal.SIGINT else 143)

    def _cleanup(self):
        """Cleanup resources on exit"""
        pass

    def publish(self, level: str, message: str):
        """Publish message to Redis if available"""
        if self.publisher:
            if level == "info":
                self.publisher.info(PROCESSOR_NAME, message)
            elif level == "warning":
                self.publisher.warning(PROCESSOR_NAME, message)
            elif level == "error":
                self.publisher.error(PROCESSOR_NAME, message)
            elif level == "complete":
                self.publisher.complete(PROCESSOR_NAME, message)
        else:
            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            print(
                f"[{timestamp}] [{PROCESSOR_NAME}] [{level.upper()}] {message}",
                file=sys.stderr,
            )

    def validate_input(self) -> Tuple[bool, str]:
        """Validate input video file"""
        if not os.path.exists(self.video_path):
            return False, f"Video file not found: {self.video_path}"

        if not self.video_path.lower().endswith(
            (".mp4", ".avi", ".mov", ".mkv", ".webm")
        ):
            return False, f"Unsupported video format: {self.video_path}"

        # Check if output directory is writable
        output_dir = os.path.dirname(self.output_path)
        if output_dir and not os.path.exists(output_dir):
            try:
                os.makedirs(output_dir, exist_ok=True)
            except Exception as e:
                return False, f"Cannot create output directory: {e}"

        return True, "Input validation passed"

    def check_dependencies(self) -> Dict[str, Any]:
        """Check if all dependencies are available"""
        dependencies = {
            "ultralytics": {"status": "unknown", "version": None},
            "opencv": {"status": "unknown", "version": None},
            "ffprobe": {"status": "unknown", "version": None},
            "redis": {
                "status": "available" if REDIS_AVAILABLE else "unavailable",
                "version": None,
            },
            "python": {"status": "available", "version": sys.version.split()[0]},
        }

        # Check ultralytics
        try:
            import ultralytics

            dependencies["ultralytics"]["status"] = "available"
            dependencies["ultralytics"]["version"] = getattr(
                ultralytics, "__version__", "unknown"
            )
        except ImportError:
            dependencies["ultralytics"]["status"] = "unavailable"

        # Check opencv
        try:
            import cv2

            dependencies["opencv"]["status"] = "available"
            dependencies["opencv"]["version"] = cv2.__version__
        except ImportError:
            dependencies["opencv"]["status"] = "unavailable"

        # Check ffprobe
        try:
            result = subprocess.run(
                ["ffprobe", "-version"],
                capture_output=True,
                text=True,
                timeout=5,
            )
            if result.returncode == 0:
                dependencies["ffprobe"]["status"] = "available"
                dependencies["ffprobe"]["version"] = result.stdout.split("\n")[0]
            else:
                dependencies["ffprobe"]["status"] = "unavailable"
        except (subprocess.SubprocessError, FileNotFoundError):
            dependencies["ffprobe"]["status"] = "unavailable"

        return dependencies

    def perform_health_check(self) -> Dict[str, Any]:
        """Perform comprehensive health check"""
        dependencies = self.check_dependencies()

        # Check if essential dependencies are available
        essential_deps = ["ultralytics", "opencv", "ffprobe"]
        all_available = all(
            dependencies.get(dep, {}).get("status") == "available"
            for dep in essential_deps
        )

        return {
            "status": "healthy" if all_available else "unhealthy",
            "dependencies": dependencies,
            "contract_version": CONTRACT_VERSION,
            "processor_name": PROCESSOR_NAME,
            "processor_version": PROCESSOR_VERSION,
            "model_name": MODEL_NAME,
            "model_version": MODEL_VERSION,
            "timestamp": datetime.now().isoformat(),
        }

    def process(self) -> Dict[str, Any]:
        """Main processing method"""
        self.start_time = time.time()
        self.publish("info", f"Starting pose estimation with model: {self.model_size}")

        # Validate input
        is_valid, message = self.validate_input()
        if not is_valid:
            return {
                "status": "error",
                "error": message,
                "processor_name": PROCESSOR_NAME,
                "processor_version": PROCESSOR_VERSION,
                "timestamp": datetime.now().isoformat(),
            }

        try:
            import cv2
            from ultralytics import YOLO

            # Load video
            cap = cv2.VideoCapture(self.video_path)
            if not cap.isOpened():
                return {
                    "status": "error",
                    "error": f"Cannot open video file: {self.video_path}",
                    "processor_name": PROCESSOR_NAME,
                    "processor_version": PROCESSOR_VERSION,
                    "timestamp": datetime.now().isoformat(),
                }

            # Get video properties
            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))
            duration = total_frames / fps if fps > 0 else 0

            self.publish(
                "info",
                f"Video: {total_frames} frames, {fps:.2f} FPS, {width}x{height}, {duration:.2f}s",
            )

            # Load YOLO Pose model
            try:
                model = YOLO(self.model_size)
                self.publish("info", f"Loaded YOLO Pose model: {self.model_size}")
            except Exception as e:
                return {
                    "status": "error",
                    "error": f"Failed to load YOLO Pose model {self.model_size}: {str(e)}",
                    "processor_name": PROCESSOR_NAME,
                    "processor_version": PROCESSOR_VERSION,
                    "timestamp": datetime.now().isoformat(),
                }

            # Process frames
            frame_count = 0
            persons_detected = 0
            pose_results = []

            while not self.is_interrupted:
                ret, frame = cap.read()
                if not ret:
                    break

                frame_count += 1

                # Report progress every 100 frames
                if frame_count % 100 == 0:
                    progress = (frame_count / total_frames) * 100
                    self.publish(
                        "info",
                        f"Processed {frame_count}/{total_frames} frames ({progress:.1f}%)",
                    )

                # Run pose estimation
                results = model(
                    frame,
                    conf=self.confidence,
                    iou=self.iou,
                    max_det=self.max_persons,
                    verbose=False,
                )

                for result in results:
                    if hasattr(result, "keypoints") and result.keypoints is not None:
                        keypoints = result.keypoints.data.cpu().numpy()
                        boxes = result.boxes.data.cpu().numpy()

                        for i, (box, kpts) in enumerate(zip(boxes, keypoints)):
                            if len(kpts) > 0:
                                persons_detected += 1

                                # Extract bounding box
                                x1, y1, x2, y2, conf, cls = box

                                # Extract keypoints
                                keypoint_data = []
                                for j, kpt in enumerate(kpts):
                                    if len(kpt) >= 3:  # x, y, confidence
                                        x, y, kpt_conf = kpt[0], kpt[1], kpt[2]
                                        if kpt_conf >= self.keypoint_confidence:
                                            keypoint_data.append(
                                                {
                                                    "name": POSE_KEYPOINT_NAMES[j]
                                                    if j < len(POSE_KEYPOINT_NAMES)
                                                    else f"keypoint_{j}",
                                                    "x": float(x),
                                                    "y": float(y),
                                                    "confidence": float(kpt_conf),
                                                }
                                            )

                                pose_results.append(
                                    {
                                        "frame": frame_count,
                                        "timestamp": frame_count / fps,
                                        "person_id": i,
                                        "bounding_box": {
                                            "x": float(x1),
                                            "y": float(y1),
                                            "width": float(x2 - x1),
                                            "height": float(y2 - y1),
                                            "confidence": float(conf),
                                        },
                                        "keypoints": keypoint_data,
                                    }
                                )

                # Check timeout
                if time.time() - self.start_time > self.timeout:
                    self.publish(
                        "warning",
                        f"Timeout reached ({self.timeout}s), stopping processing",
                    )
                    break

            cap.release()

            # Save results
            result_data = {
                "status": "success",
                "processor_name": PROCESSOR_NAME,
                "processor_version": PROCESSOR_VERSION,
                "contract_version": CONTRACT_VERSION,
                "model_name": MODEL_NAME,
                "model_version": MODEL_VERSION,
                "timestamp": datetime.now().isoformat(),
                "video_info": {
                    "path": self.video_path,
                    "frames": total_frames,
                    "fps": fps,
                    "width": width,
                    "height": height,
                    "duration": duration,
                },
                "processing_info": {
                    "model": self.model_size,
                    "confidence_threshold": self.confidence,
                    "iou_threshold": self.iou,
                    "keypoint_confidence_threshold": self.keypoint_confidence,
                    "max_persons": self.max_persons,
                    "gpu_enabled": self.gpu_enabled,
                },
                "results": {
                    "frames_processed": frame_count,
                    "persons_detected": persons_detected,
                    "poses": pose_results,
                },
            }

            # Write output
            with open(self.output_path, "w") as f:
                json.dump(result_data, f, indent=2)

            processing_time = time.time() - self.start_time
            self.publish(
                "complete",
                f"Pose estimation completed: {persons_detected} persons detected in {frame_count} frames ({processing_time:.1f}s)",
            )

            return {
                "status": "success",
                "frames_processed": frame_count,
                "persons_detected": persons_detected,
                "output_file": self.output_path,
                "processing_time": processing_time,
                "processor_name": PROCESSOR_NAME,
                "processor_version": PROCESSOR_VERSION,
                "timestamp": datetime.now().isoformat(),
            }

        except Exception as e:
            error_msg = f"Error during pose estimation: {str(e)}"
            self.publish("error", error_msg)
            return {
                "status": "error",
                "error": error_msg,
                "traceback": traceback.format_exc(),
                "processor_name": PROCESSOR_NAME,
                "processor_version": PROCESSOR_VERSION,
                "timestamp": datetime.now().isoformat(),
            }


def main():
    """Main entry point"""
    parser = argparse.ArgumentParser(
        description="Pose Processor - AI-Driven Processor Contract Version 1.0"
    )
    parser.add_argument("video_path", help="Path to input video file")
    parser.add_argument("output_path", help="Path where JSON output should be written")
    parser.add_argument("--uuid", "-u", help="UUID for Redis progress reporting")
    parser.add_argument(
        "--check-health", action="store_true", help="Perform health check and exit"
    )

    args = parser.parse_args()

    # Create processor instance
    processor = PoseProcessor(
        video_path=args.video_path,
        output_path=args.output_path,
        uuid=args.uuid,
        check_health=args.check_health,
    )

    # Health check mode
    if args.check_health:
        health_result = processor.perform_health_check()
        print(json.dumps(health_result, indent=2))
        sys.exit(0 if health_result["status"] == "healthy" else 1)

    # Process video
    try:
        result = processor.process()

        # Print result summary
        if result["status"] == "success":
            print(f"Successfully processed {result['frames_processed']} frames")
            print(f"Detected {result['persons_detected']} persons")
            print(f"Output saved to: {result['output_file']}")
        else:
            print(f"Error: {result.get('error', 'Unknown error')}")
            sys.exit(1)

    except KeyboardInterrupt:
        print("\nProcessing interrupted by user")
        sys.exit(130)
    except Exception as e:
        print(f"Fatal error: {e}")
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()
