import torch
from chatterbox.tts import ChatterboxTTS
import numpy as np
import os
import tempfile
import sys
import types

# --- Perth Watermarker Monkeypatch ---
# Chatterbox requires resemble-perth, but sometimes 'perth' package is installed instead.
# This patch prevents the AttributeError if the correct attribute is missing.
try:
    import perth
    if not hasattr(perth, 'PerthImplicitWatermarker'):
        class MockWatermarker:
            def apply_watermark(self, wav, sample_rate): return wav
            def get_watermark(self, wav, sample_rate): return 0.0
        perth.PerthImplicitWatermarker = MockWatermarker
except ImportError:
    perth = types.ModuleType("perth")
    class MockWatermarker:
        def apply_watermark(self, wav, sample_rate): return wav
        def get_watermark(self, wav, sample_rate): return 0.0
    perth.PerthImplicitWatermarker = MockWatermarker
    sys.modules["perth"] = perth
# -------------------------------------

class ChatterboxEngine:
    def __init__(self, model_type="voice_clone", device=None):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.model_type = model_type
        print(f"Initializing ChatterboxEngine on {self.device} with model type: {model_type}...")
        
        # Load the Chatterbox model (English only for voice cloning)
        self.model = ChatterboxTTS.from_pretrained(device=self.device)
        
        # Available voices (English only for voice cloning)
        self.voices = {
            "English": ["default", "male_1", "male_2", "female_1", "female_2", "child_1", "narrator", "announcer"]
        }
        
        # Custom voices storage (for cloned voices)
        self.custom_voices = {}
        
    def get_voice_list(self):
        """Get list of all available voices including custom ones."""
        all_voices = []
        for category in self.voices.values():
            all_voices.extend(category)
        all_voices.extend(self.custom_voices.keys())
        return all_voices
    
    def clone_voice(self, audio_path, voice_name):
        """
        Clone a voice from an audio file.
        For ChatterboxTTS, we just store the path for use with audio_prompt_path.
        """
        try:
            # Store the audio path for voice cloning
            self.custom_voices[voice_name] = audio_path
            print(f"Voice '{voice_name}' registered for cloning")
            return voice_name
        except Exception as e:
            print(f"Error cloning voice: {str(e)}")
            return None
    
    def generate(self, text, voice="default", speed=1.0, lang='en', custom_voice_path=None, exaggeration=0.5, cfg_weight=0.5, seed=None, temperature=1.0):
        """
        Generates audio from text using a specified voice.
        """
        try:
            # Check if this is a custom voice
            if voice in self.custom_voices:
                # Use stored audio path for voice cloning
                audio = self.model.generate(text, audio_prompt_path=self.custom_voices[voice])
            elif custom_voice_path:
                # Use provided audio path for one-shot cloning
                audio = self.model.generate(text, audio_prompt_path=custom_voice_path)
            # Ensure audio is a flat numpy array for soundfile compatibility
            if torch.is_tensor(audio):
                audio = audio.detach().cpu().numpy().flatten()
            elif isinstance(audio, np.ndarray):
                audio = audio.flatten()

            sample_rate = self.model.sr
            return audio, sample_rate
            
        except Exception as e:
            print(f"Error generating audio: {str(e)}")
            return None, 22050