import os, subprocess
import gradio as gr
import shutil
from mega import Mega
from mediafiredl import MediafireDL as MF
import os
import edge_tts
import yt_dlp
from pydub import AudioSegment
        
def download_from_url(url=None, model=None):
    if not url:
        try:
            url = model[f'{model}']
        except:
            gr.Warning("Thất Bại")
            return ''
    if model == '':
        try:
            model = url.split('/')[-1].split('?')[0]
        except:
            return gr.Warning('Vui Lòng Đưa Url')
            
    model = model.replace('.pth', '').replace('.index', '').replace('.zip', '')
    url = url.replace('/blob/main/', '/resolve/main/').replace('?download=true', '').strip()

    for directory in ["downloads", "unzips", "zip"]:
        os.makedirs(directory, exist_ok=True)

    try:
        if url.endswith('.pth'):
            subprocess.run(["wget", url, "-O", f'assets/weights/{model}.pth'])
        elif url.endswith('.index'):
            os.makedirs(f'logs/{model}', exist_ok=True)
            subprocess.run(["wget", url, "-O", f'logs/{model}/added_{model}.index'])
        elif url.endswith('.zip'):
            subprocess.run(["wget", url, "-O", f'downloads/{model}.zip'])
        else:
            if 'drive.google.com' in url:
                if '/file/d/' in url:
                    url = url.split('/d/')[1].split('/')[0]
                    subprocess.run(["gdown", url, "-O", f'downloads/{model}'])
                elif 'open?id=' in url:
                    url = url.split('open?id=')[1].split('/')[0]
                    subprocess.run(["gdown", url, "-O", f'downloads/{model}'])
                else:
                    return gr.Warning("Link Không Được Hỗ Trợ")
            elif 'drive.usercontent.google.com' in url:
                url = url.split('/download?id=')[1].split('&')[0]
                subprocess.run(["gdown", url, "-O", f'downloads/{model}'])
            elif "mega.nz" in url:
                Mega().download_url(url, 'downloads')
            elif "mediafire.com" in url:
                MF.Download(url, 'downloads')
            else:
                subprocess.run(["wget", url, "-O", f'downloads/{model}'])

        downloaded_file = next((f for f in os.listdir("downloads")), None)
        if downloaded_file:
            if downloaded_file.endswith(".zip"):
                shutil.unpack_archive(f'downloads/{downloaded_file}', "unzips", 'zip')
                for root, _, files in os.walk('unzips'):
                    for file in files:
                        file_path = os.path.join(root, file)
                        if file.endswith(".index"):
                            os.makedirs(f'logs/{model}', exist_ok=True)
                            shutil.copy2(file_path, f'logs/{model}')
                        elif file.endswith(".pth") and "G_" not in file and "D_" not in file:
                            shutil.copy(file_path, f'assets/weights/{model}.pth')
            elif downloaded_file.endswith(".pth"):
                shutil.copy(f'downloads/{downloaded_file}', f'assets/weights/{model}.pth')
            elif downloaded_file.endswith(".index"):
                os.makedirs(f'logs/{model}', exist_ok=True)
                shutil.copy(f'downloads/{downloaded_file}', f'logs/{model}/added_{model}.index')
            else:
                gr.Warning("Đã Xảy Ra Lỗi Khi Tải Model")
                return 'Thất Bại'

        gr.Info("Xong")
    except Exception as e:
        gr.Warning(f"Đã Xảy Ra Lỗi: {str(e)}")
    finally:
        shutil.rmtree("downloads", ignore_errors=True)
        shutil.rmtree("unzips", ignore_errors=True)
        shutil.rmtree("zip", ignore_errors=True)
        return 'Xong'
    
