#!/opt/homebrew/bin/python3.11
"""
YOLO 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. Resume support with auto-save
"""

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/yolo_processor_contract_v1.py"
)
PROCESSOR_VERSION = "1.0.0"
MODEL_NAME = "yolov8n.pt"
MODEL_VERSION = "8.0"

# YOLO class names (COCO dataset 80 classes)
YOLO_NAMES = [
    "person",
    "bicycle",
    "car",
    "motorbike",
    "aeroplane",
    "bus",
    "train",
    "truck",
    "boat",
    "traffic light",
    "fire hydrant",
    "stop sign",
    "parking meter",
    "bench",
    "bird",
    "cat",
    "dog",
    "horse",
    "sheep",
    "cow",
    "elephant",
    "bear",
    "zebra",
    "giraffe",
    "backpack",
    "umbrella",
    "handbag",
    "tie",
    "suitcase",
    "frisbee",
    "skis",
    "snowboard",
    "sports ball",
    "kite",
    "baseball bat",
    "baseball glove",
    "skateboard",
    "surfboard",
    "tennis racket",
    "bottle",
    "wine glass",
    "cup",
    "fork",
    "knife",
    "spoon",
    "bowl",
    "banana",
    "apple",
    "sandwich",
    "orange",
    "broccoli",
    "carrot",
    "hot dog",
    "pizza",
    "donut",
    "cake",
    "chair",
    "sofa",
    "pottedplant",
    "bed",
    "diningtable",
    "toilet",
    "tvmonitor",
    "laptop",
    "mouse",
    "remote",
    "keyboard",
    "cell phone",
    "microwave",
    "oven",
    "toaster",
    "sink",
    "refrigerator",
    "book",
    "clock",
    "vase",
    "scissors",
    "teddy bear",
    "hair drier",
    "toothbrush",
]


