import argparse
import os
import soundfile as sf

from tools.i18n.i18n import I18nAuto
from inference_webui import change_gpt_weights, change_sovits_weights, get_tts_wav
import hashlib

i18n = I18nAuto()
def get_tts_cache_key(model_name, text, prompt_audio_path):
    """
    生成 TTS 缓存 key: md5(模型+文本+md5(prompt音频内容))
    :param model_name: str
    :param text: str
    :param prompt_audio_path: str or None
    :return: str (md5 hash)
    """
    prompt_md5 = ''
    if prompt_audio_path and os.path.exists(prompt_audio_path):
        with open(prompt_audio_path, 'rb') as f:
            prompt_content = f.read()
            prompt_md5 = hashlib.md5(prompt_content).hexdigest()
    key_str = f"{model_name}::{text}::{prompt_md5}"
    return hashlib.md5(key_str.encode('utf-8')).hexdigest()

def synthesize(GPT_model_path, SoVITS_model_path, ref_audio_path, ref_text, ref_language, target_text, target_language, output_path):

    # Change model weights
    change_gpt_weights(gpt_path=GPT_model_path)
    change_sovits_weights(sovits_path=SoVITS_model_path)

    # Synthesize audio
    synthesis_result = get_tts_wav(
        ref_wav_path=ref_audio_path,
        prompt_text=ref_text,
        prompt_language=i18n(ref_language),
        text=target_text,
        text_language=i18n(target_language),
        top_p=1,
        temperature=1,
    )

    result_list = list(synthesis_result)

    if result_list:
        last_sampling_rate, last_audio_data = result_list[-1]
        output_wav_path = os.path.join(output_path, "output.wav")
        sf.write(output_wav_path, last_audio_data, last_sampling_rate)
        print(f"Audio saved to {output_wav_path}")


def main():
    parser = argparse.ArgumentParser(description="GPT-SoVITS Command Line Tool")
    parser.add_argument("--gpt_model", required=True, help="Path to the GPT model file")
    parser.add_argument("--sovits_model", required=True, help="Path to the SoVITS model file")
    parser.add_argument("--ref_audio", required=True, help="Path to the reference audio file")
    parser.add_argument("--ref_text", required=True, help="Path to the reference text file")
    parser.add_argument(
        "--ref_language", required=True, choices=["中文", "英文", "日文"], help="Language of the reference audio"
    )
    parser.add_argument("--target_text", required=True, help="Path to the target text file")
    parser.add_argument(
        "--target_language",
        required=True,
        choices=["中文", "英文", "日文", "中英混合", "日英混合", "多语种混合"],
        help="Language of the target text",
    )
    parser.add_argument("--output_path", required=True, help="Path to the output directory")

    args = parser.parse_args()

    synthesize(
        args.gpt_model,
        args.sovits_model,
        args.ref_audio,
        args.ref_text,
        args.ref_language,
        args.target_text,
        args.target_language,
        args.output_path,
    )


if __name__ == "__main__":
    main()