# lollms_client/stt_bindings/whisper/__init__.py
import os
from pathlib import Path
from typing import Optional, List, Union, Dict, Any
from ascii_colors import trace_exception, ASCIIColors
import time
import filelock

# --- Package Management and Conditional Imports ---
_whisper_installed = False
_whisper_installation_error = ""

try:
    import pipmaster as pm
    import platform # For OS detection for torch index

    # Determine initial device preference to guide torch installation
    preferred_torch_device_for_install = "cpu" # Default assumption
    
    # Tentatively set preference based on OS, assuming user might want GPU if available
    if platform.system() == "Linux" or platform.system() == "Windows":
        # On Linux/Windows, CUDA is the primary GPU acceleration for PyTorch.
        # We will try to install a CUDA version of PyTorch.
        preferred_torch_device_for_install = "cuda"
    elif platform.system() == "Darwin":
        # On macOS, MPS is the acceleration. Standard torch install usually handles this.
        preferred_torch_device_for_install = "mps" # or keep cpu if mps detection is later

    torch_pkgs = ["torch", "torchaudio","xformers"]
    audiocraft_core_pkgs = ["openai-whisper", "filelock"]

    torch_index_url = None
    if preferred_torch_device_for_install == "cuda":
        # Specify a common CUDA version index. Pip should resolve the correct torch version.
        # As of late 2023/early 2024, cu118 or cu121 are common. Let's use cu126.
        # Users with different CUDA setups might need to pre-install torch manually.
        torch_index_url = "https://download.pytorch.org/whl/cu126"
        ASCIIColors.info(f"Attempting to ensure PyTorch with CUDA support (target index: {torch_index_url})")
        # Install torch and torchaudio first from the specific index
        pm.ensure_packages(torch_pkgs, index_url=torch_index_url)
        # Then install audiocraft and other dependencies; pip should use the already installed torch
        pm.ensure_packages(audiocraft_core_pkgs)
    else:
        # For CPU, MPS, or if no specific CUDA preference was determined for install
        ASCIIColors.info("Ensuring PyTorch, AudioCraft, and dependencies using default PyPI index.")
        pm.ensure_packages(torch_pkgs + audiocraft_core_pkgs)

    import whisper
    import torch
    _whisper_installed = True
except Exception as e:
    _whisper_installation_error = str(e)
    whisper = None
    torch = None


# --- End Package Management ---

from lollms_client.lollms_stt_binding import LollmsSTTBinding

# Defines the binding name for the manager
BindingName = "WhisperSTTBinding" # Changed to avoid conflict with class name

