import os
import io
import time
import torch
import librosa
import requests
import tempfile
import threading
import numpy as np
import soundfile as sf
import gradio as gr
from transformers import AutoModel, logging as trf_logging
from huggingface_hub import login, hf_hub_download, scan_cache_dir

# Enable verbose logging for transformers
trf_logging.set_verbosity_info()

# Login (optional)
hf_token = os.getenv("HF_TOKEN")
if hf_token:
    print("🔐 Logging into Hugging Face with token...")
    login(token=hf_token)
else:
    print("⚠️ HF_TOKEN not found. Proceeding without login...")

# Load model with GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"🔧 Using device: {device}")

# Initialize model variable
model = None

# Define the repository ID
repo_id = "ai4bharat/IndicF5"

# Improved model loading with error handling
try:
    print(f"Loading {repo_id} model...")
    # Try direct loading first
    model = AutoModel.from_pretrained(
        repo_id,
        trust_remote_code=True,
        revision="main"
    ).to(device)
    print(f"Model loaded successfully! Type: {type(model)}")
    
    # Check model attributes
    model_methods = [method for method in dir(model) if not method.startswith('_') and callable(getattr(model, method))]
    print(f"Available model methods: {model_methods[:10]}...")
        
except Exception as e:
    print(f"⚠️ Error loading model directly: {e}")
    
    try:
        # Try loading with local_files_only if model is cached
        model = AutoModel.from_pretrained(
            repo_id,
            trust_remote_code=True,
            local_files_only=True
        ).to(device)
        print("Model loaded from cache!")
    except Exception as e2:
        print(f"❌ All attempts to load model failed: {e2}")

# Advanced audio processing functions
def remove_noise(audio_data, threshold=0.01):
    """Apply simple noise gate to remove low-level noise"""
    if audio_data is None:
        return np.zeros(1000)
    
    # Convert to numpy if needed
    if isinstance(audio_data, torch.Tensor):
        audio_data = audio_data.detach().cpu().numpy()
    if isinstance(audio_data, list):
        audio_data = np.array(audio_data)
    
    # Apply noise gate
    noise_mask = np.abs(audio_data) < threshold
    clean_audio = audio_data.copy()
    clean_audio[noise_mask] = 0
    
    return clean_audio

def apply_smoothing(audio_data, window_size=5):
    """Apply gentle smoothing to reduce artifacts"""
    if audio_data is None or len(audio_data) < window_size*2:
        return audio_data
    
    # Simple moving average filter
    kernel = np.ones(window_size) / window_size
    smoothed = np.convolve(audio_data, kernel, mode='same')
    
    # Keep original at the edges
    smoothed[:window_size] = audio_data[:window_size]
    smoothed[-window_size:] = audio_data[-window_size:]
    
    return smoothed

def enhance_audio(audio_data):
    """Process audio to improve quality and reduce noise"""
    if audio_data is None:
        return np.zeros(1000)
    
    # Ensure numpy array
    if isinstance(audio_data, torch.Tensor):
        audio_data = audio_data.detach().cpu().numpy()
    if isinstance(audio_data, list):
        audio_data = np.array(audio_data)
    
    # Ensure correct shape and dtype
    if len(audio_data.shape) > 1:
        audio_data = audio_data.flatten()
    if audio_data.dtype != np.float32:
        audio_data = audio_data.astype(np.float32)
    
    # Skip processing if audio is empty or too short
    if audio_data.size < 100:
        return audio_data
    
    # Check if the audio has reasonable amplitude
    rms = np.sqrt(np.mean(audio_data**2))
    print(f"Initial RMS: {rms}")
    
    # Apply gain if needed
    if rms < 0.05:  # Very quiet
        target_rms = 0.2
        gain = target_rms / max(rms, 0.0001)
        print(f"Applying gain factor: {gain}")
        audio_data = audio_data * gain
    
    # Remove DC offset
    audio_data = audio_data - np.mean(audio_data)
    
    # Apply noise gate to remove low-level noise
    audio_data = remove_noise(audio_data, threshold=0.01)
    
    # Apply gentle smoothing to reduce artifacts
    audio_data = apply_smoothing(audio_data, window_size=3)
    
    # Apply soft limiting to prevent clipping
    max_amp = np.max(np.abs(audio_data))
    if max_amp > 0.95:
        audio_data = 0.95 * audio_data / max_amp
    
    # Apply subtle compression for better audibility
    audio_data = np.tanh(audio_data * 1.1) * 0.9
    
    return audio_data

