#!/opt/homebrew/bin/python3.11
"""
OCR 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
"""

import sys
import json
import os
import argparse
import signal
import tempfile
import time
import subprocess
import traceback
import threading
from datetime import datetime
from typing import Dict, Any, List, 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/ocr_processor_contract_v1.py"
PROCESSOR_VERSION = "1.0.0"
MODEL_NAME = "easyocr"
MODEL_VERSION = "1.7"

# Unified configuration defaults
DEFAULT_TIMEOUT = 1800  # 30 minutes
DEFAULT_LANGUAGES = ["en"]
DEFAULT_CONFIDENCE = 0.7
DEFAULT_GPU = False
DEFAULT_MODEL_PATH = "~/.EasyOCR/model"


# Signal handling with timeout support
class SignalHandler:
    """Handle system signals for graceful shutdown"""

    def __init__(self):
        self.shutdown_requested = False
        self.timeout_reached = False
        self.original_handlers = {}

    def setup(self):
        """Set up signal handlers"""
        self.original_handlers[signal.SIGTERM] = signal.signal(
            signal.SIGTERM, self.handle_signal
        )
        self.original_handlers[signal.SIGINT] = signal.signal(
            signal.SIGINT, self.handle_signal
        )

    def handle_signal(self, signum, frame):
        """Handle received signal"""
        signal_name = "SIGTERM" if signum == signal.SIGTERM else "SIGINT"
        print(
            f"[{PROCESSOR_NAME}] Received {signal_name}, initiating graceful shutdown...",
            file=sys.stderr,
        )
        self.shutdown_requested = True

    def timeout_handler(self):
        """Handle timeout signal"""
        print(
            f"[{PROCESSOR_NAME}] Processing timeout reached, initiating graceful shutdown...",
            file=sys.stderr,
        )
        self.timeout_reached = True
        self.shutdown_requested = True

    def restore(self):
        """Restore original signal handlers"""
        for sig, handler in self.original_handlers.items():
            signal.signal(sig, handler)


# Timeout manager
class TimeoutManager:
    """Manage processing timeouts"""

    def __init__(self, overall_timeout: int):
        self.overall_timeout = overall_timeout
        self.start_time = time.time()
        self.timeout_thread = None
        self.timeout_event = threading.Event()

    def start_overall_timer(self):
        """Start overall timeout timer"""
        if self.overall_timeout > 0:
            self.timeout_thread = threading.Thread(
                target=self._overall_timeout_watcher, daemon=True
            )
            self.timeout_thread.start()

    def _overall_timeout_watcher(self):
        """Watch for overall timeout"""
        time.sleep(self.overall_timeout)
        if not self.timeout_event.is_set():
            self.timeout_event.set()
            print(
                f"[{PROCESSOR_NAME}] Overall timeout ({self.overall_timeout}s) reached",
                file=sys.stderr,
            )

    def check_timeout(self, operation: str = "processing") -> Tuple[bool, str]:
        """Check if timeout has been reached"""
        elapsed = time.time() - self.start_time

        if self.timeout_event.is_set():
            return True, f"{operation} timeout reached"

        if self.overall_timeout > 0 and elapsed > self.overall_timeout:
            return True, f"Overall timeout ({self.overall_timeout}s) reached"

        return False, ""

    def get_remaining_time(self) -> float:
        """Get remaining time"""
        elapsed = time.time() - self.start_time
        return max(0, self.overall_timeout - elapsed)

    def cleanup(self):
        """Clean up timeout resources"""
        self.timeout_event.set()
        if self.timeout_thread and self.timeout_thread.is_alive():
            self.timeout_thread.join(timeout=1.0)


