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

Compliant with AI-Driven Processor Contract v1.0
Effective Date: 2025-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
"""

import sys
import json
import os
import argparse
import signal
import tempfile
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/asr_processor_contract_v1.py"
PROCESSOR_VERSION = "2.0.0"
MODEL_NAME = "base"
MODEL_VERSION = "unknown"


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

    def __init__(self):
        self.shutdown_requested = 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 restore(self):
        """Restore original signal handlers"""
        for sig, handler in self.original_handlers.items():
            signal.signal(sig, handler)


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

    # Check 1: Whisper
    try:
        import whisper

        checks.append(
            {
                "name": "whisper",
                "status": "available",
                "version": whisper.__version__
                if hasattr(whisper, "__version__")
                else "unknown",
            }
        )
    except ImportError:
        checks.append(
            {
                "name": "whisper",
                "status": "missing",
                "message": "openai-whisper package not installed",
            }
        )

    # Check 2: FFmpeg/FFprobe
    try:
        result = subprocess.run(["ffprobe", "-version"], capture_output=True, text=True)
        if result.returncode == 0:
            version_line = result.stdout.split("\n")[0] if result.stdout else "unknown"
            checks.append(
                {"name": "ffprobe", "status": "available", "version": version_line}
            )
        else:
            checks.append(
                {
                    "name": "ffprobe",
                    "status": "unavailable",
                    "message": "ffprobe command failed",
                }
            )
    except Exception as e:
        checks.append(
            {
                "name": "ffprobe",
                "status": "missing",
                "message": f"ffprobe not found: {e}",
            }
        )

    # Check 3: Redis (optional)
    checks.append(
        {
            "name": "redis",
            "status": "available" if REDIS_AVAILABLE else "optional_missing",
            "message": "Redis progress reporting available"
            if REDIS_AVAILABLE
            else "Redis progress reporting disabled",
        }
    )

    # Determine overall status
    critical_checks = [
        c
        for c in checks
        if c["name"] in ["whisper", "ffprobe"]
        and c["status"] not in ["available", "optional_missing"]
    ]

    if critical_checks:
        overall_status = "unhealthy"
    else:
        overall_status = "healthy"

    return {
        "status": overall_status,
        "dependencies": checks,
        "timestamp": datetime.now().isoformat(),
    }


# Whisper model cache
_whisper_model_cache = {}


def get_whisper_model(model_name: str = "base"):
    """Get Whisper model with caching"""
    if model_name not in _whisper_model_cache:
        import whisper

        print(
            f"[{PROCESSOR_NAME}] Loading Whisper model: {model_name}", file=sys.stderr
        )
        _whisper_model_cache[model_name] = whisper.load_model(model_name)
    return _whisper_model_cache[model_name]


# Main processor class
class ASRProcessor:
    """ASR Processor compliant with AI-Driven Processor Contract"""

    def __init__(
        self,
        video_path: str,
        output_path: str,
        uuid: str = "",
        model_name: str = "base",
        chunk_size: int = 300,
        publisher=None,
    ):
        self.video_path = video_path
        self.output_path = output_path
        self.uuid = uuid
        self.model_name = model_name
        self.chunk_size = chunk_size
        self.publisher = publisher
        self.start_time = time.time()
        self.signal_handler = SignalHandler()
        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 for audio stream
        if not self._has_audio_stream():
            return False, f"No audio stream found in: {self.video_path}"

        return True, "Input validation passed"

    def _has_audio_stream(self) -> bool:
        """Check if video has audio stream"""
        try:
            cmd = [
                "ffprobe",
                "-v",
                "error",
                "-select_streams",
                "a",
                "-show_entries",
                "stream=codec_type",
                "-of",
                "csv=p=0",
                self.video_path,
            ]
            result = subprocess.run(cmd, capture_output=True, text=True)
            return "audio" in result.stdout
        except Exception:
            return False

    def _get_media_duration(self) -> float:
        """Get media duration in seconds"""
        try:
            cmd = [
                "ffprobe",
                "-v",
                "error",
                "-show_entries",
                "format=duration",
                "-of",
                "csv=p=0",
                self.video_path,
            ]
            result = subprocess.run(cmd, capture_output=True, text=True)
            return float(result.stdout.strip())
        except Exception as e:
            print(
                f"[{PROCESSOR_NAME}] Warning: Failed to get duration: {e}",
                file=sys.stderr,
            )
            return 0.0

    def _extract_audio(self, audio_path: str) -> bool:
        """Extract audio to temporary file"""
        try:
            cmd = [
                "ffmpeg",
                "-i",
                self.video_path,
                "-vn",
                "-acodec",
                "pcm_s16le",
                "-ar",
                "16000",
                "-ac",
                "1",
                "-y",
                audio_path,
            ]

            self.publish("info", f"Extracting audio to: {audio_path}")
            result = subprocess.run(cmd, capture_output=True, text=True)

            if result.returncode != 0:
                self.publish("error", f"Audio extraction failed: {result.stderr[:100]}")
                return False

            return os.path.exists(audio_path) and os.path.getsize(audio_path) > 0

        except Exception as e:
            self.publish("error", f"Audio extraction error: {e}")
            return False

    def process(self) -> Dict[str, Any]:
        """Main processing method"""
        try:
            # Check for shutdown request
            if self.signal_handler.shutdown_requested:
                raise KeyboardInterrupt("Shutdown requested by signal")

            # 1. Prepare working directory
            work_dir = tempfile.mkdtemp(prefix=f"{PROCESSOR_NAME}_")
            self.cleanup_files.append(work_dir)
            self.publish("info", f"Working directory: {work_dir}")

            # 2. Get media duration
            duration = self._get_media_duration()
            self.publish("info", f"Media duration: {duration:.2f} seconds")

            # 3. Process based on duration
            self.publish("info", "Starting transcription...")

            if duration <= self.chunk_size or self.chunk_size <= 0:
                # Single file processing
                result = self._process_single_file(work_dir)
                processing_mode = "direct"
                chunk_count = 1
            else:
                # Chunked processing (simplified for now)
                result = self._process_single_file(work_dir)
                processing_mode = "chunked"
                chunk_count = max(1, int(duration / self.chunk_size))

            # 4. Add contract-compliant metadata
            processing_time = time.time() - self.start_time
            result.update(
                {
                    "processor_name": PROCESSOR_NAME,
                    "processor_version": PROCESSOR_VERSION,
                    "contract_version": CONTRACT_VERSION,
                    "model_name": MODEL_NAME,
                    "model_version": MODEL_VERSION,
                    "processing_mode": processing_mode,
                    "chunk_count": chunk_count,
                    "chunk_duration": self.chunk_size
                    if processing_mode == "chunked"
                    else 0,
                    "metadata": {
                        "processing_time_seconds": processing_time,
                        "video_path": self.video_path,
                        "duration_seconds": duration,
                        "model": self.model_name,
                        "timestamp": datetime.now().isoformat(),
                    },
                }
            )

            # 5. Cleanup
            self.cleanup()

            self.publish(
                "complete", f"Processing completed in {processing_time:.2f} seconds"
            )
            return result

        except KeyboardInterrupt:
            self.publish("warning", "Processing interrupted by user")
            raise
        except Exception as e:
            self.publish("error", f"Processing failed: {e}")
            raise

    def _process_single_file(self, work_dir: str) -> Dict[str, Any]:
        """Process single file (no chunking)"""
        # 1. Extract audio
        audio_path = os.path.join(work_dir, "audio.wav")
        self.cleanup_files.append(audio_path)

        if not self._extract_audio(audio_path):
            raise RuntimeError("Failed to extract audio")

        # 2. Load model
        self.publish("info", f"Loading Whisper model: {self.model_name}")
        model = get_whisper_model(self.model_name)

        # 3. Transcribe
        self.publish("progress", "Transcribing audio...", 0.3)

        result = model.transcribe(audio_path)

        # 4. Format segments
        segments = []
        total_segments = len(result.get("segments", []))

        for i, segment in enumerate(result.get("segments", [])):
            segments.append(
                {
                    "start": segment.get("start", 0.0),
                    "end": segment.get("end", 0.0),
                    "text": segment.get("text", "").strip(),
                    "confidence": segment.get("confidence", 0.0),
                }
            )

            # Update progress
            if i % 10 == 0 and total_segments > 0:
                progress = 0.3 + 0.7 * (i / total_segments)
                self.publish(
                    "progress",
                    f"Transcribing segment {i + 1}/{total_segments}",
                    progress,
                )

        return {
            "language": result.get("language"),
            "language_probability": result.get("language_probability"),
            "segments": segments,
            "summary": {
                "segment_count": len(segments),
                "total_duration": result.get("duration", 0.0),
            },
        }

    def save_result(self, result: Dict[str, Any]):
        """Save result to output file"""
        # Ensure output directory exists
        output_dir = os.path.dirname(self.output_path)
        if output_dir and not os.path.exists(output_dir):
            os.makedirs(output_dir, exist_ok=True)

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

        self.publish("info", f"Result saved to: {self.output_path}")

    def cleanup(self):
        """Clean up temporary resources"""
        for file_path in self.cleanup_files:
            try:
                if os.path.isdir(file_path):
                    import shutil

                    shutil.rmtree(file_path)
                elif os.path.exists(file_path):
                    os.remove(file_path)
            except Exception as e:
                print(f"[{PROCESSOR_NAME}] Cleanup warning: {e}", file=sys.stderr)

        self.cleanup_files.clear()
        self.signal_handler.restore()


# Main function
def main():
    parser = argparse.ArgumentParser(
        description=f"{PROCESSOR_NAME} Processor - AI-Driven Processor Contract v{CONTRACT_VERSION}",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    # 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 (does not process video)",
    )

    # Hidden/configuration arguments
    parser.add_argument(
        "--model", default="base", help=argparse.SUPPRESS
    )  # Hidden from help
    parser.add_argument(
        "--chunk-size", type=int, default=300, help=argparse.SUPPRESS
    )  # Hidden from help

    args = parser.parse_args()

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

    # Create Redis publisher if UUID provided
    publisher = None
    if args.uuid and REDIS_AVAILABLE:
        try:
            publisher = RedisPublisher(args.uuid)
        except Exception as e:
            print(f"WARNING: Failed to create Redis publisher: {e}", file=sys.stderr)

    # Create and run processor
    processor = ASRProcessor(
        video_path=args.video_path,
        output_path=args.output_path,
        uuid=args.uuid,
        model_name=args.model,
        chunk_size=args.chunk_size,
        publisher=publisher,
    )

    # Validate input
    valid, msg = processor.validate_input()
    if not valid:
        print(f"ERROR: {msg}", file=sys.stderr)
        sys.exit(1)

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

        # Save result
        processor.save_result(result)

        # Print success message
        print(f"[{PROCESSOR_NAME}] Processing completed successfully", file=sys.stderr)
        print(
            f"[{PROCESSOR_NAME}] Output saved to: {args.output_path}", file=sys.stderr
        )

        sys.exit(0)

    except KeyboardInterrupt:
        print(f"[{PROCESSOR_NAME}] Processing interrupted by user", file=sys.stderr)
        sys.exit(130)

    except Exception as e:
        print(f"ERROR: {e}", file=sys.stderr)
        if os.environ.get("ASR_DEBUG") == "1":
            print(f"DEBUG: {traceback.format_exc()}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()