# Load audio from URL with improved error handling
def load_audio_from_url(url):
    print(f"Downloading reference audio from {url}")
    try:
        response = requests.get(url)
        if response.status_code == 200:
            try:
                # Save content to a temp file
                temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.wav')
                temp_file.write(response.content)
                temp_file.close()
                print(f"Saved reference audio to temp file: {temp_file.name}")

                # Try different methods to read the audio file
                audio_data = None
                sample_rate = None
                
                # Try SoundFile first
                try:
                    audio_data, sample_rate = sf.read(temp_file.name)
                    print(f"Audio loaded with SoundFile: {sample_rate}Hz, {len(audio_data)} samples")
                except Exception as sf_error:
                    print(f"SoundFile failed: {sf_error}")
                    
                    # Try librosa as fallback
                    try:
                        audio_data, sample_rate = librosa.load(temp_file.name, sr=None)
                        print(f"Audio loaded with librosa: {sample_rate}Hz, shape={audio_data.shape}")
                    except Exception as lr_error:
                        print(f"Librosa also failed: {lr_error}")
                
                # Clean up temp file
                os.unlink(temp_file.name)
                
                if audio_data is not None:
                    # Apply audio enhancement to the reference
                    audio_data = enhance_audio(audio_data)
                    return sample_rate, audio_data
                    
            except Exception as e:
                print(f"Failed to process audio data: {e}")
        else:
            print(f"Failed to download audio: status code {response.status_code}")
    except Exception as e:
        print(f"Error downloading audio: {e}")

    # Return default values as fallback
    print("⚠️ Returning default silence as reference audio")
    return 24000, np.zeros(int(24000))  # 1 second of silence at 24kHz

# Split text into chunks for streaming
def split_into_chunks(text, max_length=30):
    """Split text into smaller chunks based on punctuation and length"""
    # First split by sentences
    sentence_markers = ['.', '?', '!', ';', ':', '।', '॥']
    chunks = []
    current = ""

    # Initial coarse splitting by sentence markers
    for char in text:
        current += char
        if char in sentence_markers and current.strip():
            chunks.append(current.strip())
            current = ""

    if current.strip():
        chunks.append(current.strip())

    # Further break down long sentences
    final_chunks = []
    for chunk in chunks:
        if len(chunk) <= max_length:
            final_chunks.append(chunk)
        else:
            # Try splitting by commas for long sentences
            comma_splits = chunk.split(',')
            current_part = ""

            for part in comma_splits:
                if len(current_part) + len(part) <= max_length:
                    if current_part:
                        current_part += ","
                    current_part += part
                else:
                    if current_part:
                        final_chunks.append(current_part.strip())
                    current_part = part

            if current_part:
                final_chunks.append(current_part.strip())

    print(f"Split text into {len(final_chunks)} chunks")
    return final_chunks