# Health check functions
def check_environment() -> Dict[str, Any]:
    """Check environment and dependencies"""
    checks = []

    # Check 1: EasyOCR
    try:
        import easyocr

        checks.append(
            {
                "name": "easyocr",
                "status": "available",
                "version": easyocr.__version__
                if hasattr(easyocr, "__version__")
                else "unknown",
            }
        )
    except ImportError:
        checks.append({"name": "easyocr", "status": "missing", "version": None})

    # Check 2: OpenCV
    try:
        import cv2

        checks.append(
            {
                "name": "opencv",
                "status": "available",
                "version": cv2.__version__,
            }
        )
    except ImportError:
        checks.append({"name": "opencv", "status": "missing", "version": None})

    # Check 3: FFmpeg/FFprobe
    try:
        result = subprocess.run(["ffprobe", "-version"], capture_output=True, text=True)
        if result.returncode == 0:
            version_line = result.stdout.split("\n")[0]
            checks.append(
                {"name": "ffprobe", "status": "available", "version": version_line}
            )
        else:
            checks.append({"name": "ffprobe", "status": "error", "version": None})
    except Exception:
        checks.append({"name": "ffprobe", "status": "missing", "version": None})

    # Check 4: Redis (optional)
    if REDIS_AVAILABLE:
        checks.append({"name": "redis", "status": "available", "version": "1.0.0"})
    else:
        checks.append({"name": "redis", "status": "optional_missing", "version": None})

    # Check 5: Python version
    checks.append(
        {
            "name": "python",
            "status": "available",
            "version": f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}",
        }
    )

    # Determine overall status
    critical_deps = [c for c in checks if c["name"] in ["easyocr", "opencv", "ffprobe"]]
    missing_critical = any(c["status"] in ["missing", "error"] for c in critical_deps)

    status = "healthy" if not missing_critical else "unhealthy"

    return {"status": status, "dependencies": checks}


# Model cache for performance
_reader_cache = {}


def get_easyocr_reader(languages: List[str], gpu: bool = False):
    """Get EasyOCR reader with caching"""
    cache_key = f"{','.join(sorted(languages))}_{gpu}"

    if cache_key in _reader_cache:
        return _reader_cache[cache_key]

    try:
        import easyocr

        print(f"[{PROCESSOR_NAME}] Loading EasyOCR model for languages: {languages}")
        reader = easyocr.Reader(languages, gpu=gpu)
        _reader_cache[cache_key] = reader
        return reader
    except ImportError:
        raise RuntimeError("EasyOCR library not available")
    except Exception as e:
        raise RuntimeError(f"Failed to load EasyOCR model: {e}")