class WhisperSTTBinding(LollmsSTTBinding):
    """
    LollmsSTTBinding implementation for OpenAI's Whisper model.
    This binding runs Whisper locally.
    Requires `ffmpeg` to be installed on the system.
    """

    # Standard Whisper model sizes
    WHISPER_MODEL_SIZES = ["turbo"]

    def __init__(self,
                 **kwargs # To catch any other LollmsSTTBinding standard args
                 ):
        """
        Initialize the Whisper STT binding.

        Args:
            model_name (str): The Whisper model size to use (e.g., "turbo").
                              Defaults to "turno".
            device (Optional[str]): The device to run the model on ("cpu", "cuda", "mps").
                                    If None, `torch` will attempt to auto-detect. Defaults to None.
        """
        super().__init__(binding_name="whisper") # Not applicable
        self.default_model_name = kwargs.get("model_name", "turbo")

        if not _whisper_installed:
            raise ImportError(f"Whisper STT binding dependencies not met. Please ensure 'openai-whisper' and 'torch' are installed. Error: {_whisper_installation_error}")

        self.device = kwargs.get("device", None)
        if self.device is None or self.device == "": # Auto-detect if not specified or empty
            if torch.cuda.is_available():
                self.device = "cuda"
            elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available(): # For Apple Silicon
                self.device = "mps"
            else:
                self.device = "cpu"
        
        # Validate device string
        if not self.device or self.device.strip() == "":
            ASCIIColors.warning("Device was empty or invalid, defaulting to 'cpu'")
            self.device = "cpu"
        
        ASCIIColors.info(f"WhisperSTTBinding: Using device '{self.device}'.")
        
        self.loaded_model_name = None
        self.model = None
        try:
            self._load_whisper_model(kwargs.get("model_name", "turbo")) # Default to "turbo" if not specified
        except Exception as e:
            pass


    def _get_whisper_cache_dir(self) -> Path:
        """Get the Whisper cache directory."""
        # Whisper uses the same cache location as specified in its code
        cache_dir = os.getenv("XDG_CACHE_HOME", os.path.expanduser("~/.cache"))
        return Path(cache_dir) / "whisper"

    def _is_model_downloaded(self, model_name: str) -> bool:
        """Check if a Whisper model is already downloaded and valid."""
        cache_dir = self._get_whisper_cache_dir()
        # Whisper models are stored as {model_name}.pt
        model_file = cache_dir / f"{model_name}.pt"
        
        if not model_file.exists():
            return False
        
        # Check if file size is reasonable (at least 1MB to catch partial downloads)
        try:
            file_size = model_file.stat().st_size
            if file_size < 1_000_000:  # Less than 1MB is likely corrupted
                ASCIIColors.warning(f"Model file {model_file} appears corrupted (size: {file_size} bytes). Will re-download.")
                return False
        except Exception as e:
            ASCIIColors.warning(f"Could not check model file size: {e}")
            return False
        
        return True

    def _load_whisper_model(self, model_name_to_load: str):
        """Loads or reloads the Whisper model with file locking to prevent concurrent downloads."""
        if model_name_to_load not in whisper.available_models():
            ASCIIColors.warning(f"'{model_name_to_load}' is not a standard Whisper model size. Attempting to load anyway. Known sizes: {self.WHISPER_MODEL_SIZES}")

        if self.model is not None and self.loaded_model_name == model_name_to_load:
            ASCIIColors.info(f"Whisper model '{model_name_to_load}' already loaded.")
            return

        # Get the cache directory and create lock file path
        cache_dir = self._get_whisper_cache_dir()
        cache_dir.mkdir(parents=True, exist_ok=True)
        lock_file = cache_dir / f"{model_name_to_load}.lock"
        model_file = cache_dir / f"{model_name_to_load}.pt"

        ASCIIColors.info(f"Loading Whisper model: '{model_name_to_load}' on device '{self.device}'...")
        
        # Use file lock to prevent concurrent downloads
        lock = filelock.FileLock(lock_file, timeout=300)  # 5 minute timeout
        
        try:
            with lock:
                # Check if model is already downloaded and valid
                if self._is_model_downloaded(model_name_to_load):
                    ASCIIColors.info(f"Model '{model_name_to_load}' already downloaded, loading from cache...")
                else:
                    ASCIIColors.info(f"Downloading model '{model_name_to_load}' (this process has the lock)...")
                    # If model file exists but is corrupted, delete it first
                    if model_file.exists():
                        try:
                            ASCIIColors.warning(f"Removing corrupted model file: {model_file}")
                            model_file.unlink()
                        except Exception as e:
                            ASCIIColors.error(f"Could not remove corrupted model file: {e}")
                
                # Load the model (will download if not cached)
                self.model = whisper.load_model(model_name_to_load, device=self.device)
                self.loaded_model_name = model_name_to_load
                self.model_name = model_name_to_load
                ASCIIColors.green(f"Whisper model '{model_name_to_load}' loaded successfully.")
                
        except filelock.Timeout:
            error_msg = f"Timeout waiting for model '{model_name_to_load}' download lock. Another process may be downloading."
            ASCIIColors.error(error_msg)
            self.model = None
            self.loaded_model_name = None
            raise RuntimeError(error_msg)
        except Exception as e:
            self.model = None
            self.loaded_model_name = None
            ASCIIColors.error(f"Failed to load Whisper model '{model_name_to_load}': {e}")
            trace_exception(e)
            # If loading failed, try to clean up potentially corrupted file
            if model_file.exists():
                try:
                    ASCIIColors.warning(f"Cleaning up potentially corrupted model file after error...")
                    model_file.unlink()
                except Exception as cleanup_error:
                    ASCIIColors.error(f"Could not clean up model file: {cleanup_error}")
            raise RuntimeError(f"Failed to load Whisper model '{model_name_to_load}'") from e
        finally:
            # Clean up lock file if it exists and is not locked
            try:
                if lock_file.exists() and not lock.is_locked:
                    lock_file.unlink()
            except Exception:
                pass  # Ignore cleanup errors


    def transcribe_audio(self, audio_path: Union[str, Path], model: Optional[str] = None, **kwargs) -> str:
        """
        Transcribes the audio file at the given path using Whisper.

        Args:
            audio_path (Union[str, Path]): The path to the audio file to transcribe.
            model (Optional[str]): The specific Whisper model size to use.
                                  If None, uses the model loaded during initialization.
            **kwargs: Additional parameters for Whisper's transcribe method, e.g.:
                      `language` (str): Language code (e.g., "en", "fr"). If None, Whisper auto-detects.
                      `fp16` (bool): Whether to use fp16, defaults to True if CUDA available.
                      `task` (str): "transcribe" or "translate".

        Returns:
            str: The transcribed text.

        Raises:
            FileNotFoundError: If the audio file does not exist.
            RuntimeError: If the Whisper model is not loaded or transcription fails.
            Exception: For other errors during transcription.
        """
        if not self.model:
            self._load_whisper_model(self.default_model_name)

            
        audio_file = Path(audio_path)
        if not audio_file.exists():
            raise FileNotFoundError(f"Audio file not found at: {audio_path}")

        if model and model != self.loaded_model_name:
            ASCIIColors.info(f"Switching Whisper model to '{model}' for this transcription.")
            try:
                self._load_whisper_model(model) # Attempt to load the new model
            except RuntimeError as e:
                 # If switching fails, keep using the old model if available, or raise if none loaded
                if self.model is None:
                    raise RuntimeError(f"Failed to switch to Whisper model '{model}' and no model currently loaded.") from e
                else:
                    ASCIIColors.warning(f"Failed to switch to Whisper model '{model}'. Using previously loaded model '{self.loaded_model_name}'. Error: {e}")


        if self.model is None:
            raise RuntimeError("Whisper model is not loaded. Cannot transcribe.")

        # Prepare Whisper-specific options from kwargs
        whisper_options = {}
        if "language" in kwargs:
            whisper_options["language"] = kwargs["language"]
        if "fp16" in kwargs: # Typically handled by device selection, but allow override
            whisper_options["fp16"] = kwargs["fp16"]
        else: # Default fp16 based on device
            whisper_options["fp16"] = (self.device == "cuda")
        if "task" in kwargs: # "transcribe" or "translate"
            whisper_options["task"] = kwargs["task"]


        ASCIIColors.info(f"Transcribing '{audio_file.name}' with Whisper model '{self.loaded_model_name}' (options: {whisper_options})...")
        try:
            result = self.model.transcribe(str(audio_file), **whisper_options)
            transcribed_text = result.get("text", "")
            ASCIIColors.green("Transcription successful.")
            return transcribed_text.strip()
        except Exception as e:
            ASCIIColors.error(f"Whisper transcription failed for '{audio_file.name}': {e}")
            trace_exception(e)
            raise Exception(f"Whisper transcription error: {e}") from e

    @staticmethod
    def list_models(**kwargs) -> List[str]:
        """
        Lists the available standard Whisper model sizes.

        Args:
            **kwargs: Additional parameters (currently unused).

        Returns:
            List[str]: A list of available Whisper model size identifiers.
        """
        return whisper.available_models() # Return a copy

    def __del__(self):
        """Clean up: Unload the model to free resources."""
        if self.model is not None:
            del self.model
            self.model = None
            if torch and hasattr(torch, 'cuda') and torch.cuda.is_available():
                torch.cuda.empty_cache()
            ASCIIColors.info(f"WhisperSTTBinding for model '{self.loaded_model_name}' destroyed and resources released.")