# Improved model wrapper
class ModelWrapper:
    def __init__(self, model):
        self.model = model
        print(f"Model wrapper initialized with model type: {type(model)}")
        
        # Discover the appropriate generation method
        self.generation_method = self._find_generation_method()
        
    def _find_generation_method(self):
        """Find the appropriate method to generate speech"""
        if self.model is None:
            return None
            
        # Look for plausible generation methods
        candidates = [
            "generate_speech", "tts", "generate_audio", "synthesize", 
            "generate", "forward", "__call__"
        ]
        
        # Check for methods containing these keywords
        for name in dir(self.model):
            if any(candidate in name.lower() for candidate in candidates):
                print(f"Found potential generation method: {name}")
                return name
                
        # If nothing specific found, default to __call__
        print("No specific generation method found, will use __call__")
        return "__call__"

    def generate(self, text, ref_audio_path, ref_text, **kwargs):
        """Generate speech with improved error handling and preprocessing"""
        print(f"\n==== MODEL INFERENCE ====")
        print(f"Text input: '{text}'")
        print(f"Reference audio path: {ref_audio_path}")

        # Check if files exist
        if not os.path.exists(ref_audio_path):
            print(f"⚠️ Reference audio file not found")
            return None

        # Try different calling approaches
        result = None
        method_name = self.generation_method if self.generation_method else "__call__"
        
        # Set up different parameter combinations to try
        param_combinations = [
            # First try: standard keyword parameters
            {"text": text, "ref_audio_path": ref_audio_path, "ref_text": ref_text},
            # Second try: alternative parameter names
            {"text": text, "reference_audio": ref_audio_path, "speaker_text": ref_text},
            # Third try: just text and audio
            {"text": text, "reference_audio": ref_audio_path},
            # Fourth try: just text
            {"text": text},
            # Fifth try: positional arguments
            {}  # Will use positional below
        ]
        
        # Try each parameter combination
        for i, params in enumerate(param_combinations):
            try:
                method = getattr(self.model, method_name)
                print(f"Attempt {i+1}: Calling model.{method_name} with {list(params.keys())} parameters")
                
                # For the positional arguments case
                if not params:
                    result = method(text, ref_audio_path, ref_text, **kwargs)
                else:
                    result = method(**params, **kwargs)
                    
                print(f"✓ Call succeeded with parameters: {list(params.keys())}")
                break  # Exit loop if successful
                
            except Exception as e:
                print(f"✗ Attempt {i+1} failed: {str(e)[:100]}...")
                continue
        
        # Process the result
        if result is not None:
            # Handle tuple results (might be audio, sample_rate)
            if isinstance(result, tuple):
                result = result[0]  # Extract first element, assuming it's audio
            
            # Convert torch tensor to numpy if needed
            if isinstance(result, torch.Tensor):
                result = result.detach().cpu().numpy()
            
            # Ensure array is 1D
            if hasattr(result, 'shape') and len(result.shape) > 1:
                result = result.flatten()
            
            # Apply advanced audio processing to improve quality
            result = enhance_audio(result)
            
            return result
        else:
            print("❌ All inference attempts failed")
            return np.zeros(int(24000))  # Return 1 second of silence as fallback

# Create model wrapper
model_wrapper = ModelWrapper(model) if model is not None else None