# Main processor class
class OCRProcessor:
    """OCR Processor compliant with AI-Driven Processor Contract"""

    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 or ""
        self.check_health = check_health

        # Get unified configuration from environment
        self.timeout = int(os.environ.get("MOMENTRY_OCR_TIMEOUT", str(DEFAULT_TIMEOUT)))
        languages_str = os.environ.get(
            "MOMENTRY_OCR_LANGUAGES", ",".join(DEFAULT_LANGUAGES)
        )
        self.languages = [
            lang.strip() for lang in languages_str.split(",") if lang.strip()
        ]
        self.confidence = float(
            os.environ.get("MOMENTRY_OCR_CONFIDENCE", str(DEFAULT_CONFIDENCE))
        )
        self.gpu = (
            os.environ.get("MOMENTRY_OCR_GPU", str(DEFAULT_GPU)).lower() == "true"
        )
        self.model_path = os.environ.get("MOMENTRY_OCR_MODEL_PATH", DEFAULT_MODEL_PATH)

        # Initialize components
        self.publisher = None
        if REDIS_AVAILABLE and self.uuid:
            try:
                self.publisher = RedisPublisher(self.uuid)
            except Exception as e:
                print(
                    f"[{PROCESSOR_NAME}] Failed to initialize Redis publisher: {e}",
                    file=sys.stderr,
                )

        self.timeout_manager = TimeoutManager(self.timeout)
        self.signal_handler = SignalHandler()
        self.start_time = time.time()
        self.cleanup_files = []

        # Set up signal handling
        self.signal_handler.setup()
        atexit.register(self.cleanup)

    def publish(self, msg_type: str, message: str, progress: Optional[float] = None):
        """Publish message to Redis if available"""
        if self.publisher and REDIS_AVAILABLE:
            try:
                if msg_type == "progress" and progress is not None:
                    self.publisher.progress(
                        PROCESSOR_NAME, int(progress * 100), 0, message
                    )
                else:
                    getattr(self.publisher, msg_type)(PROCESSOR_NAME, message)
            except Exception as e:
                print(f"[{PROCESSOR_NAME}] Redis publish error: {e}", file=sys.stderr)

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

        # Check if it's a video file
        try:
            cmd = [
                "ffprobe",
                "-v",
                "error",
                "-select_streams",
                "v:0",
                "-show_entries",
                "stream=codec_type",
                "-of",
                "csv=p=0",
                self.video_path,
            ]
            result = subprocess.run(cmd, capture_output=True, text=True)
            if "video" not in result.stdout:
                return False, f"Not a video file: {self.video_path}"
        except Exception:
            # If ffprobe fails, still try to process
            pass

        return True, "Input validation passed"

    def extract_frames(self, video_path: str, interval: float = 1.0) -> List[str]:
        """Extract frames from video at specified interval"""
        temp_dir = tempfile.mkdtemp(prefix="ocr_frames_")
        self.cleanup_files.append(temp_dir)

        # Create output pattern
        output_pattern = os.path.join(temp_dir, "frame_%04d.jpg")

        cmd = [
            "ffmpeg",
            "-i",
            video_path,
            "-vf",
            f"fps=1/{interval}",  # 1 frame per `interval` seconds
            "-q:v",
            "2",  # Quality factor (2-31, lower is better)
            output_pattern,
        ]

        try:
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
            if result.returncode != 0:
                raise RuntimeError(f"FFmpeg failed: {result.stderr}")

            # Get extracted frames
            frame_files = sorted(
                [f for f in os.listdir(temp_dir) if f.endswith(".jpg")]
            )
            return [os.path.join(temp_dir, f) for f in frame_files]

        except subprocess.TimeoutExpired:
            raise RuntimeError("Frame extraction timeout after 300s")
        except Exception as e:
            raise RuntimeError(f"Frame extraction failed: {e}")

    def process_frame(
        self, reader, frame_path: str, frame_idx: int, timestamp: float
    ) -> Dict[str, Any]:
        """Process a single frame for OCR"""
        try:
            import cv2

            # Read image
            image = cv2.imread(frame_path)
            if image is None:
                return {"frame": frame_idx, "timestamp": timestamp, "texts": []}

            # Perform OCR
            results = reader.readtext(
                image,
                detail=1,
                paragraph=False,
                contrast_ths=0.1,
                adjust_contrast=0.5,
                text_threshold=self.confidence,
            )

            # Format results
            texts = []
            for bbox, text, prob in results:
                if prob >= self.confidence:
                    # Convert bounding box to x, y, width, height
                    xs = [point[0] for point in bbox]
                    ys = [point[1] for point in bbox]
                    x_min, x_max = min(xs), max(xs)
                    y_min, y_max = min(ys), max(ys)

                    texts.append(
                        {
                            "text": text.strip(),
                            "x": int(x_min),
                            "y": int(y_min),
                            "width": int(x_max - x_min),
                            "height": int(y_max - y_min),
                            "confidence": float(prob),
                        }
                    )

            return {"frame": frame_idx, "timestamp": timestamp, "texts": texts}

        except Exception as e:
            print(
                f"[{PROCESSOR_NAME}] Error processing frame {frame_idx}: {e}",
                file=sys.stderr,
            )
            return {"frame": frame_idx, "timestamp": timestamp, "texts": []}

    def process(self) -> Dict[str, Any]:
        """Main processing method"""
        self.publish("info", f"Starting OCR processing: {self.video_path}")
        self.publish(
            "info",
            f"Configuration: timeout={self.timeout}s, languages={self.languages}, confidence={self.confidence}",
        )

        # Validate input
        is_valid, validation_msg = self.validate_input()
        if not is_valid:
            raise RuntimeError(f"Input validation failed: {validation_msg}")

        self.publish("info", "Input validation passed")

        # Start timeout monitoring
        self.timeout_manager.start_overall_timer()

        # Load OCR model
        self.publish("info", f"Loading OCR model for languages: {self.languages}")
        try:
            reader = get_easyocr_reader(self.languages, self.gpu)
        except RuntimeError as e:
            raise RuntimeError(f"Failed to load OCR model: {e}")

        self.publish("progress", "OCR model loaded", 0.1)

        # Extract frames (1 frame per second)
        self.publish("info", "Extracting frames from video...")
        try:
            frame_files = self.extract_frames(self.video_path, interval=1.0)
        except RuntimeError as e:
            raise RuntimeError(f"Frame extraction failed: {e}")

        self.publish("info", f"Extracted {len(frame_files)} frames")
        self.publish("progress", "Frame extraction complete", 0.3)

        # Check for timeout
        timeout_reached, timeout_msg = self.timeout_manager.check_timeout(
            "frame extraction"
        )
        if timeout_reached:
            raise RuntimeError(f"Frame extraction {timeout_msg}")

        # Process frames
        self.publish("info", "Processing frames with OCR...")
        frames = []
        total_frames = len(frame_files)

        for idx, frame_file in enumerate(frame_files):
            # Check for shutdown request
            if self.signal_handler.shutdown_requested:
                raise RuntimeError("Processing interrupted by signal")

            # Check for timeout
            timeout_reached, timeout_msg = self.timeout_manager.check_timeout(
                "frame processing"
            )
            if timeout_reached:
                raise RuntimeError(f"Frame processing {timeout_msg}")

            # Process frame
            timestamp = idx * 1.0  # 1 frame per second
            frame_result = self.process_frame(reader, frame_file, idx, timestamp)
            frames.append(frame_result)

            # Report progress
            progress = 0.3 + (idx / total_frames) * 0.6
            self.publish(
                "progress", f"Processed frame {idx + 1}/{total_frames}", progress
            )

        self.publish("progress", "Frame processing complete", 0.9)

        # Prepare final result
        result = {
            "processor_name": PROCESSOR_NAME,
            "processor_version": PROCESSOR_VERSION,
            "contract_version": CONTRACT_VERSION,
            "model_name": MODEL_NAME,
            "model_version": MODEL_VERSION,
            "video_path": self.video_path,
            "timestamp": datetime.utcnow().isoformat() + "Z",
            "processing_time_seconds": time.time() - self.start_time,
            "configuration": {
                "timeout_seconds": self.timeout,
                "languages": self.languages,
                "confidence_threshold": self.confidence,
                "gpu_enabled": self.gpu,
                "frame_interval_seconds": 1.0,
            },
            "frame_count": len(frames),
            "fps": len(frames) / (len(frames) * 1.0) if frames else 0.0,
            "frames": frames,
        }

        self.publish("progress", "OCR processing complete", 1.0)
        self.publish(
            "complete",
            f"OCR processing completed successfully in {result['processing_time_seconds']:.1f}s",
        )

        return result

    def cleanup(self):
        """Clean up temporary resources"""
        self.timeout_manager.cleanup()
        self.signal_handler.restore()

        # Clean up temporary files
        for path in self.cleanup_files:
            try:
                if os.path.isdir(path):
                    import shutil

                    shutil.rmtree(path, ignore_errors=True)
                elif os.path.exists(path):
                    os.unlink(path)
            except Exception:
                pass


