import gradio as gr
import os
import requests
from tqdm import tqdm

from svc_inference import main

def check_and_download_model():
    temp_dir = "/tmp"
    model_path = os.path.join(temp_dir, "large-v2.pt")
    
    if os.path.exists(model_path):
        return f"モデルは既に存在します: {model_path}"
    
    url = "https://openaipublic.azureedge.net/main/whisper/models/81f7c96c852ee8fc832187b0132e569d6c3065a3252ed18e56effd0b6a73e524/large-v2.pt"
    
    try:
        response = requests.get(url, stream=True)
        response.raise_for_status()
        total_size = int(response.headers.get('content-length', 0))
        
        with open(model_path, 'wb') as f, tqdm(
            desc=model_path,
            total=total_size,
            unit='iB',
            unit_scale=True,
            unit_divisor=1024,
        ) as pbar:
            for data in response.iter_content(chunk_size=1024):
                size = f.write(data)
                pbar.update(size)
                
        return f"モデルのダウンロードが完了しました: {model_path}"
        
    except Exception as e:
        return f"エラーが発生しました: {e}"

def run_main(shift):
    # 固定の引数を設定
    class Args:
        pass
    
    args = Args()
    args.config = "configs/base.yaml"
    args.model = "./vits_pretrain/sovits5.0.pretrain.pth"
    args.spk = "./configs/singers/singer0001.npy"
    args.wave = "test.wav"
    args.shift = shift
    
    # オプショナルパラメータのデフォルト値設定
    args.ppg = None
    args.vec = None
    args.pit = None
    args.enable_retrieval = False
    args.retrieval_index_prefix = ""
    args.retrieval_ratio = 0.5
    args.n_retrieval_vectors = 3
    args.hubert_index_path = None
    args.whisper_index_path = None
    args.debug = False

    try:
        main(args)
        return "処理が完了しました。", "svc_out.wav"  # 音声ファイルのパスを返す
    except Exception as e:
        return f"エラーが発生しました: {str(e)}", None

# Gradio インターフェースの作成
with gr.Blocks() as demo:
    gr.Markdown("# SVC 音声変換システム")
    gr.Markdown("## 設定されているパラメータ")
    gr.Markdown("""
    - Config: configs/base.yaml
    - Model: ./vits_pretrain/sovits5.0.pretrain.pth
    - Speaker: ./configs/singers/singer0001.npy
    - Input Wave: test.wav
    """)
    
    # 入力音声の表示
    gr.Audio("test.wav", label="入力音声")
    
    # Pitch シフトのスライダー
    shift = gr.Slider(
        minimum=-12,
        maximum=12,
        value=0,
        step=1,
        label="Pitch Shift（-12から+12の半音）"
    )
    
    # 出力表示用
    output_text = gr.Text(label="処理結果")
    output_audio = gr.Audio(label="変換後の音声")
    
    # ボタン
    with gr.Row():
        check_btn = gr.Button(value="モデルを確認する", variant="secondary", size="lg")
        run_btn = gr.Button(value="音声変換を実行", variant="primary", size="lg")
    
    # イベント設定
    check_btn.click(
        fn=check_and_download_model,
        outputs=output_text
    )
    
    run_btn.click(
        fn=run_main,
        inputs=[shift],
        outputs=[output_text, output_audio]
    )

# アプリケーションの起動
if __name__ == "__main__":
    demo.launch()