# Streaming TTS class with improved audio quality
class StreamingTTS:
    def __init__(self):
        self.is_generating = False
        self.should_stop = False
        self.temp_dir = None
        self.ref_audio_path = None
        self.output_file = None
        self.all_chunks = []
        self.sample_rate = 24000  # Default sample rate

        # Create temp directory
        self.temp_dir = tempfile.mkdtemp()
        print(f"Created temp directory: {self.temp_dir}")

    def prepare_ref_audio(self, ref_audio, ref_sr):
        """Prepare reference audio with enhanced quality"""
        try:
            if self.ref_audio_path is None:
                self.ref_audio_path = os.path.join(self.temp_dir, "ref_audio.wav")

                # Process the reference audio to ensure clean quality
                ref_audio = enhance_audio(ref_audio)
                
                # Save the reference audio
                sf.write(self.ref_audio_path, ref_audio, ref_sr, format='WAV', subtype='FLOAT')
                print(f"Saved reference audio to: {self.ref_audio_path}")

                # Verify file was created
                if os.path.exists(self.ref_audio_path):
                    print(f"Reference audio saved successfully: {os.path.getsize(self.ref_audio_path)} bytes")
                else:
                    print("⚠️ Failed to create reference audio file!")

            # Create output file
            if self.output_file is None:
                self.output_file = os.path.join(self.temp_dir, "output.wav")
                print(f"Output will be saved to: {self.output_file}")
        except Exception as e:
            print(f"Error preparing reference audio: {e}")

    def cleanup(self):
        """Clean up temporary files"""
        if self.temp_dir:
            try:
                if os.path.exists(self.ref_audio_path):
                    os.remove(self.ref_audio_path)
                if os.path.exists(self.output_file):
                    os.remove(self.output_file)
                os.rmdir(self.temp_dir)
                self.temp_dir = None
                print("Cleaned up temporary files")
            except Exception as e:
                print(f"Error cleaning up: {e}")

    def generate(self, text, ref_audio, ref_sr, ref_text):
        """Start generation in a new thread"""
        if self.is_generating:
            print("Already generating speech, please wait")
            return

        # Check model is loaded
        if model_wrapper is None:
            print("⚠️ Model is not loaded. Cannot generate speech.")
            return

        self.is_generating = True
        self.should_stop = False
        self.all_chunks = []

        # Start in a new thread
        threading.Thread(
            target=self._process_streaming,
            args=(text, ref_audio, ref_sr, ref_text),
            daemon=True
        ).start()

    def _process_streaming(self, text, ref_audio, ref_sr, ref_text):
        """Process text in chunks with high-quality audio generation"""
        try:
            # Prepare reference audio
            self.prepare_ref_audio(ref_audio, ref_sr)

            # Split text into smaller chunks for faster processing
            chunks = split_into_chunks(text)
            print(f"Processing {len(chunks)} chunks")

            combined_audio = None
            total_start_time = time.time()

            # Process each chunk
            for i, chunk in enumerate(chunks):
                if self.should_stop:
                    print("Stopping generation as requested")
                    break

                chunk_start = time.time()
                print(f"Processing chunk {i+1}/{len(chunks)}: {chunk}")

                # Generate speech for this chunk
                try:
                    with torch.inference_mode():
                        chunk_audio = model_wrapper.generate(
                            chunk,
                            self.ref_audio_path,
                            ref_text
                        )

                        if chunk_audio is None or (hasattr(chunk_audio, 'size') and chunk_audio.size == 0):
                            print("⚠️ Empty audio returned for this chunk")
                            chunk_audio = np.zeros(int(24000 * 0.5))  # 0.5s silence
                    
                    # Process the audio to improve quality
                    chunk_audio = enhance_audio(chunk_audio)
                    
                    chunk_time = time.time() - chunk_start
                    print(f"✓ Chunk {i+1} processed in {chunk_time:.2f}s")

                    # Add small silence between chunks
                    silence = np.zeros(int(24000 * 0.1))  # 0.1s silence
                    chunk_audio = np.concatenate([chunk_audio, silence])

                    # Add to our collection
                    self.all_chunks.append(chunk_audio)

                    # Combine all chunks so far
                    if combined_audio is None:
                        combined_audio = chunk_audio
                    else:
                        combined_audio = np.concatenate([combined_audio, chunk_audio])

                    # Process combined audio for consistent quality
                    processed_audio = enhance_audio(combined_audio)
                    
                    # Write intermediate output
                    sf.write(self.output_file, processed_audio, 24000, format='WAV', subtype='FLOAT')

                except Exception as e:
                    print(f"Error processing chunk {i+1}: {str(e)[:100]}")
                    continue

            total_time = time.time() - total_start_time
            print(f"Total generation time: {total_time:.2f}s")

        except Exception as e:
            print(f"Error in streaming TTS: {str(e)[:100]}")
        finally:
            self.is_generating = False
            print("Generation complete")

    def get_current_audio(self):
        """Get current audio file path for Gradio"""
        if self.output_file and os.path.exists(self.output_file):
            file_size = os.path.getsize(self.output_file)
            if file_size > 0:
                return self.output_file
        return None

    def stop(self):
        """Stop generation"""
        self.should_stop = True
        print("Stop request received")

