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

Compliant with AI-Driven Processor Contract v1.0
With unified configuration and timeout handling

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 with timeout handling
8. Model caching for performance
"""

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, Optional, Tuple
import atexit

# Whisper import at module level for proper error handling
try:
    import whisper

    WHISPER_AVAILABLE = True
    WHISPER_VERSION = getattr(whisper, "__version__", "unknown")
except ImportError:
    WHISPER_AVAILABLE = False
    WHISPER_VERSION = None

# 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 = "asr"
PROCESSOR_VERSION = "2.1.0"

# Unified configuration defaults
DEFAULT_OVERALL_TIMEOUT = 3600  # 1 hour
DEFAULT_PROCESS_TIMEOUT = 1800  # 30 minutes
DEFAULT_CHUNK_TIMEOUT = 300  # 5 minutes
DEFAULT_MODEL_SIZE = "medium"
DEFAULT_DEVICE = "cpu"
DEFAULT_LANGUAGE = "auto"


# 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, process_timeout: int, chunk_timeout: int):
        self.overall_timeout = overall_timeout
        self.process_timeout = process_timeout
        self.chunk_timeout = chunk_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, timeout_type: str = "overall") -> float:
        """Get remaining time for specified timeout type"""
        elapsed = time.time() - self.start_time

        if timeout_type == "overall":
            return max(0, self.overall_timeout - elapsed)
        elif timeout_type == "process":
            return max(0, self.process_timeout - elapsed)
        elif timeout_type == "chunk":
            return max(0, self.chunk_timeout - elapsed)

        return 0.0

    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: Whisper
    if WHISPER_AVAILABLE:
        checks.append(
            {
                "name": "whisper",
                "status": "available",
                "version": WHISPER_VERSION,
            }
        )
    else:
        checks.append({"name": "whisper", "status": "missing", "version": None})

    # 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]
            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 3: 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 4: Python version
    checks.append(
        {
            "name": "python",
            "status": "available",
            "version": f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}",
        }
    )

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


# Model cache for performance
_model_cache = {}


def get_whisper_model(model_size: str = "medium", device: str = "cpu"):
    """Get Whisper model with caching"""
    if not WHISPER_AVAILABLE:
        raise RuntimeError("Whisper library not available")

    cache_key = f"{model_size}_{device}"

    if cache_key in _model_cache:
        return _model_cache[cache_key]

    try:
        print(f"[{PROCESSOR_NAME}] Loading Whisper model: {model_size} on {device}")
        model = whisper.load_model(model_size, device=device)
        _model_cache[cache_key] = model
        return model
    except Exception as e:
        raise RuntimeError(f"Failed to load Whisper model: {e}")


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

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

        # Get unified configuration: command-line args override environment variables
        self.overall_timeout = int(
            os.environ.get("MOMENTRY_ASR_TIMEOUT", str(DEFAULT_OVERALL_TIMEOUT))
        )
        self.process_timeout = int(
            os.environ.get("MOMENTRY_ASR_PROCESS_TIMEOUT", str(DEFAULT_PROCESS_TIMEOUT))
        )
        self.chunk_timeout = int(
            os.environ.get("MOMENTRY_ASR_CHUNK_TIMEOUT", str(DEFAULT_CHUNK_TIMEOUT))
        )
        self.model_size = model_size or os.environ.get("MOMENTRY_ASR_MODEL_SIZE", DEFAULT_MODEL_SIZE)
        self.device = device or os.environ.get("MOMENTRY_ASR_DEVICE", DEFAULT_DEVICE)
        self.language = language or os.environ.get("MOMENTRY_ASR_LANGUAGE", DEFAULT_LANGUAGE)

        # 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.overall_timeout, self.process_timeout, self.chunk_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 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 extract_audio(self, video_path: str) -> str:
        """Extract audio from video file"""
        temp_dir = tempfile.mkdtemp(prefix="asr_audio_")
        audio_path = os.path.join(temp_dir, "audio.wav")
        self.cleanup_files.append(temp_dir)

        cmd = [
            "ffmpeg",
            "-i",
            video_path,
            "-vn",
            "-acodec",
            "pcm_s16le",
            "-ar",
            "16000",
            "-ac",
            "1",
            "-y",
            audio_path,
        ]

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

            return audio_path
        except subprocess.TimeoutExpired:
            raise RuntimeError(f"Audio extraction timeout after {self.chunk_timeout}s")
        except Exception as e:
            raise RuntimeError(f"Audio extraction failed: {e}")

    def transcribe_audio(self, audio_path: str) -> Dict[str, Any]:
        """Transcribe audio using Whisper"""
        if not WHISPER_AVAILABLE:
            raise RuntimeError("Whisper library not available")

        self.publish("info", f"Starting transcription with model: {self.model_size}")
        print(
            f"[DEBUG] WHISPER_AVAILABLE: {WHISPER_AVAILABLE}, whisper module: {'available' if 'whisper' in globals() else 'not in globals'}"
        )

        try:
            model = get_whisper_model(self.model_size, self.device)
            print(f"[DEBUG] Model loaded: {model}")

            # Start timeout monitoring for transcription
            self.timeout_manager.start_overall_timer()

            # Set language for transcription
            language = self.language
            if language == "auto":
                # For auto, let Whisper handle language detection internally
                language = None
                self.publish("info", "Language detection will be handled by Whisper")
            else:
                self.publish("info", f"Using specified language: {language}")

            # Perform transcription
            transcribe_language = language if language != "auto" else None
            self.publish(
                "info",
                f"Transcribing audio (language: {transcribe_language if transcribe_language else 'auto'})...",
            )

            result = model.transcribe(
                audio_path,
                language=transcribe_language,
                task="transcribe",
                beam_size=5,
                best_of=5,
            )

            # Check for timeout during transcription
            timeout_reached, timeout_msg = self.timeout_manager.check_timeout(
                "transcription"
            )
            if timeout_reached:
                raise RuntimeError(f"Transcription {timeout_msg}")

            return {
                "language": result.get("language"),
                "language_probability": result.get("language_probability"),
                "segments": [
                    {
                        "start": segment["start"],
                        "end": segment["end"],
                        "text": segment["text"].strip(),
                    }
                    for segment in result.get("segments", [])
                ],
            }

        except RuntimeError as e:
            if "timeout" in str(e).lower():
                raise
            else:
                raise RuntimeError(f"Transcription failed: {e}")
        except Exception as e:
            raise RuntimeError(f"Transcription error: {e}")

    def process(self) -> Dict[str, Any]:
        """Main processing method"""
        self.publish("info", f"Starting ASR processing: {self.video_path}")
        self.publish(
            "info",
            f"Configuration: timeout={self.overall_timeout}s, model={self.model_size}, device={self.device}",
        )

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

        # Extract audio
        self.publish("info", "Extracting audio from video...")
        audio_path = self.extract_audio(self.video_path)
        self.publish("progress", "Audio extraction complete", 0.3)

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

        # Transcribe audio
        self.publish("info", "Transcribing audio...")
        transcription_result = self.transcribe_audio(audio_path)
        self.publish("progress", "Transcription complete", 0.8)

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

        # Prepare final result
        result = {
            "processor_name": PROCESSOR_NAME,
            "processor_version": PROCESSOR_VERSION,
            "contract_version": CONTRACT_VERSION,
            "video_path": self.video_path,
            "timestamp": datetime.utcnow().isoformat() + "Z",
            "processing_time_seconds": time.time() - self.start_time,
            "configuration": {
                "model_size": self.model_size,
                "device": self.device,
                "language": self.language,
                "timeout_seconds": self.overall_timeout,
            },
            **transcription_result,
        }

        self.publish("progress", "ASR processing complete", 1.0)
        self.publish(
            "complete",
            f"ASR 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="ASR Processor - AI-Driven Processor Contract Version 2.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("--model-size", help=argparse.SUPPRESS)
    parser.add_argument("--device", help=argparse.SUPPRESS)
    parser.add_argument("--language", help=argparse.SUPPRESS)
    parser.add_argument("--timeout", type=int, 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 = ASRProcessor(
        video_path=args.video_path,
        output_path=args.output_path,
        uuid=args.uuid if args.uuid else None,
        check_health=args.check_health,
        model_size=args.model_size,
        device=args.device,
        language=args.language,
    )

    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"ASR 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()