class YOLOProcessor:
    """YOLO Object Detection Processor with Resume Support"""

    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_YOLO_TIMEOUT", "7200"))
        self.model_size = os.environ.get("MOMENTRY_YOLO_MODEL_SIZE", "yolov8n.pt")
        self.confidence = float(os.environ.get("MOMENTRY_YOLO_CONFIDENCE", "0.25"))
        self.iou = float(os.environ.get("MOMENTRY_YOLO_IOU", "0.45"))
        self.gpu_enabled = (
            os.environ.get("MOMENTRY_YOLO_GPU", "false").lower() == "true"
        )
        self.auto_save_interval = int(
            os.environ.get("MOMENTRY_YOLO_AUTO_SAVE_INTERVAL", "30")
        )
        self.auto_save_frames = int(
            os.environ.get("MOMENTRY_YOLO_AUTO_SAVE_FRAMES", "300")
        )

        # Parse classes to detect (empty list = all classes)
        classes_str = os.environ.get("MOMENTRY_YOLO_CLASSES", "")
        self.classes_to_detect = (
            [c.strip() for c in classes_str.split(",") if c.strip()]
            if classes_str
            else []
        )

        # 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
        self.last_save_time = time.time()
        self.last_save_frame = 0
        self.detection_data = None
        self.last_processed_frame = 0

        # 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..."
        )

        # Save current progress
        if self.detection_data:
            self._save_detection_data(is_interrupted=True)

        sys.exit(130 if signum == signal.SIGINT else 143)

    def _cleanup(self):
        """Cleanup resources on exit"""
        if self.detection_data and not self.is_interrupted:
            self._save_detection_data(is_interrupted=False)

    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 _load_existing_data(self) -> Tuple[Optional[Dict], int]:
        """Load existing detection data from file"""
        if not os.path.exists(self.output_path):
            return None, 0

        try:
            with open(self.output_path, "r", encoding="utf-8") as f:
                data = json.load(f)

            frames = data.get("frames", {})
            if frames:
                last_frame = max(int(k) for k in frames.keys())
                return data, last_frame
        except (json.JSONDecodeError, KeyError, ValueError) as e:
            self.publish("warning", f"Could not load existing file: {e}")

        return None, 0

    def _save_detection_data(self, is_interrupted: bool = False):
        """Save detection data to file"""
        if not self.detection_data:
            return

        try:
            # Update metadata
            metadata = self.detection_data.get("metadata", {})
            metadata["last_saved_at"] = datetime.now().isoformat()
            metadata["last_saved_frame"] = self.last_processed_frame
            if is_interrupted:
                metadata["status"] = "interrupted"
            self.detection_data["metadata"] = metadata

            # Save to file
            with open(self.output_path, "w", encoding="utf-8") as f:
                json.dump(self.detection_data, f, indent=2, default=str)

            self.last_save_time = time.time()
            self.publish(
                "info",
                f"Saved progress to {self.output_path} (frame {self.last_processed_frame})",
            )
        except Exception as e:
            self.publish("error", f"Failed to save detection data: {e}")

    def _should_auto_save(self, current_frame: int) -> bool:
        """Check if we should auto-save based on time or frame count"""
        time_elapsed = time.time() - self.last_save_time
        frames_elapsed = current_frame - self.last_save_frame

        return (
            time_elapsed >= self.auto_save_interval
            or frames_elapsed >= self.auto_save_frames
        )

    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
        )

        status = "healthy" if all_available else "unhealthy"

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

    def get_video_info(self) -> Dict[str, Any]:
        """Get video information using ffprobe"""
        try:
            cmd = [
                "ffprobe",
                "-v",
                "quiet",
                "-print_format",
                "json",
                "-show_format",
                "-show_streams",
                self.video_path,
            ]

            result = subprocess.run(cmd, capture_output=True, text=True, timeout=10)
            if result.returncode != 0:
                raise subprocess.CalledProcessError(result.returncode, cmd)

            info = json.loads(result.stdout)

            # Extract video stream
            video_stream = None
            for stream in info.get("streams", []):
                if stream.get("codec_type") == "video":
                    video_stream = stream
                    break

            if not video_stream:
                raise ValueError("No video stream found")

            return {
                "duration": float(info.get("format", {}).get("duration", 0)),
                "width": int(video_stream.get("width", 0)),
                "height": int(video_stream.get("height", 0)),
                "fps": eval(video_stream.get("r_frame_rate", "0/1")),
                "codec": video_stream.get("codec_name", "unknown"),
                "bitrate": int(info.get("format", {}).get("bit_rate", 0)),
                "size": int(info.get("format", {}).get("size", 0)),
            }
        except Exception as e:
            self.publish("warning", f"Could not get video info: {e}")
            return {
                "duration": 0,
                "width": 0,
                "height": 0,
                "fps": 0,
                "codec": "unknown",
                "bitrate": 0,
                "size": 0,
            }

    def process(self) -> Dict[str, Any]:
        """Main processing method"""
        self.start_time = time.time()
        self.publish("info", f"Starting YOLO processing: {self.video_path}")
        self.publish(
            "info",
            f"Configuration: timeout={self.timeout}s, model={self.model_size}, "
            f"confidence={self.confidence}, iou={self.iou}, gpu={self.gpu_enabled}",
        )

        # Validate input
        is_valid, validation_msg = self.validate_input()
        if not is_valid:
            self.publish("error", f"Input validation failed: {validation_msg}")
            return {
                "processor_name": PROCESSOR_NAME,
                "processor_version": PROCESSOR_VERSION,
                "contract_version": CONTRACT_VERSION,
                "model_name": MODEL_NAME,
                "model_version": MODEL_VERSION,
                "status": "error",
                "error": validation_msg,
                "timestamp": datetime.now().isoformat(),
            }

        # Check for existing results (resume support)
        existing_data, self.last_processed_frame = self._load_existing_data()
        resume_mode = existing_data is not None and self.last_processed_frame > 0

        if resume_mode:
            self.detection_data = existing_data
            self.publish("info", f"Resuming from frame {self.last_processed_frame + 1}")
        else:
            # Initialize new detection data
            video_info = self.get_video_info()
            self.detection_data = {
                "metadata": {
                    "video_path": self.video_path,
                    "fps": video_info["fps"],
                    "width": video_info["width"],
                    "height": video_info["height"],
                    "total_frames": 0,  # Will be updated during processing
                    "total_duration": video_info["duration"],
                    "processed_at": datetime.now().isoformat(),
                    "auto_save_interval": self.auto_save_interval,
                    "auto_save_frames": self.auto_save_frames,
                    "status": "processing",
                    "last_saved_at": datetime.now().isoformat(),
                    "last_saved_frame": 0,
                    "completed_at": None,
                    "model": self.model_size,
                    "confidence_threshold": self.confidence,
                    "iou_threshold": self.iou,
                    "classes_to_detect": self.classes_to_detect,
                },
                "frames": {},
            }

        try:
            # Import ultralytics
            from ultralytics import YOLO

            # Load model
            self.publish("info", f"Loading YOLO model: {self.model_size}")
            model = YOLO(self.model_size)

            # Set device
            device = "0" if self.gpu_enabled else "cpu"
            self.publish("info", f"Using device: {device}")

            # Process video
            self.publish("info", "Starting video processing...")

            # Use ultralytics to process video
            results = model.track(
                source=self.video_path,
                conf=self.confidence,
                iou=self.iou,
                device=device,
                classes=self.classes_to_detect if self.classes_to_detect else None,
                verbose=False,
                stream=True,  # Stream results to save memory
            )

            frame_count = 0
            for frame_idx, result in enumerate(results):
                if self.is_interrupted:
                    break

                # Skip frames if resuming
                if resume_mode and frame_idx <= self.last_processed_frame:
                    continue

                frame_count += 1
                self.last_processed_frame = frame_idx

                # Extract detections
                frame_detections = []
                if result.boxes is not None:
                    boxes = result.boxes
                    for box_idx in range(len(boxes)):
                        class_id = int(boxes.cls[box_idx])
                        confidence = float(boxes.conf[box_idx])
                        bbox = boxes.xyxy[box_idx].cpu().numpy()

                        # Convert to YOLO format
                        frame_detections.append(
                            {
                                "class_name": YOLO_NAMES[class_id]
                                if class_id < len(YOLO_NAMES)
                                else f"class_{class_id}",
                                "class_id": class_id,
                                "x": int(bbox[0]),
                                "y": int(bbox[1]),
                                "width": int(bbox[2] - bbox[0]),
                                "height": int(bbox[3] - bbox[1]),
                                "confidence": confidence,
                            }
                        )

                # Store frame results
                self.detection_data["frames"][str(frame_idx)] = {
                    "frame": frame_idx,
                    "timestamp": frame_idx / video_info["fps"]
                    if video_info["fps"] > 0
                    else 0,
                    "objects": frame_detections,
                }

                # Auto-save check
                if self._should_auto_save(frame_idx):
                    self._save_detection_data()
                    self.last_save_frame = frame_idx

                # Progress reporting
                if frame_count % 100 == 0:
                    elapsed = time.time() - self.start_time
                    fps = frame_count / elapsed if elapsed > 0 else 0
                    self.publish(
                        "info", f"Processed {frame_count} frames ({fps:.1f} fps)"
                    )

            # Update metadata
            self.detection_data["metadata"]["total_frames"] = frame_count
            self.detection_data["metadata"]["status"] = (
                "completed" if not self.is_interrupted else "interrupted"
            )
            self.detection_data["metadata"]["completed_at"] = (
                datetime.now().isoformat() if not self.is_interrupted else None
            )

            # Final save
            self._save_detection_data()

            elapsed = time.time() - self.start_time
            self.publish(
                "complete",
                f"Processing completed: {frame_count} frames in {elapsed:.1f}s",
            )

            return {
                "processor_name": PROCESSOR_NAME,
                "processor_version": PROCESSOR_VERSION,
                "contract_version": CONTRACT_VERSION,
                "model_name": MODEL_NAME,
                "model_version": MODEL_VERSION,
                "status": "success",
                "frames_processed": frame_count,
                "total_frames": frame_count,
                "elapsed_time": elapsed,
                "timestamp": datetime.now().isoformat(),
                "output_file": self.output_path,
            }

        except ImportError:
            error_msg = "ultralytics package not installed"
            self.publish("error", error_msg)
            return {
                "processor_name": PROCESSOR_NAME,
                "processor_version": PROCESSOR_VERSION,
                "contract_version": CONTRACT_VERSION,
                "model_name": MODEL_NAME,
                "model_version": MODEL_VERSION,
                "status": "error",
                "error": error_msg,
                "timestamp": datetime.now().isoformat(),
            }
        except Exception as e:
            error_msg = f"Processing error: {str(e)}"
            self.publish("error", error_msg)
            self.publish("error", traceback.format_exc())
            return {
                "processor_name": PROCESSOR_NAME,
                "processor_version": PROCESSOR_VERSION,
                "contract_version": CONTRACT_VERSION,
                "model_name": MODEL_NAME,
                "model_version": MODEL_VERSION,
                "status": "error",
                "error": error_msg,
                "timestamp": datetime.now().isoformat(),
            }


def main():
    """Main entry point"""
    parser = argparse.ArgumentParser(
        description="YOLO 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 = YOLOProcessor(
        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"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()
