"""Plugin implementation for Mistral Voxtral transcription through vLLM server"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/plugin.ipynb.

# %% auto 0
__all__ = ['VLLMServer', 'VoxtralVLLMPlugin']

# %% ../nbs/plugin.ipynb 3
import json
import logging
from pathlib import Path
from typing import Dict, Any, Optional, List, Union, Generator
import tempfile
import warnings
from threading import Thread
import subprocess
import time
import requests
import atexit
import signal
import threading
import queue
import re
from datetime import datetime

from fastcore.basics import patch

import numpy as np
import soundfile as sf

try:
    from openai import OpenAI
    OPENAI_AVAILABLE = True
except ImportError:
    OPENAI_AVAILABLE = False

try:
    from mistral_common.protocol.transcription.request import TranscriptionRequest
    from mistral_common.protocol.instruct.chunk import RawAudio
    from mistral_common.audio import Audio
    MISTRAL_COMMON_AVAILABLE = True
except ImportError:
    MISTRAL_COMMON_AVAILABLE = False
    
from cjm_transcription_plugin_system.plugin_interface import PluginInterface
from cjm_transcription_plugin_system.core import AudioData, TranscriptionResult

# %% ../nbs/plugin.ipynb 4
import subprocess
import time
import requests
import atexit
import signal
import threading
import queue
import re
from contextlib import contextmanager
from typing import Optional, Callable, List
from datetime import datetime

class VLLMServer:
    def __init__(
        self,
        model: str = "mistralai/Voxtral-Mini-3B-2507",
        port: int = 8000,
        host: str = "0.0.0.0",
        gpu_memory_utilization: float = 0.85,
        log_level: str = "INFO",  # DEBUG, INFO, WARNING, ERROR
        capture_logs: bool = True,
        **kwargs
    ):
        self.model = model
        self.port = port
        self.host = host
        self.base_url = f"http://localhost:{self.port}"
        self.process: Optional[subprocess.Popen] = None
        self.capture_logs = capture_logs
        self.log_level = log_level
        
        # Log management
        self.log_queue = queue.Queue()
        self.log_thread = None
        self.stop_logging = threading.Event()
        self.log_callbacks: List[Callable] = []
        
        # Build command
        self.cmd = [
            "python", "-m", "vllm.entrypoints.openai.api_server",
            "--model", model,
            "--port", str(self.port),
            "--host", host,
            "--gpu-memory-utilization", str(gpu_memory_utilization),
            "--tokenizer-mode", "mistral",
            "--config-format", "mistral",
            "--load-format", "mistral",
        ]
        
        # Add any additional arguments
        for key, value in kwargs.items():
            self.cmd.extend([f"--{key.replace('_', '-')}", str(value)])
    
    def add_log_callback(self, callback: Callable[[str], None]):
        """Add a callback function that will be called for each log line.
        
        Args:
            callback: Function that takes a log line string as input
        """
        self.log_callbacks.append(callback)
    
    def _process_log_line(self, line: str):
        """Process a single log line."""
        if not line.strip():
            return
        
        # Add timestamp if not present
        if not re.match(r'^\d{2}-\d{2} \d{2}:\d{2}:\d{2}', line):
            timestamp = datetime.now().strftime("%m-%d %H:%M:%S")
            line = f"{timestamp} {line}"
        
        # Store in queue for retrieval
        self.log_queue.put(line)
        
        # Call callbacks
        for callback in self.log_callbacks:
            try:
                callback(line)
            except Exception as e:
                print(f"Error in log callback: {e}")
        
        # Print if we're capturing logs
        if self.capture_logs:
            # Color code based on log level
            if "ERROR" in line:
                print(f"\033[91m{line}\033[0m")  # Red
            elif "WARNING" in line:
                print(f"\033[93m{line}\033[0m")  # Yellow
            elif "INFO" in line:
                print(f"\033[94m{line}\033[0m")  # Blue
            else:
                print(line)
    
    def _log_reader(self, pipe, pipe_name: str):
        """Read logs from a pipe in a separate thread."""
        for line in iter(pipe.readline, ''):
            if self.stop_logging.is_set():
                break
            if line:
                # Add pipe identifier for debugging
                if pipe_name == "stderr" and "INFO" in line:
                    # Most INFO logs come through stderr in vLLM
                    self._process_log_line(line.strip())
                elif pipe_name == "stdout":
                    self._process_log_line(line.strip())
    
    def start(self, wait_for_ready: bool = True, timeout: int = 120, show_progress: bool = True):
        """Start the vLLM server.
        
        Args:
            wait_for_ready: Wait for server to be ready before returning
            timeout: Maximum time to wait for server to be ready
            show_progress: Show progress indicators during startup
        """
        if self.is_running():
            print("Server is already running")
            return
        
        print(f"Starting vLLM server with model {self.model}...")
        
        # Start process with pipes for output capture
        self.process = subprocess.Popen(
            self.cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            bufsize=1,  # Line buffered
            universal_newlines=True,
            preexec_fn=lambda: signal.signal(signal.SIGINT, signal.SIG_IGN)
        )
        
        # Start log reading threads
        self.stop_logging.clear()
        
        if self.capture_logs:
            # Create threads for reading stdout and stderr
            stdout_thread = threading.Thread(
                target=self._log_reader,
                args=(self.process.stdout, "stdout"),
                daemon=True
            )
            stderr_thread = threading.Thread(
                target=self._log_reader,
                args=(self.process.stderr, "stderr"),
                daemon=True
            )
            stdout_thread.start()
            stderr_thread.start()
        
        # Register cleanup on exit
        atexit.register(self.stop)
        
        if wait_for_ready:
            self._wait_for_server(timeout, show_progress)
    
    def _wait_for_server(self, timeout: int, show_progress: bool):
        """Wait for server to be ready to accept requests."""
        start_time = time.time()
        last_status = ""
        
        # Key phases to look for in logs
        phases = {
            "detected platform": "🔍 Detecting platform...",
            "Loading model": "📦 Loading model weights...",
            "downloading weights": "⬇️ Downloading model weights...",
            "Loading weights took": "✅ Model weights loaded",
            "Graph capturing": "📊 Capturing CUDA graphs...",
            "Graph capturing finished": "✅ CUDA graphs ready",
            "Starting vLLM API server": "🚀 Starting API server...",
            "Available routes": "✅ Server ready!"
        }
        
        while time.time() - start_time < timeout:
            # Check for new status in logs
            if show_progress and not self.log_queue.empty():
                try:
                    log_line = self.log_queue.get_nowait()
                    for key, status_msg in phases.items():
                        if key in log_line:
                            if status_msg != last_status:
                                print(f"\n  {status_msg}")
                                last_status = status_msg
                            break
                except queue.Empty:
                    pass
            
            # Try to connect to the server
            try:
                response = requests.get(f"{self.base_url}/health", timeout=1)
                if response.status_code == 200:
                    print(f"\n✅ vLLM server is ready at {self.base_url}")
                    return
            except requests.exceptions.RequestException:
                pass
            
            # Check if process has crashed
            if self.process.poll() is not None:
                raise RuntimeError(f"Server process exited with code {self.process.poll()}")
            
            time.sleep(0.5)
        
        raise TimeoutError(f"Server did not start within {timeout} seconds")
    
    def stop(self):
        """Stop the vLLM server."""
        if self.process and self.process.poll() is None:
            print("Stopping vLLM server...")
            
            # Signal threads to stop
            self.stop_logging.set()
            
            # Terminate process
            self.process.terminate()
            try:
                self.process.wait(timeout=10)
            except subprocess.TimeoutExpired:
                print("Force killing server...")
                self.process.kill()
                self.process.wait()
            
            self.process = None
            print("Server stopped")
    
    def restart(self):
        """Restart the server."""
        self.stop()
        time.sleep(2)
        self.start()
    
    def is_running(self) -> bool:
        """Check if server is running and responsive.
        
        This method checks both if the process is alive and if the server
        is actually responding to health checks.
        """
        # First check if process exists and hasn't exited
        if self.process is None or self.process.poll() is not None:
            return False
        
        # Try to connect to the server
        try:
            response = requests.get(f"{self.base_url}/health", timeout=1)
            if response.status_code == 200:
                # print(f"\n✅ vLLM server is ready at {self.base_url}")
                return True
            else:
                print(f"\n❌ vLLM server is not ready at {self.base_url}")
                return False
        except requests.exceptions.RequestException as e:
            print(f"\n❌ The following exception occurred: {e}")
            return False
    
    def get_recent_logs(self, n: int = 100) -> List[str]:
        """Get the most recent n log lines.
        
        Args:
            n: Number of recent log lines to retrieve
            
        Returns:
            List of recent log lines
        """
        logs = []
        while not self.log_queue.empty() and len(logs) < n:
            try:
                logs.append(self.log_queue.get_nowait())
            except queue.Empty:
                break
        return logs
    
    def get_metrics_from_logs(self) -> dict:
        """Parse recent logs to extract performance metrics.
        
        Returns:
            Dictionary with metrics like throughput, GPU usage, etc.
        """
        metrics = {
            "prompt_throughput": 0.0,
            "generation_throughput": 0.0,
            "running_requests": 0,
            "waiting_requests": 0,
            "gpu_kv_cache_usage": 0.0,
            "prefix_cache_hit_rate": 0.0,
        }
        
        # Look for metric lines in recent logs
        recent_logs = self.get_recent_logs(50)
        for log in recent_logs:
            if "Avg prompt throughput:" in log:
                # Parse metrics line
                match = re.search(r'Avg prompt throughput: ([\d.]+) tokens/s', log)
                if match:
                    metrics["prompt_throughput"] = float(match.group(1))
                
                match = re.search(r'Avg generation throughput: ([\d.]+) tokens/s', log)
                if match:
                    metrics["generation_throughput"] = float(match.group(1))
                
                match = re.search(r'Running: (\d+) reqs', log)
                if match:
                    metrics["running_requests"] = int(match.group(1))
                
                match = re.search(r'Waiting: (\d+) reqs', log)
                if match:
                    metrics["waiting_requests"] = int(match.group(1))
                
                match = re.search(r'GPU KV cache usage: ([\d.]+)%', log)
                if match:
                    metrics["gpu_kv_cache_usage"] = float(match.group(1))
                
                match = re.search(r'Prefix cache hit rate: ([\d.]+)%', log)
                if match:
                    metrics["prefix_cache_hit_rate"] = float(match.group(1))
                
                break  # Use the most recent metrics
        
        return metrics
    
    def tail_logs(self, follow: bool = True, n: int = 10):
        """Tail the server logs (similar to tail -f).
        
        Args:
            follow: Continue displaying new logs as they arrive
            n: Number of initial lines to display
        """
        # Display recent logs
        recent = self.get_recent_logs(n)
        for line in recent:
            print(line)
        
        if follow:
            print("\n--- Following logs (Ctrl+C to stop) ---\n")
            try:
                while self.is_running():
                    if not self.log_queue.empty():
                        print(self.log_queue.get())
                    else:
                        time.sleep(0.1)
            except KeyboardInterrupt:
                print("\n--- Stopped following logs ---")
    
    def __enter__(self):
        """Context manager entry."""
        self.start()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.stop()

# %% ../nbs/plugin.ipynb 5
class VoxtralVLLMPlugin(PluginInterface):
    """Mistral Voxtral transcription plugin via vLLM server."""
    
    def __init__(self):
        """Initialize the Voxtral VLLM plugin with default configuration."""
        self.logger = logging.getLogger(f"{__name__}.{type(self).__name__}")
        self.config = {}
        self.server: Optional[VLLMServer] = None
        self.client: Optional[OpenAI] = None
        self.model_id: Optional[str] = None
    
    @property
    def name(
        self
    ) -> str:  # Returns the plugin name
        """Get the plugin name identifier."""
        return "voxtral_vllm"
    
    @property
    def version(
        self
    ) -> str:  # Returns the plugin version
        """Get the plugin version string."""
        return "1.0.0"
    
    @property
    def supported_formats(
        self
    ) -> List[str]:  # Returns list of supported audio formats
        """Get the list of supported audio file formats."""
        return ["wav", "mp3", "flac", "m4a", "ogg", "webm", "mp4", "avi", "mov"]
    
    def get_config_schema(
        self
    ) -> Dict[str, Any]:  # Returns the configuration schema dictionary
        """Return configuration schema for Voxtral VLLM."""
        return {
            "$schema": "http://json-schema.org/draft-07/schema#",
            "type": "object",
            "title": "Voxtral VLLM Configuration",
            "properties": {
                "model_id": {
                    "type": "string",
                    "enum": ["mistralai/Voxtral-Mini-3B-2507", "mistralai/Voxtral-Small-24B-2507"],
                    "default": "mistralai/Voxtral-Mini-3B-2507",
                    "description": "Voxtral model to use. Mini is faster, Small is more accurate."
                },
                "server_mode": {
                    "type": "string",
                    "enum": ["managed", "external"],
                    "default": "managed",
                    "description": "'managed': plugin manages server lifecycle, 'external': connect to existing server"
                },
                "server_url": {
                    "type": "string",
                    "default": "http://localhost:8000",
                    "description": "vLLM server URL (for external mode)"
                },
                "server_port": {
                    "type": "integer",
                    "minimum": 1024,
                    "maximum": 65535,
                    "default": 8000,
                    "description": "Port for managed vLLM server"
                },
                "gpu_memory_utilization": {
                    "type": "number",
                    "minimum": 0.1,
                    "maximum": 1.0,
                    "default": 0.85,
                    "description": "Fraction of GPU memory to use (managed mode)"
                },
                "max_model_len": {
                    "type": "integer",
                    "minimum": 1024,
                    "maximum": 131072,
                    "default": 32768,
                    "description": "Maximum sequence length for the model"
                },
                "language": {
                    "type": ["string", "null"],
                    "default": "en",
                    "description": "Language code for transcription (e.g., 'en', 'es', 'fr')",
                    "examples": ['ar', 'nl', 'en', 'fr', 'de', 'hi', 'it', 'pt', 'es']
                },
                "temperature": {
                    "type": "number",
                    "minimum": 0.0,
                    "maximum": 2.0,
                    "default": 0.0,
                    "description": "Temperature for sampling (0.0 for deterministic)"
                },
                "streaming": {
                    "type": "boolean",
                    "default": False,
                    "description": "Enable streaming output by default"
                },
                "server_startup_timeout": {
                    "type": "integer",
                    "minimum": 30,
                    "maximum": 600,
                    "default": 120,
                    "description": "Timeout in seconds for server startup (managed mode)"
                },
                "auto_start_server": {
                    "type": "boolean",
                    "default": True,
                    "description": "Automatically start server on first use (managed mode)"
                },
                "capture_server_logs": {
                    "type": "boolean",
                    "default": True,
                    "description": "Capture vLLM server logs (managed mode)"
                },
                "dtype": {
                    "type": "string",
                    "enum": ["auto", "half", "float16", "bfloat16", "float32"],
                    "default": "auto",
                    "description": "Data type for model weights"
                },
                "tensor_parallel_size": {
                    "type": "integer",
                    "minimum": 1,
                    "maximum": 8,
                    "default": 1,
                    "description": "Number of GPUs for tensor parallelism"
                }
            },
            "required": ["model_id"],
            "additionalProperties": False
        }
    
    def get_current_config(
        self
    ) -> Dict[str, Any]:  # Returns the current configuration dictionary
        """Return current configuration."""
        defaults = self.get_config_defaults()
        return {**defaults, **self.config}
    
    def initialize(
        self,
        config: Optional[Dict[str, Any]] = None  # Configuration dictionary to initialize the plugin
    ) -> None:
        """Initialize the plugin with configuration."""
        if config:
            is_valid, error = self.validate_config(config)
            if not is_valid:
                raise ValueError(f"Invalid configuration: {error}")
        
        # Merge with defaults
        defaults = self.get_config_defaults()
        self.config = {**defaults, **(config or {})}
        
        self.model_id = self.config["model_id"]
        
        # Initialize based on server mode
        if self.config["server_mode"] == "managed":
            # Create managed server instance (but don't start yet)
            self.server = VLLMServer(
                model=self.model_id,
                port=self.config["server_port"],
                gpu_memory_utilization=self.config["gpu_memory_utilization"],
                max_model_len=self.config["max_model_len"],
                capture_logs=self.config["capture_server_logs"],
                dtype=self.config["dtype"],
                tensor_parallel_size=self.config["tensor_parallel_size"]
            )
            server_url = f"http://localhost:{self.config['server_port']}"
        else:
            # External server mode
            server_url = self.config["server_url"]
        
        # Create OpenAI client
        self.client = OpenAI(
            api_key="EMPTY",  # vLLM doesn't require API key
            base_url=f"{server_url}/v1"
        )
        
        self.logger.info(
            f"Initialized Voxtral VLLM plugin with model '{self.model_id}' "
            f"in {self.config['server_mode']} mode"
        )
    
    def _ensure_server_running(self) -> None:
        """Ensure the vLLM server is running (for managed mode)."""
        if self.config["server_mode"] == "managed" and self.server:
            if not self.server.is_running():
                print("\n\nSERVER IS NOT RUNNING\n\n")
                if self.config["auto_start_server"]:
                    self.logger.info("Starting vLLM server...")
                    self.server.start(
                        wait_for_ready=True,
                        timeout=self.config["server_startup_timeout"],
                        show_progress=self.config.get("capture_server_logs", True)
                    )
                else:
                    raise RuntimeError("vLLM server is not running and auto_start_server is disabled")
        elif self.config["server_mode"] == "external":
            # Check if external server is accessible
            try:
                response = requests.get(f"{self.config['server_url']}/health", timeout=5)
                if response.status_code != 200:
                    raise RuntimeError(f"External vLLM server returned status {response.status_code}")
            except requests.exceptions.RequestException as e:
                raise RuntimeError(f"Cannot connect to external vLLM server: {e}")
    
    def _prepare_audio(
        self,
        audio: Union[AudioData, str, Path]  # Audio data, file path, or Path object to prepare
    ) -> str:  # Returns path to the prepared audio file
        """Prepare audio for Voxtral processing."""
        if isinstance(audio, (str, Path)):
            # Already a file path
            return str(audio)
        
        elif isinstance(audio, AudioData):
            # Save AudioData to temporary file
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_file:
                # Ensure audio is in the correct format
                audio_array = audio.samples
                
                # If stereo, convert to mono
                if audio_array.ndim > 1:
                    audio_array = audio_array.mean(axis=1)
                
                # Ensure float32 and normalized
                if audio_array.dtype != np.float32:
                    audio_array = audio_array.astype(np.float32)
                
                # Normalize if needed
                if audio_array.max() > 1.0:
                    audio_array = audio_array / np.abs(audio_array).max()
                
                # Save to file
                sf.write(tmp_file.name, audio_array, audio.sample_rate)
                return tmp_file.name
        else:
            raise ValueError(f"Unsupported audio input type: {type(audio)}")
    
    def execute(
        self,
        audio: Union[AudioData, str, Path],  # Audio data or path to audio file to transcribe
        **kwargs  # Additional arguments to override config
    ) -> TranscriptionResult:  # Returns transcription result with text and metadata
        """Transcribe audio using Voxtral via vLLM."""
        # Ensure server is running
        self._ensure_server_running()
        
        # Prepare audio file
        audio_path = self._prepare_audio(audio)
        temp_file_created = not isinstance(audio, (str, Path))
        
        try:
            # Merge runtime kwargs with config
            exec_config = {**self.config, **kwargs}
            
            # Prepare inputs using mistral_common
            self.logger.info(f"Processing audio with Voxtral {self.model_id} via vLLM")
            
            input_audio = Audio.from_file(audio_path, strict=False)
            input_audio = RawAudio.from_audio(input_audio)
            
            req = TranscriptionRequest(
                model=self.model_id,
                audio=input_audio,
                language=exec_config.get("language", "en"),
                temperature=exec_config.get("temperature", 0.0)
            ).to_openai(exclude=("top_p", "seed"))
            
            # Get transcription from vLLM server
            response = self.client.audio.transcriptions.create(**req)
            
            # Create transcription result
            transcription_result = TranscriptionResult(
                text=response.text.strip(),
                confidence=None,  # vLLM doesn't provide confidence scores
                segments=None,  # vLLM doesn't provide segments by default
                metadata={
                    "model": self.model_id,
                    "language": exec_config.get("language", "en"),
                    "server_mode": self.config["server_mode"],
                    "temperature": exec_config.get("temperature", 0.0),
                }
            )
            
            self.logger.info(f"Transcription completed: {len(response.text.split())} words")
            return transcription_result
            
        finally:
            # Clean up temporary file if created
            if temp_file_created:
                try:
                    Path(audio_path).unlink()
                except Exception:
                    pass
    
    def is_available(
        self
    ) -> bool:  # Returns True if vLLM and its dependencies are available
        """Check if vLLM and required dependencies are available."""
        if not OPENAI_AVAILABLE:
            return False
        if not MISTRAL_COMMON_AVAILABLE:
            return False
        
        # Check if vLLM is installed
        try:
            import vllm
            return True
        except ImportError:
            return False
    
    def cleanup(
        self
    ) -> None:
        """Clean up resources."""
        self.logger.info("Cleaning up Voxtral VLLM plugin")
        
        # Stop managed server if running
        if self.config.get("server_mode") == "managed" and self.server:
            if self.server.is_running():
                self.logger.info("Stopping managed vLLM server")
                self.server.stop()
            self.server = None
        
        # Clear client
        self.client = None
        
        self.logger.info("Cleanup completed successfully")

# %% ../nbs/plugin.ipynb 6
@patch
def supports_streaming(
    self: VoxtralVLLMPlugin
) -> bool:
    """Check if this plugin supports streaming transcription."""
    return True

@patch
def execute_stream(
    self: VoxtralVLLMPlugin,
    audio: Union[AudioData, str, Path],  # Audio data or path to audio file
    **kwargs  # Additional plugin-specific parameters
) -> Generator[str, None, TranscriptionResult]:  # Yields text chunks, returns final result
    """Stream transcription results chunk by chunk.
    
    Args:
        audio: Audio data or path to audio file
        **kwargs: Additional plugin-specific parameters
        
    Yields:
        str: Partial transcription text chunks as they become available
        
    Returns:
        TranscriptionResult: Final complete transcription with metadata
    """
    # Ensure server is running
    self._ensure_server_running()
    
    # Prepare audio file
    audio_path = self._prepare_audio(audio)
    temp_file_created = not isinstance(audio, (str, Path))
    
    try:
        # Merge runtime kwargs with config
        exec_config = {**self.config, **kwargs}
        
        # Prepare inputs using mistral_common
        self.logger.info(f"Streaming transcription with Voxtral {self.model_id} via vLLM")
        
        input_audio = Audio.from_file(audio_path, strict=False)
        input_audio = RawAudio.from_audio(input_audio)
        
        req = TranscriptionRequest(
            model=self.model_id,
            audio=input_audio,
            language=exec_config.get("language", "en"),
            temperature=exec_config.get("temperature", 0.0)
        ).to_openai(exclude=("top_p", "seed"))
        
        # Get streaming transcription from vLLM server
        response = self.client.audio.transcriptions.create(**req, stream=True)
        
        # Collect generated text
        generated_text = ""
        for chunk in response:
            if chunk.choices[0]['delta']['content']:
                text_chunk = chunk.choices[0]['delta']['content']
                generated_text += text_chunk
                yield text_chunk
        
        # Return final result
        return TranscriptionResult(
            text=generated_text.strip(),
            confidence=None,
            segments=None,
            metadata={
                "model": self.model_id,
                "language": exec_config.get("language", "en"),
                "server_mode": self.config["server_mode"],
                "temperature": exec_config.get("temperature", 0.0),
                "streaming": True,
            }
        )
        
    finally:
        # Clean up temporary file if created
        if temp_file_created:
            try:
                Path(audio_path).unlink()
            except Exception:
                pass
