import os
import io
import base64
import tempfile
from flask import Flask, request, jsonify, send_file
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import torchaudio
import numpy as np
from typing import Optional, Dict, Any
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = Flask(__name__)

# Track app start time for uptime calculation
import time
app.start_time = time.time()

class SopranoTTS:
    def __init__(self, model_path: str = "Gaston895/aegis001", device: str = "auto"):
        """Initialize Soprano TTS model from Hugging Face repository"""
        self.device = self._get_device(device)
        logger.info(f"Loading model from {model_path} on device: {self.device}")
        
        try:
            # Load tokenizer and model from Hugging Face
            logger.info("Loading tokenizer...")
            self.tokenizer = AutoTokenizer.from_pretrained(
                model_path,
                trust_remote_code=True,
                use_fast=False
            )
            
            logger.info("Loading model...")
            self.model = AutoModelForCausalLM.from_pretrained(
                model_path,
                torch_dtype=torch.bfloat16 if self.device != "cpu" else torch.float32,
                device_map=self.device if self.device != "cpu" else None,
                trust_remote_code=True,
                low_cpu_mem_usage=True
            )
            
            if self.device == "cpu":
                self.model = self.model.to(self.device)
            
            self.model.eval()
            logger.info("Model loaded successfully")
            
        except Exception as e:
            logger.error(f"Error loading model: {e}")
            raise
    
    def _get_device(self, device: str) -> str:
        """Determine the best device to use"""
        if device == "auto":
            if torch.cuda.is_available():
                return "cuda"
            elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
                return "mps"
            else:
                return "cpu"
        return device
    
    def generate_speech(self, text: str, temperature: float = 0.7, top_p: float = 0.9) -> np.ndarray:
        """Generate speech from text using TTS synthesis"""
        try:
            # Try gTTS first (works better in cloud environments)
            try:
                from gtts import gTTS
                logger.info("Using Google Text-to-Speech (gTTS)")
                return self._generate_with_gtts(text)
            except ImportError:
                logger.warning("gTTS not available, trying pyttsx3")
                pass
            except Exception as e:
                logger.warning(f"gTTS failed: {e}, trying pyttsx3")
                pass
            
            # Try pyttsx3 as fallback
            try:
                import pyttsx3
                logger.info("Using pyttsx3 for text-to-speech synthesis")
                return self._generate_with_pyttsx3(text, temperature, top_p)
            except ImportError:
                logger.warning("pyttsx3 not available, using language model approach")
                pass
            except Exception as e:
                logger.warning(f"pyttsx3 failed: {e}, using language model approach")
                pass
            
            # Final fallback to language model approach
            logger.info("Using enhanced language model speech synthesis")
            return self._generate_with_language_model(text, temperature, top_p)
            
        except Exception as e:
            logger.error(f"Error generating speech: {e}")
            raise
    
    def _generate_with_pyttsx3(self, text: str, temperature: float, top_p: float) -> np.ndarray:
        """Generate speech using pyttsx3"""
        import pyttsx3
        import tempfile
        import soundfile as sf
        
        # Initialize TTS engine
        engine = pyttsx3.init()
        
        # Set voice properties based on temperature
        voices = engine.getProperty('voices')
        if voices:
            # Use temperature to select voice characteristics
            voice_index = int(temperature * len(voices)) % len(voices)
            engine.setProperty('voice', voices[voice_index].id)
        
        # Set speech rate based on top_p
        rate = engine.getProperty('rate')
        new_rate = int(rate * (0.5 + top_p * 0.8))  # Vary rate based on top_p
        engine.setProperty('rate', new_rate)
        
        # Set volume
        engine.setProperty('volume', 0.8)
        
        # Generate speech to temporary WAV file first
        with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as tmp_wav:
            engine.save_to_file(text, tmp_wav.name)
            engine.runAndWait()
            
            # Load the generated audio
            try:
                audio_data, sample_rate = sf.read(tmp_wav.name)
                os.unlink(tmp_wav.name)
                
                # Ensure mono audio
                if len(audio_data.shape) > 1:
                    audio_data = np.mean(audio_data, axis=1)
                
                # Resample to 22kHz for better MP3 compatibility
                if sample_rate != 22050:
                    import librosa
                    audio_data = librosa.resample(audio_data, orig_sr=sample_rate, target_sr=22050)
                
                return audio_data.astype(np.float32)
                
            except Exception as e:
                os.unlink(tmp_wav.name)
                logger.error(f"Error loading pyttsx3 audio: {e}")
                raise
    
    def _generate_with_gtts(self, text: str) -> np.ndarray:
        """Generate speech using Google Text-to-Speech"""
        from gtts import gTTS
        import tempfile
        import soundfile as sf
        
        try:
            # Generate speech with gTTS
            logger.info(f"Generating speech with gTTS for text: {text[:50]}...")
            tts = gTTS(text=text, lang='en', slow=False)
            
            with tempfile.NamedTemporaryFile(suffix='.mp3', delete=False) as tmp_file:
                tts.save(tmp_file.name)
                
                try:
                    # Convert MP3 to audio array and load
                    audio_data, sample_rate = sf.read(tmp_file.name)
                    os.unlink(tmp_file.name)
                    
                    # Ensure mono audio
                    if len(audio_data.shape) > 1:
                        audio_data = np.mean(audio_data, axis=1)
                    
                    # Resample to 22kHz for better MP3 compatibility
                    if sample_rate != 22050:
                        try:
                            import librosa
                            audio_data = librosa.resample(audio_data, orig_sr=sample_rate, target_sr=22050)
                        except ImportError:
                            # If librosa not available, keep original sample rate
                            logger.warning("librosa not available, keeping original sample rate")
                    
                    logger.info(f"gTTS synthesis successful, duration: {len(audio_data)/22050:.2f}s")
                    return audio_data.astype(np.float32)
                    
                except Exception as e:
                    if os.path.exists(tmp_file.name):
                        os.unlink(tmp_file.name)
                    logger.error(f"Error loading gTTS audio: {e}")
                    raise
                    
        except Exception as e:
            logger.error(f"gTTS synthesis failed: {e}")
            raise
    
    def _generate_with_language_model(self, text: str, temperature: float, top_p: float) -> np.ndarray:
        """Enhanced speech-like synthesis using the language model"""
        try:
            logger.info(f"Using enhanced language model synthesis for: {text[:50]}...")
            
            # Tokenize input text
            inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
            inputs = {k: v.to(self.device) for k, v in inputs.items()}
            
            # Generate with the model
            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    max_length=min(1024, inputs['input_ids'].shape[1] + 200),
                    temperature=temperature,
                    top_p=top_p,
                    do_sample=True,
                    pad_token_id=self.tokenizer.eos_token_id
                )
            
            # Convert output tokens to speech-like audio
            audio_tokens = outputs[0][inputs['input_ids'].shape[1]:]
            
            # Create more realistic speech-like audio from tokens
            sample_rate = 22050
            duration = max(2.0, len(text) * 0.08)  # More realistic duration
            num_samples = int(sample_rate * duration)
            
            # Generate speech-like waveform from tokens
            audio_data = np.zeros(num_samples, dtype=np.float32)
            
            # Analyze text for phonetic patterns
            words = text.lower().split()
            
            # Use token values to create more realistic speech patterns
            for i, token_id in enumerate(audio_tokens[:min(100, len(audio_tokens))]):
                token_val = float(token_id.item())
                
                # Create formant frequencies based on token values and text analysis
                # Simulate human speech formants
                f0 = 120 + (token_val % 80)  # Fundamental frequency (pitch)
                f1 = 300 + (token_val % 400)  # First formant (vowel quality)
                f2 = 900 + (token_val % 800)  # Second formant (vowel quality)
                f3 = 2200 + (token_val % 600)  # Third formant (consonant quality)
                
                # Time segment for this token
                start_idx = int(i * num_samples / len(audio_tokens))
                end_idx = int((i + 1) * num_samples / len(audio_tokens))
                
                if start_idx < num_samples and end_idx <= num_samples:
                    t_segment = np.linspace(0, (end_idx - start_idx) / sample_rate, end_idx - start_idx)
                    
                    # Create more realistic formant-based audio segment
                    segment = (
                        0.4 * np.sin(2 * np.pi * f0 * t_segment) +  # Fundamental
                        0.3 * np.sin(2 * np.pi * f1 * t_segment) +  # First formant
                        0.2 * np.sin(2 * np.pi * f2 * t_segment) +  # Second formant
                        0.1 * np.sin(2 * np.pi * f3 * t_segment)    # Third formant
                    )
                    
                    # Apply realistic speech envelope
                    attack_time = 0.02  # 20ms attack
                    decay_time = 0.05   # 50ms decay
                    
                    attack_samples = int(attack_time * sample_rate)
                    decay_samples = int(decay_time * sample_rate)
                    
                    envelope = np.ones_like(t_segment)
                    
                    # Attack phase
                    if attack_samples > 0 and len(envelope) > attack_samples:
                        envelope[:attack_samples] = np.linspace(0, 1, attack_samples)
                    
                    # Decay phase
                    if decay_samples > 0 and len(envelope) > decay_samples:
                        envelope[-decay_samples:] = np.linspace(1, 0.3, decay_samples)
                    
                    segment *= envelope
                    
                    # Add some realistic noise and breathiness
                    noise = np.random.normal(0, 0.01, len(segment))
                    segment += noise
                    
                    audio_data[start_idx:end_idx] += segment
            
            # Apply overall speech-like processing
            # Normalize
            if np.max(np.abs(audio_data)) > 0:
                audio_data = audio_data / np.max(np.abs(audio_data)) * 0.7
            
            # Apply simple low-pass filter to simulate vocal tract
            try:
                from scipy import signal
                # Simple low-pass filter at 4kHz
                nyquist = sample_rate / 2
                cutoff = 4000 / nyquist
                b, a = signal.butter(4, cutoff, btype='low')
                audio_data = signal.filtfilt(b, a, audio_data)
            except ImportError:
                # If scipy not available, skip filtering
                pass
            
            logger.info(f"Enhanced language model synthesis completed, duration: {len(audio_data)/sample_rate:.2f}s")
            return audio_data
            
        except Exception as e:
            logger.error(f"Error in enhanced language model speech generation: {e}")
            # Final fallback: simple speech-like synthesis
            return self._generate_simple_speech(text)
    
    def _generate_simple_speech(self, text: str) -> np.ndarray:
        """Simple speech-like synthesis as final fallback"""
        sample_rate = 22050
        duration = max(1.0, len(text) * 0.08)
        num_samples = int(sample_rate * duration)
        
        logger.info(f"Using simple speech synthesis fallback for: {text[:50]}...")
        
        # Create more realistic speech patterns
        audio_data = np.zeros(num_samples, dtype=np.float32)
        
        # Analyze text for speech patterns
        words = text.lower().split()
        
        for i, word in enumerate(words):
            # Time segment for this word
            start_idx = int(i * num_samples / len(words))
            end_idx = int((i + 1) * num_samples / len(words))
            
            if start_idx < num_samples and end_idx <= num_samples:
                word_duration = (end_idx - start_idx) / sample_rate
                t_word = np.linspace(0, word_duration, end_idx - start_idx)
                
                # Create word-specific frequencies based on vowels and consonants
                vowel_count = sum(1 for c in word if c in 'aeiou')
                consonant_count = len(word) - vowel_count
                
                # Base frequency varies with word characteristics
                base_freq = 120 + (vowel_count * 30) + (consonant_count * 15)
                
                # Create formant-like structure
                f1 = base_freq
                f2 = base_freq * 2.2
                f3 = base_freq * 3.8
                
                # Generate word audio with formants
                word_audio = (
                    0.5 * np.sin(2 * np.pi * f1 * t_word) +
                    0.3 * np.sin(2 * np.pi * f2 * t_word) +
                    0.2 * np.sin(2 * np.pi * f3 * t_word)
                )
                
                # Apply word envelope (attack, sustain, decay)
                envelope = np.ones_like(t_word)
                attack_samples = max(1, len(t_word) // 15)
                decay_samples = max(1, len(t_word) // 12)
                
                if attack_samples > 0 and len(envelope) > attack_samples:
                    envelope[:attack_samples] = np.linspace(0, 1, attack_samples)
                if decay_samples > 0 and len(envelope) > decay_samples:
                    envelope[-decay_samples:] = np.linspace(1, 0, decay_samples)
                
                word_audio *= envelope
                
                # Add some noise for realism
                noise = np.random.normal(0, 0.02, len(word_audio))
                word_audio += noise
                
                audio_data[start_idx:end_idx] = word_audio
        
        # Normalize
        if np.max(np.abs(audio_data)) > 0:
            audio_data = audio_data / np.max(np.abs(audio_data)) * 0.6
        
        logger.info(f"Simple speech synthesis completed, duration: {len(audio_data)/sample_rate:.2f}s")
        return audio_data

# Initialize the model globally
try:
    tts_model = SopranoTTS(model_path="Gaston895/aegis001")
    logger.info("TTS model initialized successfully")
except Exception as e:
    logger.error(f"Failed to initialize TTS model: {e}")
    tts_model = None

@app.route('/', methods=['GET'])
def home():
    """Health check endpoint"""
    return jsonify({
        "status": "healthy",
        "model": "Soprano TTS",
        "version": "1.0.0",
        "endpoints": {
            "synthesize": "/synthesize",
            "batch_synthesize": "/batch_synthesize",
            "health": "/health",
            "wakeup": "/wakeup"
        }
    })

@app.route('/health', methods=['GET'])
def health():
    """Health check endpoint"""
    model_status = "loaded" if tts_model is not None else "failed"
    return jsonify({
        "status": "healthy",
        "model_status": model_status,
        "device": tts_model.device if tts_model else "unknown"
    })

@app.route('/wakeup', methods=['GET', 'POST'])
def wakeup():
    """Wake-up endpoint to prevent space from sleeping"""
    import time
    import datetime
    
    current_time = datetime.datetime.now().isoformat()
    uptime = time.time() - app.start_time if hasattr(app, 'start_time') else 0
    
    # Log the wake-up call
    logger.info(f"Wake-up signal received at {current_time}")
    
    # Perform a quick model check to ensure it's still loaded
    model_status = "loaded" if tts_model is not None else "failed"
    
    # Optional: Generate a very short test audio to keep the model warm
    test_audio_generated = False
    if tts_model is not None:
        try:
            # Generate a very short test phrase to keep the model active
            test_audio = tts_model.generate_speech("Wake up", temperature=0.7, top_p=0.9)
            test_audio_generated = len(test_audio) > 0
            logger.info("Model warm-up test completed successfully")
        except Exception as e:
            logger.warning(f"Model warm-up test failed: {e}")
    
    return jsonify({
        "status": "awake",
        "message": "Space is now awake and ready",
        "timestamp": current_time,
        "uptime_seconds": round(uptime, 2),
        "model_status": model_status,
        "model_warmed": test_audio_generated,
        "device": tts_model.device if tts_model else "unknown",
        "endpoints": {
            "synthesize": "/synthesize",
            "batch_synthesize": "/batch_synthesize",
            "health": "/health",
            "wakeup": "/wakeup"
        }
    })

@app.route('/synthesize', methods=['POST'])
def synthesize():
    """Text-to-speech synthesis endpoint"""
    if tts_model is None:
        return jsonify({"error": "Model not loaded"}), 500
    
    try:
        data = request.get_json()
        if not data or 'text' not in data:
            return jsonify({"error": "Missing 'text' field in request"}), 400
        
        text = data['text']
        if not text.strip():
            return jsonify({"error": "Text cannot be empty"}), 400
        
        # Optional parameters
        temperature = data.get('temperature', 0.7)
        top_p = data.get('top_p', 0.9)
        output_format = data.get('format', 'mp3')  # mp3 or base64
        
        # Validate parameters
        if not 0.1 <= temperature <= 2.0:
            return jsonify({"error": "Temperature must be between 0.1 and 2.0"}), 400
        if not 0.1 <= top_p <= 1.0:
            return jsonify({"error": "Top_p must be between 0.1 and 1.0"}), 400
        
        logger.info(f"Synthesizing text: {text[:50]}...")
        
        # Generate speech
        audio_data = tts_model.generate_speech(text, temperature, top_p)
        
        # Convert to audio file
        sample_rate = 22050  # Better for MP3 compression
        
        if output_format == 'base64':
            # Return as base64 encoded MP3 audio
            with tempfile.NamedTemporaryFile(suffix='.mp3', delete=False) as tmp_file:
                # Convert to MP3 using pydub
                try:
                    from pydub import AudioSegment
                    import io
                    
                    # Convert numpy array to audio segment
                    audio_int16 = (audio_data * 32767).astype(np.int16)
                    audio_segment = AudioSegment(
                        audio_int16.tobytes(),
                        frame_rate=sample_rate,
                        sample_width=2,
                        channels=1
                    )
                    
                    # Export as MP3
                    audio_segment.export(tmp_file.name, format="mp3", bitrate="128k")
                    
                except ImportError:
                    # Fallback to soundfile with WAV if pydub not available
                    import soundfile as sf
                    tmp_file.name = tmp_file.name.replace('.mp3', '.wav')
                    sf.write(tmp_file.name, audio_data, sample_rate)
                
                with open(tmp_file.name, 'rb') as f:
                    audio_bytes = f.read()
                
                os.unlink(tmp_file.name)
                
                audio_b64 = base64.b64encode(audio_bytes).decode('utf-8')
                
                return jsonify({
                    "success": True,
                    "audio_base64": audio_b64,
                    "sample_rate": sample_rate,
                    "duration": len(audio_data) / sample_rate,
                    "text": text,
                    "format": "mp3" if "mp3" in tmp_file.name else "wav"
                })
        
        else:
            # Return as MP3 file
            with tempfile.NamedTemporaryFile(suffix='.mp3', delete=False) as tmp_file:
                try:
                    from pydub import AudioSegment
                    
                    # Convert numpy array to audio segment
                    audio_int16 = (audio_data * 32767).astype(np.int16)
                    audio_segment = AudioSegment(
                        audio_int16.tobytes(),
                        frame_rate=sample_rate,
                        sample_width=2,
                        channels=1
                    )
                    
                    # Export as MP3
                    audio_segment.export(tmp_file.name, format="mp3", bitrate="128k")
                    
                    return send_file(
                        tmp_file.name,
                        mimetype='audio/mpeg',
                        as_attachment=True,
                        download_name=f'speech_{hash(text)}.mp3'
                    )
                    
                except ImportError:
                    # Fallback to WAV if pydub not available
                    import soundfile as sf
                    tmp_file.name = tmp_file.name.replace('.mp3', '.wav')
                    sf.write(tmp_file.name, audio_data, sample_rate)
                    
                    return send_file(
                        tmp_file.name,
                        mimetype='audio/wav',
                        as_attachment=True,
                        download_name=f'speech_{hash(text)}.wav'
                    )
    
    except Exception as e:
        logger.error(f"Error in synthesis: {e}")
        return jsonify({"error": f"Synthesis failed: {str(e)}"}), 500

@app.route('/batch_synthesize', methods=['POST'])
def batch_synthesize():
    """Batch text-to-speech synthesis endpoint"""
    if tts_model is None:
        return jsonify({"error": "Model not loaded"}), 500
    
    try:
        data = request.get_json()
        if not data or 'texts' not in data:
            return jsonify({"error": "Missing 'texts' field in request"}), 400
        
        texts = data['texts']
        if not isinstance(texts, list) or len(texts) == 0:
            return jsonify({"error": "Texts must be a non-empty list"}), 400
        
        if len(texts) > 10:  # Limit batch size
            return jsonify({"error": "Maximum 10 texts per batch"}), 400
        
        # Optional parameters
        temperature = data.get('temperature', 0.7)
        top_p = data.get('top_p', 0.9)
        
        results = []
        
        for i, text in enumerate(texts):
            if not text.strip():
                results.append({"error": f"Text {i} is empty"})
                continue
            
            try:
                logger.info(f"Synthesizing batch text {i+1}/{len(texts)}: {text[:30]}...")
                audio_data = tts_model.generate_speech(text, temperature, top_p)
                
                # Convert to base64 MP3
                with tempfile.NamedTemporaryFile(suffix='.mp3', delete=False) as tmp_file:
                    try:
                        from pydub import AudioSegment
                        
                        # Convert numpy array to audio segment
                        audio_int16 = (audio_data * 32767).astype(np.int16)
                        audio_segment = AudioSegment(
                            audio_int16.tobytes(),
                            frame_rate=22050,
                            sample_width=2,
                            channels=1
                        )
                        
                        # Export as MP3
                        audio_segment.export(tmp_file.name, format="mp3", bitrate="128k")
                        
                    except ImportError:
                        # Fallback to WAV if pydub not available
                        import soundfile as sf
                        tmp_file.name = tmp_file.name.replace('.mp3', '.wav')
                        sf.write(tmp_file.name, audio_data, 22050)
                    
                    with open(tmp_file.name, 'rb') as f:
                        audio_bytes = f.read()
                    
                    os.unlink(tmp_file.name)
                    
                    audio_b64 = base64.b64encode(audio_bytes).decode('utf-8')
                    
                    results.append({
                        "success": True,
                        "audio_base64": audio_b64,
                        "text": text,
                        "duration": len(audio_data) / 22050,
                        "format": "mp3" if "mp3" in tmp_file.name else "wav"
                    })
            
            except Exception as e:
                logger.error(f"Error synthesizing text {i}: {e}")
                results.append({"error": f"Failed to synthesize text {i}: {str(e)}"})
        
        return jsonify({
            "success": True,
            "results": results,
            "sample_rate": 22050,
            "format": "mp3"
        })
    
    except Exception as e:
        logger.error(f"Error in batch synthesis: {e}")
        return jsonify({"error": f"Batch synthesis failed: {str(e)}"}), 500

if __name__ == '__main__':
    port = int(os.environ.get('PORT', 7860))
    app.run(host='0.0.0.0', port=port, debug=False)