from spleeter import Splitter
import torchaudio
from torchaudio.transforms import Resample
import torch
import gradio as gr
from io import BytesIO
import requests

def yt2audio(search:str,format:str='mp3',audioBitrate='320',videoQuality='720'):
    return requests.get(requests.post('https://cnv.cx/v2/converter', headers={
    'Accept': '*/*',
    'Accept-Language': 'en-GB,en;q=0.9',
    'Connection': 'keep-alive',
    'Content-Type': 'application/x-www-form-urlencoded',
    'DNT': '1',
    'Origin': 'https://frame.y2meta-uk.com',
    'Referer': 'https://frame.y2meta-uk.com/',
    'Sec-Fetch-Dest': 'empty',
    'Sec-Fetch-Mode': 'cors',
    'Sec-Fetch-Site': 'cross-site',
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.0.0 Safari/537.36',
    'key': requests.get('https://cnv.cx/v2/sanity/key', headers={
    'Accept': '*/*',
    'Accept-Language': 'en-GB,en;q=0.9',
    'Connection': 'keep-alive',
    'Content-Type': 'application/json',
    'DNT': '1',
    'If-None-Match': requests.get('https://cdn.jsdelivr.net/npm/@iframe-resizer/child', headers={
    'sec-ch-ua-platform': '"Windows"',
    'Referer': 'https://frame.y2meta-uk.com/',
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.0.0 Safari/537.36',
    'sec-ch-ua': '"Google Chrome";v="141", "Not?A_Brand";v="8", "Chromium";v="141"',
    'DNT': '1',
    'sec-ch-ua-mobile': '?0',
}).headers['etag'],
    'Origin': 'https://frame.y2meta-uk.com',
    'Referer': 'https://frame.y2meta-uk.com/',
    'Sec-Fetch-Dest': 'empty',
    'Sec-Fetch-Mode': 'cors',
    'Sec-Fetch-Site': 'cross-site',
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.0.0 Safari/537.36',
    'sec-ch-ua': '"Google Chrome";v="141", "Not?A_Brand";v="8", "Chromium";v="141"',
    'sec-ch-ua-mobile': '?0',
    'sec-ch-ua-platform': '"Windows"',
},verify=False).json()['key'],
    'sec-ch-ua': '"Google Chrome";v="141", "Not?A_Brand";v="8", "Chromium";v="141"',
    'sec-ch-ua-mobile': '?0',
    'sec-ch-ua-platform': '"Windows"',
}, data={
    'link': 'https://youtu.be/'+requests.get(f'https://wwd.mp3juice.blog/search.php?q={search}').json()['items'][0]['id'],
    'format': format,
    'audioBitrate': audioBitrate,
    'videoQuality': videoQuality,
    'filenameStyle': 'pretty',
    'vCodec': 'h264',
},verify=False).json()['url'], headers={
    'DNT': '1',
    'Referer': 'https://frame.y2meta-uk.com/',
    'Upgrade-Insecure-Requests': '1',
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.0.0 Safari/537.36',
    'sec-ch-ua': '"Google Chrome";v="141", "Not?A_Brand";v="8", "Chromium";v="141"',
    'sec-ch-ua-mobile': '?0',
    'sec-ch-ua-platform': '"Windows"',
},verify=False).content


def separate_audio(audio_path:str,inst_no:int,progress=gr.Progress(True)):
    """
    Separate audio into instrument tracks.

    Args:
        audio_path (str): Path to input audio.
        inst_no (int): Number of instruments to separate.

    Returns:
        tuple: Up to 5 MP3 file paths for separated tracks.
    """
    
    model = Splitter(inst_no)
    wav, sr = torchaudio.load(audio_path)
    target_sr = 44100
    if sr != target_sr:
        resampler = Resample(sr, target_sr)
        wav = resampler(wav)
        sr = target_sr
    with torch.inference_mode():
        results = model.forward(wav)
    for i in results:
        torchaudio.save(f"{i}.mp3", results[i], sr)
    return tuple([i+".mp3" for i in results] + [None for _ in range(5-len(results))])

def separate_audio_by_search(query:str,inst_no:int,progress=gr.Progress(True)):
    """
    Searches YouTube for the given text query, downloads the audio,
    and separates it into multiple instruments (e.g., vocals, drums, bass, etc.).

    Args:
        query (str): Text query to search on YouTube.
        inst_no (int): Number of instruments to separate.

    Returns:
        tuple: Up to 5 MP3 file paths for separated tracks.
    """
    
    model = Splitter(inst_no)
    wav, sr = torchaudio.load(BytesIO(yt2audio(query)),format="mp3")
    target_sr = 44100
    if sr != target_sr:
        resampler = Resample(sr, target_sr)
        wav = resampler(wav)
        sr = target_sr
    with torch.inference_mode():
        results = model.forward(wav)
    for i in results:
        torchaudio.save(f"{i}.mp3", results[i], sr)
    return tuple([i+".mp3" for i in results] + [None for _ in range(5-len(results))])

gr.TabbedInterface([gr.Interface(separate_audio_by_search, [gr.Textbox(),gr.Dropdown([2,4,5])], [gr.Audio(type="filepath"), gr.Audio(type="filepath"),gr.Audio(type="filepath"),gr.Audio(type="filepath"),gr.Audio(type="filepath")]),gr.Interface(separate_audio, [gr.Audio(type="filepath"),gr.Dropdown([2,4,5])], [gr.Audio(type="filepath"), gr.Audio(type="filepath"),gr.Audio(type="filepath"),gr.Audio(type="filepath"),gr.Audio(type="filepath")])],["YT2STEM","Audio2STEM"]).launch(mcp_server=True)