import torch
import logging
import soundfile as sf
import numpy as np
from io import BytesIO
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from diffusers import StableDiffusionPipeline
from telegram import Update
from telegram.ext import Application, MessageHandler, CommandHandler, filters, CallbackContext

# 📌 توکن ربات تلگرام را اینجا بگذار
TOKEN = "توکن_ربات_تلگرام"

# 📌 مدل‌های رایگان از Hugging Face
CHAT_MODEL = "mistralai/Mistral-7B-v0.1"  # چت
IMAGE_MODEL = "stabilityai/stable-diffusion-2"  # تصویرسازی
TTS_MODEL = "facebook/mms-tts-eng"  # متن به صدا
STT_MODEL = "openai/whisper-base"  # صدا به متن

# 📌 بارگذاری مدل‌های هوش مصنوعی
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(CHAT_MODEL)
chat_model = AutoModelForCausalLM.from_pretrained(CHAT_MODEL, device_map="auto")
image_pipeline = StableDiffusionPipeline.from_pretrained(IMAGE_MODEL).to(device)
tts_pipeline = pipeline("text-to-speech", model=TTS_MODEL)
stt_pipeline = pipeline("automatic-speech-recognition", model=STT_MODEL)

# 📌 تنظیمات لاگ
logging.basicConfig(format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO)

# 🔹 چت هوش مصنوعی
async def chat(update: Update, context: CallbackContext):
    text = update.message.text
    inputs = tokenizer(text, return_tensors="pt").to(device)
    outputs = chat_model.generate(**inputs, max_new_tokens=100)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    await update.message.reply_text(response)

# 🔹 تصویرسازی هوش مصنوعی
async def generate_image(update: Update, context: CallbackContext):
    text = update.message.text.replace("/image", "").strip()
    image = image_pipeline(text).images[0]
    
    bio = BytesIO()
    image.save(bio, format="PNG")
    bio.seek(0)
    
    await update.message.reply_photo(photo=bio)

# 🔹 تبدیل متن به صدا
async def text_to_speech(update: Update, context: CallbackContext):
    text = update.message.text.replace("/tts", "").strip()
    audio = tts_pipeline(text)
    
    bio = BytesIO()
    sf.write(bio, np.array(audio["audio"]), samplerate=audio["sampling_rate"], format="WAV")
    bio.seek(0)
    
    await update.message.reply_audio(audio=bio)

# 🔹 تبدیل صدا به متن
async def speech_to_text(update: Update, context: CallbackContext):
    file = await update.message.voice.get_file()
    bio = BytesIO()
    await file.download_to_memory(bio)
    bio.seek(0)

    text = stt_pipeline(bio)
    await update.message.reply_text(f"متن تبدیل شده: {text['text']}")

# 🔹 اجرای کد پایتون
async def run_code(update: Update, context: CallbackContext):
    code = update.message.text.replace("/code", "").strip()
    
    try:
        exec_globals = {}
        exec(code, exec_globals)
        result = exec_globals.get("result", "کدی اجرا شد اما خروجی‌ای ندارد.")
    except Exception as e:
        result = f"خطا: {e}"
    
    await update.message.reply_text(f"نتیجه:\n{result}")

# 📌 راه‌اندازی ربات تلگرام
def main():
    app = Application.builder().token(TOKEN).build()
    
    app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, chat))
    app.add_handler(CommandHandler("image", generate_image))
    app.add_handler(CommandHandler("tts", text_to_speech))
    app.add_handler(CommandHandler("stt", speech_to_text))
    app.add_handler(CommandHandler("code", run_code))

    app.run_polling()

if __name__ == "__main__":
    main()