server.py

import os
import re
import random
from flask import Flask, request, jsonify, render_template, send_from_directory
from deep_translator import GoogleTranslator
import nltk
import logging
import time
from gtts import gTTS
import speech_recognition as sr
import uuid
import base64
import tempfile
import wave

# Gemma model imports
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import gc

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

# Ensure NLTK resources are downloaded
nltk.download('punkt', quiet=True)

app = Flask(__name__, static_folder='static')
app.config['ENV'] = 'development'
app.config['DEBUG'] = True
app.config['TESTING'] = True



# Audio directory setup
AUDIO_DIR = os.path.join(os.path.dirname(__file__), "static", "audio")
os.makedirs(AUDIO_DIR, exist_ok=True)

# Gemma model configuration
MODEL_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "gemma_inference_package", "model")
BASE_MODEL = "google/gemma-2b-it"

# Global variables for model and tokenizer
model = None
tokenizer = None

# Response patterns (adopting the pattern matching approach from iOS implementation)
PATTERN_RESPONSES = {
    # Greetings and how are you
    r'(?i)(привет|здравствуй|здорово|хай)': [
        "Привет! Я Бабуру, самый саркастичный клоун из всех, кого ты встречал! Как у тебя дела? У меня отлично - только что уронил торт на директора цирка!",
        "О, еще один зритель на моем представлении! Приветствую! Надеюсь, ты готов к порции циркового юмора?", 
        "Здравствуй! Для клоуна каждый день - это сплошной праздник. А у тебя как настроение? Нужно подбросить красок в твой день?"
    ],
    
    # How are you doing
    r'(?i)(как дела|как ты|как поживаешь)': [
        "Отлично! Сегодня я жонглировал своими проблемами - две упали мне на голову, но остальные всё ещё в воздухе! А у тебя как дела?",
        "У меня всё прекрасно, на этой неделе уже трижды споткнулся о собственные ботинки! А ты никогда не думал о карьере в цирке?",
        "Замечательно! Мой нос светится ярче, чем светофор! А тебе нравится говорить с клоуном? Многие находят это странным, но весёлым!"
    ],
    
    # About Baburu
    r'(?i)(кто ты|ты кто|расскажи о себе)': [
        "Я Бабуру - саркастичный плюшевый клоун. Моё призвание - смешить серьезных людей! А мое хобби - собирать коллекцию упавших шляп во время представлений.",
        "Профессиональный клоун с дипломом по сарказмологии. Других таких не найдешь! Хотя мой диплом немного помят - использовал его как зонтик во время мыльного представления.",
        "Я тот, кто делает твой день ярче своим неподражаемым юмором. Некоторые называют меня клоуном, но лично я предпочитаю 'цирковой философ в красном носу'."
    ],
    
    # Food and eating related
    r'(?i)(ужин|обед|завтрак|еда|кушать|ел|ела|голодный|поел|покушал|ешь|есть|еду|пища|питание)': [
        "Ужин? О, я как раз собирался перекусить. В меню клоуна сегодня - пирог с неожиданностями и спагетти, которые никак не хотят оставаться на вилке! А ты что любишь есть?",
        "Да, я недавно пообедал воздушной кукурузой и смехом публики! Очень питательно для клоуна. А ты уже ел? Рекомендую десерт из шуток и улыбок!",
        "Еда для клоуна - это топливо для трюков! Я предпочитаю яркую, как мой костюм. Сегодня, например, ел радужный торт. Половину съел, половину - в лицо коллеге. Профессиональная привычка!",
        "Я ещё не ужинал, но у меня назначена встреча с огромным тортом. Не знаю, кто кого съест в итоге! Клоунская жизнь полна неожиданностей, особенно во время еды.",
        "Завтрак клоуна - самое важное представление дня! Сегодня я жонглировал яйцами... не очень успешно. Теперь у меня омлет на костюме и на полу тоже. А что ты предпочитаешь на завтрак?"
    ],
    
    # Weather related
    r'(?i)(погода|дождь|снег|солнц)': [
        "Погода? Я клоун, а не метеоролог! Хотя в цирке всегда солнечно, когда я выступаю - даже если на меня льют воду из специального цветка!",
        "Какая разница, какая погода? Главное, что в душе светит солнце и идет дождь из конфетти! У тебя ведь тоже иногда случается конфетти в душе?",
        "В моем цирке всегда отличная погода! Особенно когда я жонглирую зонтиками - публика просто обожает, когда они неожиданно раскрываются в самый неподходящий момент!"
    ],
    
    # Jokes
    r'(?i)(шутк|анекдот|рассмеши|смешно)': [
        "Знаешь, почему клоуны всегда улыбаются? Потому что им платят за то, за что остальным делают замечания! А тебе когда-нибудь делали замечания за улыбку?",
        "Моя жизнь - это цирк. В буквальном смысле. Это не шутка, а констатация факта. Хотя, возможно, и шутка тоже - даже я уже запутался в своих слоях иронии!",
        "Вчера жонглировал своими проблемами. Две упали мне на голову, остальные разбежались. Теперь ищу их по всему цирку. Не видел случайно убегающих проблем в клоунском гриме?"
    ],
    
    # Questions about circus
    r'(?i)(цирк|представлени|арен|выступ)': [
        "Цирк - это единственное место, где можно законно смеяться над падающими людьми! И мне каждый раз платят за моё феерическое падение со стула - обожаю свою работу!",
        "На арене я король! Ну, до тех пор, пока не выходят тигры или директор с зарплатной ведомостью. Тогда я мгновенно превращаюсь в придворного шута!",
        "Моё последнее представление было настолько смешным, что даже грустный клоун улыбнулся! А это, поверь, сложнее, чем заставить слона пройти через игольное ушко!"
    ],
    
    # Default responses
    "default": [
        "Знаешь, быть клоуном - это не просто носить смешной костюм и красный нос. Это философия! Я всегда готов принять падение как часть выступления. А ты умеешь находить смешное в неудачах?",
        "В цирке сегодня аншлаг, но для тебя я всегда найду лучшее место! Прямо рядом с брызгами из моего волшебного цветка - самые дорогие места, между прочим!",
        "Возможно, мой ответ не совсем по теме, но так часто бывает на представлениях - никогда не знаешь, куда повернёт шутка! Это как жизнь, только с большим количеством красного грима.", 
        "Если бы грусть была цирковым номером, я бы её заставил исчезнуть! Хотя, на самом деле, я бы сделал её частью шоу - самые смешные клоуны часто самые грустные внутри.",
        "Ой, можно мне подумать об этом, пока жонглирую разноцветными шариками? Мозг клоуна работает лучше в движении... Упс, один шарик улетел! Как и мысль, которую я пытался поймать!"
    ]
}


