import os
import logging
import threading
import numpy as np
import torch
import soundfile as sf
from pydub import AudioSegment
from telegram import Update
from telegram.ext import ApplicationBuilder, MessageHandler, filters, CommandHandler
from transformers import AutoTokenizer, AutoModelForCausalLM
from TTS.api import TTS
from huggingface_hub import login
import asyncio
import noisereduce as nr
from scipy.io import wavfile

# ===== تهيئة التوكن =====
login(token=os.getenv("HF_TOKEN"))

# ===== إعدادات النظام =====
logging.basicConfig(
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    level=logging.INFO
)
logger = logging.getLogger(__name__)

# ===== تحميل النماذج =====
try:
    # 1. نموذج التعرف على الكلام (محدث)
    tts = TTS(model_name="tts_models/multilingual/multi-dataset/xtts_v2", progress_bar=False, gpu=False)
    
    # 2. نموذج الذكاء الاصطناعي (محدث)
    tokenizer = AutoTokenizer.from_pretrained("aubmindlab/aragpt2-mega")
    model = AutoModelForCausalLM.from_pretrained("aubmindlab/aragpt2-mega")

except Exception as e:
    logger.error(f"فشل تحميل النماذج: {str(e)}")
    raise

# ===== ذاكرة المحادثة =====
conversation_history = {}
MAX_HISTORY = 3

# ===== دوال معالجة الصوت =====
def enhance_audio(input_path: str, output_path: str) -> bool:
    try:
        # تقليل الضوضاء وتحسين الجودة
        rate, data = wavfile.read(input_path)
        reduced_noise = nr.reduce_noise(y=data, sr=rate, stationary=True)
        normalized = reduced_noise / np.max(np.abs(reduced_noise))
        wavfile.write(output_path, rate, normalized)
        return True
    except Exception as e:
        logger.error(f"خطأ في تحسين الصوت: {str(e)}")
        return False

async def speech_to_text(audio_path: str) -> str:
    try:
        return tts.transcribe(audio_path, language="ar")
    except Exception as e:
        logger.error(f"فشل التعرف على الصوت: {str(e)}")
        return ""

async def generate_response(text: str, user_id: str) -> str:
    try:
        # إدارة ذاكرة المحادثة
        history = conversation_history.get(user_id, [])
        history.append(text)
        history = history[-MAX_HISTORY:]
        conversation_history[user_id] = history
        
        # توليد الرد مع ضوابط الجودة
        input_text = "المحادثة:\n" + "\n".join([f"أنت: {msg}" for msg in history]) + "\nالبوت:"
        
        inputs = tokenizer.encode(input_text, return_tensors="pt", max_length=512, truncation=True)
        
        outputs = model.generate(
            inputs,
            max_length=200,
            num_return_sequences=1,
            temperature=0.7,
            repetition_penalty=1.5,
            pad_token_id=tokenizer.eos_token_id
        )
        
        response = tokenizer.decode(outputs[0], skip_special_tokens=True).split("البوت:")[-1].strip()
        return response.split("\n")[0]
    except Exception as e:
        logger.error(f"فشل توليد الرد: {str(e)}")
        return "عذرًا، حدث خطأ ما. يرجى المحاولة لاحقًا."

async def text_to_speech(text: str) -> None:
    try:
        tts.tts_to_file(
            text=text,
            file_path="bot_response.wav",
            speaker="ar_female_1",  # صوت أنثوي
            language="ar"
        )
    except Exception as e:
        logger.error(f"فشل تحويل النص إلى صوت: {str(e)}")

# ===== دوال التفاعل مع المستخدم =====
async def start(update: Update, context):
    await update.message.reply_text("مرحبًا! أنا البوت الصوتي الأنثوي 🎤\nأرسل لي رسالة صوتية وسأرد عليك بصوت أنثوي واضح.")

async def process_voice(update: Update, context):
    try:
        user_id = str(update.message.from_user.id)
        voice_file = await update.message.voice.get_file()
        await voice_file.download_to_drive("user_voice.ogg")
        
        # تحويل الصوت إلى نص
        user_text = await speech_to_text("user_voice.ogg")
        if not user_text:
            await update.message.reply_text("لم أتمكن من فهم الصوت. يرجى المحاولة مرة أخرى.")
            return
        
        # توليد الرد
        bot_response = await generate_response(user_text, user_id)
        
        # تحويل الرد إلى صوت
        await text_to_speech(bot_response)
        
        # تحسين الصوت
        if enhance_audio("bot_response.wav", "bot_response_enhanced.wav"):
            await update.message.reply_voice("bot_response_enhanced.wav")
        else:
            await update.message.reply_voice("bot_response.wav")
            
    except Exception as e:
        logger.error(f"خطأ غير متوقع: {str(e)}")
        await update.message.reply_text("⚠️ عذرًا، حدث خطأ في المعالجة.")

# ===== التشغيل الرئيسي =====
def run_bot():
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    
    application = ApplicationBuilder().token(os.getenv("TELEGRAM_TOKEN")).build()
    application.add_handler(CommandHandler("start", start))
    application.add_handler(MessageHandler(filters.VOICE, process_voice))
    
    application.run_polling()

if __name__ == "__main__":
    bot_thread = threading.Thread(target=run_bot, daemon=True)
    bot_thread.start()
    bot_thread.join()