def main():
    """Main entry point"""
    parser = argparse.ArgumentParser(
        description="OCR Processor - AI-Driven Processor Contract Version 1.0"
    )

    # Required arguments
    parser.add_argument("video_path", help="Path to input video file")
    parser.add_argument("output_path", help="Path where JSON output should be written")

    # Optional arguments
    parser.add_argument(
        "--uuid", "-u", default="", help="UUID for Redis progress reporting"
    )
    parser.add_argument(
        "--check-health", action="store_true", help="Perform health check and exit"
    )

    # Hidden configuration arguments (following contract)
    parser.add_argument("--timeout", type=int, help=argparse.SUPPRESS)
    parser.add_argument("--languages", help=argparse.SUPPRESS)
    parser.add_argument("--confidence", type=float, help=argparse.SUPPRESS)
    parser.add_argument("--gpu", action="store_true", help=argparse.SUPPRESS)

    args = parser.parse_args()

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

    # Create processor
    processor = OCRProcessor(
        video_path=args.video_path,
        output_path=args.output_path,
        uuid=args.uuid if args.uuid else None,
        check_health=args.check_health,
    )

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

        # Write output
        with open(args.output_path, "w", encoding="utf-8") as f:
            json.dump(result, f, indent=2, ensure_ascii=False)

        print(f"[{PROCESSOR_NAME}] Processing completed successfully")
        print(f"[{PROCESSOR_NAME}] Output written to: {args.output_path}")

        sys.exit(0)

    except RuntimeError as e:
        error_msg = f"OCR processing failed: {e}"
        processor.publish("error", error_msg)
        print(f"[{PROCESSOR_NAME}] ERROR: {error_msg}", file=sys.stderr)
        sys.exit(1)

    except KeyboardInterrupt:
        processor.publish("warning", "Processing interrupted by user")
        print(f"[{PROCESSOR_NAME}] Processing interrupted by user", file=sys.stderr)
        sys.exit(130)  # Standard exit code for SIGINT

    except Exception as e:
        error_msg = f"Unexpected error: {e}\n{traceback.format_exc()}"
        processor.publish("error", error_msg)
        print(f"[{PROCESSOR_NAME}] CRITICAL ERROR: {error_msg}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()
