import gradio as gr
import torch
import torchaudio
from torchaudio.transforms import Resample
import numpy as np
import os
from datetime import datetime
import soundfile as sf

# بررسی وجود GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def load_audio(file_path, target_sr=16000):
    """بارگذاری فایل صوتی و تبدیل به تنسور"""
    try:
        waveform, sample_rate = torchaudio.load(file_path)
        
        # تغییر نرخ نمونه‌برداری در صورت نیاز
        if sample_rate != target_sr:
            resampler = Resample(sample_rate, target_sr)
            waveform = resampler(waveform)
            
        return waveform.to(device), target_sr
    except Exception as e:
        raise gr.Error(f"خطا در بارگذاری فایل صوتی: {str(e)}")

def preprocess_audio(waveform, sr):
    """پیش‌پردازش سیگنال صوتی"""
    # نرمالایز کردن سیگنال
    waveform = waveform / torch.max(torch.abs(waveform))
    
    # تبدیل به مونو اگر استریو باشد
    if waveform.shape[0] > 1:
        waveform = torch.mean(waveform, dim=0, keepdim=True)
        
    return waveform

def clone_voice(source_audio, target_audio, output_format="wav"):
    """عمل کلون کردن صدا"""
    try:
        # بارگذاری فایل‌های صوتی
        source_waveform, source_sr = load_audio(source_audio)
        target_waveform, target_sr = load_audio(target_audio)
        
        # پیش‌پردازش
        source_waveform = preprocess_audio(source_waveform, source_sr)
        target_waveform = preprocess_audio(target_waveform, target_sr)
        
        # در اینجا باید مدل تبدیل صدا اعمال شود
        # این قسمت ساده‌سازی شده و نیاز به پیاده‌سازی واقعی دارد
        # برای نمونه، فقط طول موج هدف را با طول موج منبع هماهنگ می‌کنیم
        min_len = min(source_waveform.shape[1], target_waveform.shape[1])
        converted_waveform = target_waveform[:, :min_len]
        
        # ذخیره فایل نتیجه
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_filename = f"output_{timestamp}.{output_format.lower()}"
        
        # ذخیره با soundfile که با فرمت‌های مختلف کار می‌کند
        sf.write(output_filename, converted_waveform.cpu().numpy().T, target_sr)
        
        return output_filename
    except Exception as e:
        raise gr.Error(f"خطا در پردازش صدا: {str(e)}")

# رابط Gradio
with gr.Blocks(title="Voice Cloner") as demo:
    gr.Markdown("# 🎤 Voice Cloner")
    gr.Markdown("بارگذاری فایل صوتی منبع و فایل صوتی هدف برای کلون کردن صدا")
    
    with gr.Row():
        with gr.Column():
            source_audio = gr.Audio(label="فایل صوتی منبع (صدا برای کپی کردن)", type="filepath")
            target_audio = gr.Audio(label="فایل صوتی هدف (محتوا برای تبدیل)", type="filepath")
            output_format = gr.Dropdown(
                choices=["wav", "mp3"],
                value="wav",
                label="فرمت فایل خروجی",
                interactive=True
            )
            submit_btn = gr.Button("شروع تبدیل صدا")
        
        with gr.Column():
            output_audio = gr.Audio(label="فایل صوتی نتیجه", interactive=False)
            download_link = gr.File(label="دانلود فایل نتیجه")
    
    submit_btn.click(
        fn=clone_voice,
        inputs=[source_audio, target_audio, output_format],
        outputs=[output_audio]
    )
    
    output_audio.change(
        lambda x: gr.File(value=x),
        inputs=[output_audio],
        outputs=[download_link]
    )

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