# --- Main Test Block (Example Usage) ---
if __name__ == '__main__':
    if not _whisper_installed:
        print(f"{ASCIIColors.RED}Whisper dependencies not met. Skipping tests. Error: {_whisper_installation_error}{ASCIIColors.RESET}")
        print(f"{ASCIIColors.YELLOW}Please ensure 'openai-whisper' and 'torch' are installed, and 'ffmpeg' is in your system PATH.{ASCIIColors.RESET}")
        exit()

    ASCIIColors.yellow("--- WhisperSTTBinding Test ---")
    
    # --- Prerequisites for testing ---
    # 1. Create a dummy WAV file for testing, or provide a path to a real one.
    #    You'll need `scipy` to create a dummy WAV easily, or use an external tool.
    #    Let's assume a simple way to signal a missing test file.
    test_audio_file = Path("test_audio_for_whisper.wav")
    
    # Try to create a dummy file if it doesn't exist (requires scipy)
    if not test_audio_file.exists():
        try:
            import numpy as np
            from scipy.io.wavfile import write as write_wav
            samplerate = 44100; fs = 100
            t = np.linspace(0., 1., samplerate)
            amplitude = np.iinfo(np.int16).max
            data = amplitude * np.sin(2. * np.pi * fs * t)
            write_wav(test_audio_file, samplerate, data.astype(np.int16))
            ASCIIColors.green(f"Created dummy audio file: {test_audio_file}")
        except ImportError:
            ASCIIColors.warning(f"SciPy not installed. Cannot create dummy audio file.")
            ASCIIColors.warning(f"Please place a '{test_audio_file.name}' in the current directory or modify the path.")
        except Exception as e_dummy_audio:
            ASCIIColors.error(f"Could not create dummy audio file: {e_dummy_audio}")


    if not test_audio_file.exists():
        ASCIIColors.error(f"Test audio file '{test_audio_file}' not found. Skipping transcription test.")
    else:
        try:
            ASCIIColors.cyan("\n--- Initializing WhisperSTTBinding (model: 'tiny') ---")
            # Using 'tiny' model for faster testing. Change to 'turbo' or 'small' for better quality.
            stt_binding = WhisperSTTBinding(model_name="tiny")

            ASCIIColors.cyan("\n--- Listing available Whisper models ---")
            models = stt_binding.list_models()
            print(f"Available models: {models}")

            ASCIIColors.cyan(f"\n--- Transcribing '{test_audio_file.name}' with 'tiny' model ---")
            transcription = stt_binding.transcribe_audio(test_audio_file)
            print(f"Transcription (tiny): '{transcription}'")

            # Test with a specific language hint (if your audio is not English or for robustness)
            # ASCIIColors.cyan(f"\n--- Transcribing '{test_audio_file.name}' with 'tiny' model and language hint 'en' ---")
            # transcription_lang_hint = stt_binding.transcribe_audio(test_audio_file, language="en")
            # print(f"Transcription (tiny, lang='en'): '{transcription_lang_hint}'")

            # Test switching model dynamically (optional, will re-download/load if different)
            # ASCIIColors.cyan(f"\n--- Transcribing '{test_audio_file.name}' by switching to 'turbo' model ---")
            # transcription_base = stt_binding.transcribe_audio(test_audio_file, model="turbo")
            # print(f"Transcription (turbo): '{transcription_base}'")


        except ImportError as e_imp:
            ASCIIColors.error(f"Import error during test: {e_imp}")
            ASCIIColors.info("This might be due to `openai-whisper` or `torch` not being installed correctly.")
        except FileNotFoundError as e_fnf:
            ASCIIColors.error(f"File not found during test: {e_fnf}")
        except RuntimeError as e_rt:
            ASCIIColors.error(f"Runtime error during test (often model load or ffmpeg issue): {e_rt}")
            if "ffmpeg" in str(e_rt).lower():
                ASCIIColors.yellow("This error often means 'ffmpeg' is not installed or not found in your system's PATH.")
                ASCIIColors.yellow("Please install ffmpeg: https://ffmpeg.org/download.html")
        except Exception as e:
            ASCIIColors.error(f"An unexpected error occurred during testing: {e}")
            trace_exception(e)
        finally:
            # Clean up dummy audio file if we created it for this test
            # (Be careful if you are using a real test_audio_file you want to keep)
            # if "samplerate" in locals() and test_audio_file.exists(): # Simple check if we likely created it
            #     try:
            #         os.remove(test_audio_file)
            #         ASCIIColors.info(f"Removed dummy audio file: {test_audio_file}")
            #     except Exception as e_del:
            #         ASCIIColors.warning(f"Could not remove dummy audio file {test_audio_file}: {e_del}")
            pass # For this example, let's not auto-delete. User can manage it.


    ASCIIColors.yellow("\n--- WhisperSTTBinding Test Finished ---")