import os
import torch
from openvoice import se_extractor
from openvoice.api import ToneColorConverter
import requests
import zipfile
import shutil
device="cuda:0" if torch.cuda.is_available() else "cpu"
output_dir = 'outputs'

# os.system("wget https://myshell-public-repo-host.s3.amazonaws.com/openvoice/checkpoints_v2_0417.zip")
# os.system("unzip checkpoints_v2_0417.zip")

ROOT = os.getcwd()
cache_dir = f"{ROOT}/.cache/torch/hub"
extract_dir = os.path.join(cache_dir, "snakers4_silero-vad_master")
url = "https://github.com/snakers4/silero-vad/zipball/master"

if not os.path.exists(extract_dir):
    os.makedirs(cache_dir, exist_ok=True)
    zip_path = os.path.join(cache_dir, "master.zip")
    with requests.get(url, stream=True) as r:
        r.raise_for_status()
        with open(zip_path, 'wb') as f:
            for chunk in r.iter_content(chunk_size=8192):
                f.write(chunk)
    temp_extract_dir = os.path.join(cache_dir, "temp_extract")
    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        zip_ref.extractall(temp_extract_dir)

    # Find the extracted folder and rename it
    extracted_folder = os.path.join(temp_extract_dir, os.listdir(temp_extract_dir)[0])
    shutil.move(extracted_folder, extract_dir)
    shutil.rmtree(temp_extract_dir)

class Postprocessor:
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.tone_color_converter = ToneColorConverter('checkpoints_v2/converter/config.json', device=device)
        self.tone_color_converter.load_ckpt('checkpoints_v2/converter/checkpoint.pth')
        os.makedirs(output_dir, exist_ok=True)

    def convert_tone_color(self, reference_speaker, src_path="outputs/xtts.wav", save_path="outputs/openvoice.wav", base_speaker="outputs/xtts.wav"):
        source_se, audio_name = se_extractor.get_se(base_speaker, self.tone_color_converter, vad=True)
        target_se, audio_name = se_extractor.get_se(reference_speaker, self.tone_color_converter, vad=True)

        encode_message = "@MyShell"

        self.tone_color_converter.convert(
            audio_src_path=src_path,
            src_se=source_se,
            tgt_se=target_se,
            output_path=save_path,
            message=encode_message)

if __name__ == "__main__":
    postprocessor = Postprocessor()
    postprocessor.convert_tone_color(reference_speaker="outputs/xtts.wav")
    print("Done!")