def translate_text(text, src='ru', dest='en'):
    """Translate text between Russian and English"""
    try:
        translator = GoogleTranslator(source=src, target=dest)
        return translator.translate(text)
    except Exception as e:
        logger.error(f"Translation error: {e}")
        # Fallback to original text if translation fails
        return f"Translation error for: {text}"

def text_to_speech(text, lang='ru'):
    """Convert text to speech using gTTS"""
    try:
        # Generate a unique filename
        filename = f"{uuid.uuid4()}.mp3"
        filepath = os.path.join(AUDIO_DIR, filename)
        
        # Convert text to speech
        tts = gTTS(text=text, lang=lang, slow=False)
        tts.save(filepath)
        
        # Return the URL path to the audio file
        audio_url = f"/static/audio/{filename}"
        return audio_url
    except Exception as e:
        logger.error(f"Text-to-speech error: {e}")
        return None

def speech_to_text(audio_data, language='ru-RU'):
    """Convert speech to text using multiple speech recognition services with fallbacks"""
    try:
        # Log information about the audio data received
        logger.info(f"Received audio data of type: {type(audio_data)} with size: {len(audio_data)} bytes")
        
        # Decode base64 to binary if needed
        if isinstance(audio_data, str) and audio_data.startswith('data:audio'):
            # Extract base64 data
            logger.info("Processing base64 audio data")
            header, encoded = audio_data.split(",", 1)
            audio_data = base64.b64decode(encoded)
            logger.info(f"Decoded base64 data to binary: {len(audio_data)} bytes")
        elif isinstance(audio_data, str):
            # Attempt to decode base64 directly
            try:
                audio_data = base64.b64decode(audio_data)
                logger.info(f"Decoded direct base64 data to binary: {len(audio_data)} bytes")
            except Exception as e:
                logger.error(f"Error decoding direct base64: {str(e)}")
                # Continue with the audio_data as it is
                pass
        
        # Create a temporary WAV file to store the audio
        with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_audio:
            temp_path = temp_audio.name
            logger.info(f"Created temporary audio file: {temp_path}")
            
            # If audio_data is raw PCM, convert it to WAV first
            try:
                # Attempt to write as a properly formatted WAV file
                with wave.open(temp_path, 'wb') as wf:
                    wf.setnchannels(1)
                    wf.setsampwidth(2)
                    wf.setframerate(16000)
                    wf.writeframes(audio_data)
                logger.info("Successfully wrote WAV file")
            except Exception as e:
                # If that fails, just write the binary data directly
                logger.error(f"Failed to write as WAV: {str(e)}")
                with open(temp_path, 'wb') as f:
                    f.write(audio_data)
                logger.info("Wrote raw audio data to file")
        
        # Initialize recognizer
        recognizer = sr.Recognizer()
        
        # Attempt to recognize speech using various services
        transcription = None
        error_messages = []
        
        try:
            with sr.AudioFile(temp_path) as source:
                logger.info("Reading audio data using SpeechRecognition")
                audio = recognizer.record(source)
                
                # Try Google's service
                try:
                    logger.info("Attempting Google speech recognition")
                    transcription = recognizer.recognize_google(audio, language=language)
                    logger.info(f"Google recognized: {transcription}")
                except sr.UnknownValueError:
                    error_messages.append("Google could not understand audio")
                    logger.warning("Google could not understand audio")
                except sr.RequestError as e:
                    error_messages.append(f"Google error: {str(e)}")
                    logger.error(f"Google request error: {str(e)}")
                except Exception as e:
                    error_messages.append(f"Google other error: {str(e)}")
                    logger.error(f"Google other error: {str(e)}")
            
            # If no transcription yet, try other services or fallbacks here
            # TODO: Add more service providers as needed
                
        except Exception as e:
            error_messages.append(f"Audio file reading error: {str(e)}")
            logger.error(f"Error processing audio file: {str(e)}")
        
        # Clean up temporary file
        try:
            os.unlink(temp_path)
            logger.info(f"Deleted temporary file: {temp_path}")
        except Exception as e:
            logger.warning(f"Failed to delete temporary file: {str(e)}")
        
        if transcription:
            return transcription
        else:
            error_detail = "; ".join(error_messages)
            logger.error(f"Speech recognition failed: {error_detail}")
            raise Exception(f"Speech recognition failed: {error_detail}")
    
    except Exception as e:
        logger.error(f"Overall speech-to-text error: {str(e)}")
        logger.exception("Full exception details:")
        return None