def tach_nhac(input, format, segments_size, overlap, denoise):
    if not os.path.exists(input): return gr.Warning('Vui lòng nhập Input')

    output = 'audios'
    if os.path.exists(f"audios/instrumental.{format}"): os.remove(f"audios/instrumental.{format}")
    if os.path.exists(f"audios/original_vocal.{format}"): os.remove(f"audios/original_vocal.{format}")
    if os.path.exists(f"audios/backup_vocals.{format}"): os.remove(f"audios/backup_vocals.{format}")
    if os.path.exists(f"audios/main_vocals.{format}"): os.remove(f"audios/main_vocals.{format}")

    prompt1 = f'audio-separator "{input}" --model_filename UVR-MDX-NET-Voc_FT.onnx --output_dir={output} --output_format={format} --normalization=0.9 --mdx_segment_size={segments_size} --mdx_overlap={overlap}'
    if denoise: 
        prompt1 += " --mdx_enable_denoise"
    os.system(prompt1)

    instruments = [f for f in os.listdir(output) if '(Instrumental)_UVR-MDX-NET-Voc_FT' in f]
    if instruments: os.rename(f"audios/{instruments[0]}", f'audios/instrumental.{format}')
    original_vocals = [f for f in os.listdir(output) if '(Vocals)_UVR-MDX-NET-Voc_FT' in f]
    if original_vocals: os.rename(f"audios/{original_vocals[0]}", f'audios/original_vocal.{format}')

    original_vocal = f"audios/original_vocal.{format}"
    prompt2 = f'audio-separator "{original_vocal}" --model_filename UVR_MDXNET_KARA_2.onnx --output_dir={output} --output_format={format} --normalization=0.9 --mdx_segment_size={segments_size} --mdx_overlap={overlap}'
    if denoise: 
        prompt2 += " --mdx_enable_denoise"
    os.system(prompt2)

    backup_vocals = [f for f in os.listdir(output) if '(Instrumental)_UVR_MDXNET_KARA_2' in f]
    if backup_vocals: os.rename(f"audios/{backup_vocals[0]}", f'audios/backup_vocals.{format}')

    main_vocals = [f for f in os.listdir(output) if '(Vocals)_UVR_MDXNET_KARA_2' in f]
    if main_vocals: os.rename(f"audios/{main_vocals[0]}", f'audios/main_vocals.{format}')

    gr.Info("Xong")
    return f'audios/instrumental.{format}', f'audios/original_vocal.{format}', f'audios/backup_vocals.{format}', f'audios/main_vocals.{format}'

async def TTS(text, giong, speed):
    if text == '': return gr.Warning('Vui lòng nhập prompt')
    if os.path.exists(f"audios/text.wav"): os.remove(f"audios/text.wav")
    output_file = f"audios/text.wav"
    rates = f"+{speed}%" if speed >= 0 else f"{speed}%"
    communicate = edge_tts.Communicate(text, giong, rate=rates)
    await communicate.save(output_file)
    gr.Info("Xong")
    return output_file, output_file, "Xong"

def download_url(url):
  if url == '': return gr.Warning("Vui lòng nhập Url")
  if os.path.exists("audios/audio.wav"): os.remove("audios/audio.wav")
  ydl_opts = {
    'format': 'bestaudio/best',
    'outtmpl': 'audios/audio',
    'postprocessors': [{
        'key': 'FFmpegExtractAudio',
        'preferredcodec': 'wav',
        'preferredquality': '192',
    }],
    'noplaylist': True,
    'verbose': True,
  }

  with yt_dlp.YoutubeDL(ydl_opts) as ydl:
    ydl.download([url])
  gr.Info("Xong")
  return "audios/audio.wav", "Xong"

def merge_audio(file1, file2, file3, file4, format):
    if os.path.exists(f"audios/combined_audio.{format}"): os.remove(f"audios/combined_audio.{format}")
    audio_segments = []
    if file1:
        audio_segments.append(AudioSegment.from_file(file1.name))
    if file2:
        audio_segments.append(AudioSegment.from_file(file2.name))
    if file3:
        audio_segments.append(AudioSegment.from_file(file3.name))
    if file4:
        audio_segments.append(AudioSegment.from_file(file4.name))
    
    if len(audio_segments) < 2:
        return gr.Warning("Cần Ít Nhất 2 Đoạn Âm Thanh Để Sử Dụng")

    combined = audio_segments[0]
    for segment in audio_segments[1:]:
        combined = combined.overlay(segment)
    
    combined.export(f"audios/combined_audio.{format}", format=format)
    gr.Info("Xong")
    return f"audios/combined_audio.{format}"