import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from diffusers import StableDiffusionPipeline
import numpy as np
import soundfile as sf
import os

# -----------------------------
# تنظیمات پایه
# -----------------------------
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# -----------------------------
# 1) چت‌بات (مدل سبک‌تر Qwen)
# -----------------------------
CHAT_MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"

chat_tokenizer = AutoTokenizer.from_pretrained(CHAT_MODEL_ID)
chat_model = AutoModelForCausalLM.from_pretrained(
    CHAT_MODEL_ID,
    torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
    device_map="auto"
)

def chat_fn(user_text, history):
    if history is None:
        history = []
    history.append({"role": "user", "content": user_text})

    prompt = ""
    for turn in history:
        if turn["role"] == "user":
            prompt += f"User: {turn['content']}\n"
        else:
            prompt += f"Assistant: {turn['content']}\n"
    prompt += "Assistant:"

    inputs = chat_tokenizer(prompt, return_tensors="pt").to(DEVICE)
    outputs = chat_model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=True,
        temperature=0.8,
        top_p=0.9
    )
    text = chat_tokenizer.decode(outputs[0], skip_special_tokens=True)
    if "Assistant:" in text:
        text = text.split("Assistant:")[-1].strip()

    history.append({"role": "assistant", "content": text})
    return text, history

# -----------------------------
# 2) تصویر (مدل سبک‌تر SD Turbo)
# -----------------------------
IMG_MODEL_ID = "stabilityai/sd-turbo"

sd_pipe = StableDiffusionPipeline.from_pretrained(
    IMG_MODEL_ID,
    torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32
)
sd_pipe = sd_pipe.to(DEVICE)

def image_fn(prompt, steps, guidance, width, height):
    image = sd_pipe(
        prompt,
        num_inference_steps=steps,
        guidance_scale=guidance,
        width=width,
        height=height
    ).images[0]
    return image

# -----------------------------
# 3) موسیقی (MusicGen نسخه مناسب‌تر)
# -----------------------------
MUSIC_MODEL_ID = "facebook/musicgen-medium"

try:
    music_pipe = pipeline("text-to-audio", model=MUSIC_MODEL_ID, device=0 if DEVICE == "cuda" else -1)
except Exception:
    music_pipe = None

def music_fn(prompt, duration):
    if music_pipe is None:
        return None, "مدل MusicGen لود نشد. بعداً دوباره تلاش کن یا مدل را عوض کن."

    result = music_pipe(
        prompt,
        forward_params={"do_sample": True, "max_new_tokens": 512}
    )
    audio = result["audio"]
    sr = result["sampling_rate"]

    out_path = "music_output.wav"
    sf.write(out_path, audio, sr)
    return out_path, "آهنگ ساخته شد."

# -----------------------------
# 4) TTS (ساخت صدا از متن)
# -----------------------------
TTS_MODEL_ID = "facebook/mms-tts-eng"

try:
    tts_pipe = pipeline("text-to-speech", model=TTS_MODEL_ID, device=0 if DEVICE == "cuda" else -1)
except Exception:
    tts_pipe = None

def tts_fn(text):
    if tts_pipe is None:
        return None, "مدل TTS لود نشد. بعداً دوباره تلاش کن یا مدل را عوض کن."

    result = tts_pipe(text)
    audio = result["audio"]
    sr = result["sampling_rate"]

    out_path = "tts_output.wav"
    sf.write(out_path, audio, sr)
    return out_path, "صدا ساخته شد."

# -----------------------------
# 5) ویدیو (دمو ساده)
# -----------------------------
def video_fn(prompt, duration):
    msg = (
        "ساخت ویدیو مدل‌محور (مثل ModelScope Video / AnimateDiff) "
        "نیاز به GPU قوی دارد. این تب فعلاً دمو است. "
        "می‌توانیم بعداً روی سیستم قوی یا Space جداگانه فعالش کنیم."
    )
    return msg

# -----------------------------
# رابط کاربری Gradio
# -----------------------------
with gr.Blocks(title="Nexuz Multi-AI Studio") as demo:
    gr.Markdown("# 🌌 Nexuz Multi-AI Studio\nیک محیط شخصی، رایگان و چندمدلی برای تو ✨")

    with gr.Tab("💬 چت‌بات"):
        chat_state = gr.State([])
        chat_input = gr.Textbox(label="پیام تو", placeholder="هرچی دوست داری بنویس...")
        chat_output = gr.Textbox(label="جواب مدل", lines=8)
        chat_btn = gr.Button("ارسال")

        def chat_wrapper(user_text, state):
            reply, new_state = chat_fn(user_text, state)
            return reply, new_state

        chat_btn.click(
            chat_wrapper,
            inputs=[chat_input, chat_state],
            outputs=[chat_output, chat_state]
        )

    with gr.Tab("🖼 تصویر (Stable Diffusion)"):
        img_prompt = gr.Textbox(label="پرامپت تصویر", placeholder="مثلاً: anime girl in neon city...")
        img_steps = gr.Slider(5, 30, value=10, step=1, label="تعداد گام‌ها")
        img_guidance = gr.Slider(0.5, 10.0, value=3.0, step=0.5, label="Guidance Scale")
        img_width = gr.Slider(256, 768, value=512, step=64, label="عرض")
        img_height = gr.Slider(256, 768, value=512, step=64, label="ارتفاع")
        img_btn = gr.Button("ساخت تصویر")
        img_out = gr.Image(label="خروجی تصویر")

        img_btn.click(
            image_fn,
            inputs=[img_prompt, img_steps, img_guidance, img_width, img_height],
            outputs=img_out
        )

    with gr.Tab("🎵 موسیقی (MusicGen)"):
        music_prompt = gr.Textbox(label="پرامپت موسیقی", placeholder="soft romantic lo-fi beat with piano...")
        music_duration = gr.Slider(5, 60, value=15, step=5, label="مدت تقریبی (ثانیه، دمو)")
        music_btn = gr.Button("ساخت آهنگ")
        music_audio = gr.Audio(label="خروجی موسیقی", type="filepath")
        music_msg = gr.Textbox(label="وضعیت", interactive=False)

        music_btn.click(
            music_fn,
            inputs=[music_prompt, music_duration],
            outputs=[music_audio, music_msg]
        )

    with gr.Tab("🎤 TTS (ساخت صدا از متن)"):
        tts_text = gr.Textbox(label="متن برای خواندن", placeholder="هر جمله‌ای که دوست داری...")
        tts_btn = gr.Button("ساخت صدا")
        tts_audio = gr.Audio(label="خروجی صدا", type="filepath")
        tts_msg = gr.Textbox(label="وضعیت", interactive=False)

        tts_btn.click(
            tts_fn,
            inputs=[tts_text],
            outputs=[tts_audio, tts_msg]
        )

    with gr.Tab("🎥 ویدیو (دمو)"):
        video_prompt = gr.Textbox(label="پرامپت ویدیو", placeholder="مثلاً: anime girl walking in neon city...")
        video_duration = gr.Slider(2, 20, value=5, step=1, label="مدت (ثانیه، تئوری)")
        video_btn = gr.Button("توضیح وضعیت ویدیو")
        video_out = gr.Textbox(label="پیام", interactive=False)

        video_btn.click(
            video_fn,
            inputs=[video_prompt, video_duration],
            outputs=video_out
        )

    gr.Markdown("### 🔧 نکته\nاین محیط پایه است؛ هر مدل را می‌توانیم عوض کنیم، اضافه کنیم، یا شخصی‌سازی کنیم.")

if __name__ == "__main__":
    demo.launch()