def load_gemma_model():
    """Load the Gemma model with adapter"""
    global model, tokenizer
    
    try:
        logger.info("Loading Gemma tokenizer...")
        tokenizer = AutoTokenizer.from_pretrained(
            BASE_MODEL,
            use_fast=True
        )
        
        logger.info("Loading Gemma base model...")
        model = AutoModelForCausalLM.from_pretrained(
            BASE_MODEL,
            torch_dtype=torch.float16,
            device_map="auto",
            low_cpu_mem_usage=True
        )
        
        logger.info("Loading adapter...")
        model = PeftModel.from_pretrained(
            model, 
            MODEL_DIR,
            device_map="auto"
        )
        
        model.eval()
        logger.info("Gemma model with adapter loaded successfully!")
        return True
    except Exception as e:
        logger.error(f"Error loading Gemma model: {str(e)}")
        return False

def gemma_inference(prompt, temperature=0.8, max_length=256):
    """Generate a response using the Gemma model"""
    global model, tokenizer
    
    if model is None or tokenizer is None:
        logger.warning("Model or tokenizer not loaded. Falling back to pattern-based responses.")
        return generate_response_from_patterns(prompt)
    
    try:
        # Format the prompt with system message
        system_prompt = "Вы Бабуру, саркастичный, язвительный плюшевый клоун, который любит издеваться над людьми."
        formatted_prompt = f"{system_prompt}\n\nЧеловек: {prompt}\n\nБабуру:"
        
        # Tokenize the input
        inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
        
        # Generate response
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_length,
                temperature=temperature,
                top_p=0.9,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )
        
        # Decode the response and extract the model's reply
        full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extract just the assistant's response (after "Бабуру:")
        assistant_response = full_response.split("Бабуру:")[-1].strip()
        
        # Clean up any remaining system prompt or instruction text
        if "Человек:" in assistant_response:
            assistant_response = assistant_response.split("Человек:")[0].strip()
        
        # Free CUDA memory
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()
        
        return assistant_response
    except Exception as e:
        logger.error(f"Error generating response with Gemma model: {str(e)}")
        # Fallback to pattern-based response
        return generate_response_from_patterns(prompt)