# Load reference example (Malayalam)
EXAMPLES = [{
    "audio_url": "https://raw.githubusercontent.com/Aparna0112/voicerecording-_TTS/main/KC%20Voice.wav",
    "ref_text": "ഹലോ ഇത് അപരനെ അല്ലേ ഞാൻ ജഗദീപ് ആണ് വിളിക്കുന്നത് ഇപ്പോൾ ഫ്രീയാണോ സംസാരിക്കാമോ ",
    "synth_text": "ഞാൻ മലയാളം സംസാരിക്കാൻ കഴിയുന്നു."
}]

print("\nPreloading reference audio...")
ref_sr, ref_audio = load_audio_from_url(EXAMPLES[0]["audio_url"])

if ref_audio is None:
    print("⚠️ Failed to load reference audio. Using silence instead.")
    ref_audio = np.zeros(int(24000))
    ref_sr = 24000

# Initialize streaming TTS
streaming_tts = StreamingTTS()

# Add a stop button functionality
def stop_generation():
    streaming_tts.stop()
    return "Generation stopped"

# Gradio interface
with gr.Blocks() as iface:
    gr.Markdown("## 🚀 IndicF5 Malayalam TTS")

    with gr.Row():
        gr.Markdown("### System Status:")
        system_status = gr.Markdown(f"- Device: {device}\n- Model loaded: {'Yes' if model is not None else 'No'}\n- Reference audio: {'Loaded' if ref_audio is not None else 'Not loaded'}")

    with gr.Row():
        text_input = gr.Textbox(
            label="Malayalam Text",
            placeholder="Enter text here...",
            lines=3,
            value=EXAMPLES[0]["synth_text"] if EXAMPLES else "ഹലോ, എന്തൊക്കെ ഉണ്ട് വിശേഷം?"
        )

    with gr.Row():
        generate_btn = gr.Button("🎤 Generate Speech", variant="primary")
        stop_btn = gr.Button("🛑 Stop Generation", variant="secondary")

    # Status indicator
    status_text = gr.Textbox(label="Status", value="Ready", interactive=False)

    # Audio output
    output_audio = gr.Audio(
        label="Generated Speech",
        type="filepath",
        autoplay=True
    )

    # Debug information (hidden by default)
    with gr.Accordion("Advanced", open=False):
        debug_output = gr.Textbox(label="Debug Log", value="", lines=5)

    def start_generation(text):
        if not text.strip():
            return None, "Please enter some text", "Error: Empty text input"

        if model is None:
            return None, "⚠️ Model not loaded. Cannot generate speech.", "Error: Model not loaded"

        if ref_audio is None:
            return None, "⚠️ Reference audio not loaded. Cannot generate speech.", "Error: Reference audio not loaded"

        # Capture stdout for debug purposes
        import io
        from contextlib import redirect_stdout
        f = io.StringIO()
        with redirect_stdout(f):
            streaming_tts.generate(text, ref_audio, ref_sr, EXAMPLES[0]["ref_text"] if EXAMPLES else "")

        debug_log = f.getvalue()

        # Add a delay to ensure file is created
        time.sleep(1.5)

        audio_path = streaming_tts.get_current_audio()
        if audio_path and os.path.exists(audio_path) and os.path.getsize(audio_path) > 0:
            return audio_path, "Generation started - audio playing", debug_log
        else:
            return None, "Starting generation... please wait", debug_log

    generate_btn.click(start_generation, inputs=text_input, outputs=[output_audio, status_text, debug_output])
    stop_btn.click(stop_generation, inputs=None, outputs=status_text)

# Cleanup on exit
def exit_handler():
    streaming_tts.cleanup()

import atexit
atexit.register(exit_handler)

print("Starting Gradio interface...")
iface.launch()