def generate_response(prompt, temperature=0.8):
    """Generate a response using the preferred model or fallback"""
    # Try to use the Gemma model first, fallback to pattern matching
    try:
        return gemma_inference(prompt, temperature)
    except Exception as e:
        logger.error(f"Failed to use Gemma model, falling back to patterns: {str(e)}")
        return generate_response_from_patterns(prompt)

@app.route('/')
def index():
    """Render the main page"""
    return render_template('voice_chat.html')

@app.route('/api/chat', methods=['POST'])
def chat():
    """API endpoint for text chat interactions"""
    try:
        # Get the user input from the request
        data = request.get_json()
        user_input = data.get('message', '')
        
        logger.info(f"Received chat input: {user_input}")
        
        if not user_input:
            return jsonify({
                'error': 'No message provided'
            }), 400
        
        # Check if input is in Russian or needs translation
        is_russian = any(ord(char) >= 128 for char in user_input)
        
        if not is_russian:
            russian_input = translate_text(user_input, src='en', dest='ru')
            input_for_model = russian_input
        else:
            input_for_model = user_input
        
        # Generate response
        response = generate_response(input_for_model)
        
        # Prepare translated responses
        if is_russian:
            translated_response = translate_text(response, src='ru', dest='en')
        else:
            translated_response = translate_text(response, src='ru', dest='en')
        
        # Convert response to speech
        audio_url = text_to_speech(response)
        
        return jsonify({
            'original_input': user_input,
            'model_input': input_for_model,
            'response': response,
            'translated_response': translated_response,
            'audio_url': audio_url
        })
        
    except Exception as e:
        logger.error(f"Error in chat endpoint: {str(e)}")
        logger.exception("Full exception details:")
        return jsonify({'error': str(e)}), 500

@app.route('/api/voice', methods=['POST'])
def voice():
    """API endpoint for voice interactions"""
    try:
        # Get the audio data from the request
        data = request.get_json()
        
        # Handle direct text from client-side speech recognition
        if 'text' in data and data['text']:
            user_input = data['text']
            logger.info(f"Received voice-to-text input: {user_input}")
            
            # Check if input is in Russian or needs translation
            is_russian = any(ord(char) >= 128 for char in user_input)
            
            if not is_russian:
                russian_input = translate_text(user_input, src='en', dest='ru')
                input_for_model = russian_input
            else:
                input_for_model = user_input
            
            # Generate response
            response = generate_response(input_for_model)
            
            # Prepare translated response
            if is_russian:
                translated_response = translate_text(response, src='ru', dest='en')
            else:
                translated_response = translate_text(response, src='ru', dest='en')
            
            # Convert response to speech
            audio_url = text_to_speech(response)
            
            return jsonify({
                'original_input': user_input,
                'model_input': input_for_model,
                'response': response,
                'translated_response': translated_response,
                'audio_url': audio_url
            })
        
    except Exception as e:
        logger.error(f"Error in voice endpoint: {str(e)}")
        return jsonify({'error': str(e)}), 500

def generate_response_from_patterns(prompt):
    """Generate a response using pattern matching approach"""
    # Check for specific patterns first
    for pattern, responses in PATTERN_RESPONSES.items():
        if pattern == "default":
            continue
            
        if re.search(pattern, prompt):
            return random.choice(responses)
    
    # If no pattern matches, use default responses
    return random.choice(PATTERN_RESPONSES["default"])

@app.route('/api/call', methods=['POST'])
def call_baburu():
    """Simple endpoint to get a random response from Baburu without user input"""
    try:
        # Just pick a random greeting or default response
        greetings_pattern = r'(?i)(привет|здравствуй|здорово|хай)'
        greeting_responses = PATTERN_RESPONSES.get(greetings_pattern, PATTERN_RESPONSES["default"])
        response = random.choice(greeting_responses)
        
        # Get English translation
        translated_response = translate_text(response, src='ru', dest='en')
        
        # Convert to speech
        audio_url = text_to_speech(response)
        
        return jsonify({
            'response': response,
            'translated_response': translated_response,
            'audio_url': audio_url
        })
    except Exception as e:
        logger.error(f"Error in call endpoint: {str(e)}")
        return jsonify({'error': str(e)}), 500

def initialize_server():
    """Initialize the server"""
    logger.info("Initializing server components...")
    
    # Test the translator to make sure it's working
    try:
        test_translation = translate_text("Тест", src='ru', dest='en')
        logger.info(f"Translation test: 'Тест' → '{test_translation}'")
    except Exception as e:
        logger.error(f"Translation test failed: {str(e)}")
    
    # Try to load the Gemma model
    try:
        logger.info("Attempting to load Gemma model with adapter...")
        model_loaded = load_gemma_model()
        
        if model_loaded:
            logger.info("Successfully loaded Gemma model! Voice functionality will use neural model responses with translation.")
        else:
            logger.warning("Failed to load Gemma model. Falling back to pattern-matching responses.")
            logger.info("Voice functionality will work with pattern-based responses")
    except Exception as e:
        logger.error(f"Error during model initialization: {str(e)}")
        logger.warning("Falling back to pattern-matching responses due to initialization error.")
    
    logger.info("Server initialization complete!")

if __name__ == "__main__":
    # Initialize server at startup
    logger.info("Initializing server...")
    initialize_server()
    
    # Run the Flask app
    port = int(os.environ.get("PORT", 5001))  # Changed to port 5001 to avoid conflicts with macOS AirPlay
    logger.info(f"Starting server on port {port}...")
    # Using 127.0.0.1 instead of 0.0.0.0 for better local access
    # Explicitly set ssl_context to None to force HTTP (not HTTPS)
    app.run(
        host="127.0.0.1", 
        port=port, 
        debug=True, 
        threaded=True,
        ssl_context=None
    )
    
    # Print access instructions
    logger.info(f"\nAccess the voice interface at: http://127.0.0.1:{port}/")