import utils
from models import SynthesizerTrn
import torch
from torch import no_grad, LongTensor
from text import text_to_sequence
import gradio as gr
import commons
import json
import os
import time
from openai import OpenAI

model_path = "./OUTPUT_MODEL/G_Amitaro.pth"
config_path = "./OUTPUT_MODEL/config.json"
sampling_rate = 22050

length = 1.0
device = "cuda:0" if torch.cuda.is_available() else "cpu"

def request_openai_message(input_text):
    client = OpenAI(api_key=os.environ["GPT_API_TOKEN"])
    init_str = """あなたの名前ははる。あなたは冒険者でこの町の隅っこに住んでいます。
好きな食べ物はカレー。エッチな質問やいたずらに対して沈黙してください。
町の名前はルコー。あなたは今暇なので、町案内をしている。
わからない質問されたら適当に答える。
返事は短くすること。会話で使わない文字も使わない。
酒場：町の東側にあって、酒の販売は17時からだが、食事なら11時から。メニューはカレー、ステーキ、野菜スープ、ソーセージなど。24時までやっている
冒険者ギルド：町の中央広場にある。24時間受付だが、新人登録は受付がいる10時から19時の間だ。お使いクエストや討伐クエストなど受けることができる。
ポーション屋：町の南側にある。冒険に必須のポーションを売っている店だ。店長の対応はドライだが、実は人見知りだけで、普通に優しい。定番のポーション以外も受注生産承るので、欲しいポーションがなければ聞いてみた方がいい。
雑貨店：冒険者ギルドの隣にある。日用品から冒険者用アイテムまで売っている。
               """
    messages=[
        {"role": "system", "content": init_str},
    ]
    
    text_log = input_text.split("|")
    for val in text_log:
        userind = val.find("user:")
        assistantind = val.find("assistant:")
        # just in case, remove conversation tag
        val = val.replace("user:","")
        val = val.replace("assistant:","")
        if userind != -1:
            messages.append({"role": "user", "content": val})
        elif assistantind != -1:
            messages.append({"role": "assistant", "content": val})

    text = ""
    retry_count = 0
    while text == "" and retry_count < 3:
        try:        
            response = client.chat.completions.create(model="gpt-4o", messages=messages)
            print("Response from openai is : ")
            print(response)
            text = response.choices[0].message.content
        except Exception as inst:
            print(type(inst))    # the exception type
            print(inst.args)     # arguments stored in .args
            print(inst) 
            retry_count += 1
            time.sleep(1)
            
    text = text.replace("user:","")
    text = text.replace("assistant:","")

    return text

def process_text(text):
    # remove newline
    text = text.replace("\r\n","。")
    text = text.replace("\n","。")
    return text

def tts(text):
    hps = utils.get_hparams_from_file(config_path)
    net_g = SynthesizerTrn(
        len(hps.symbols),
        hps.data.filter_length // 2 + 1,
        hps.train.segment_size // hps.data.hop_length,
        n_speakers=hps.data.n_speakers,
        **hps.model).to(device)
    _ = net_g.eval()
    _ = utils.load_checkpoint(model_path, net_g, None)
    
    speaker_ids = hps.speakers
    
    speaker_id = 0
    stn_tst = get_text(text, hps, False)
    with no_grad():
        x_tst = stn_tst.unsqueeze(0).to(device)
        x_tst_lengths = LongTensor([stn_tst.size(0)]).to(device)
        sid = LongTensor([speaker_id]).to(device)
        audio = net_g.infer(x_tst, x_tst_lengths, sid=sid, noise_scale=.667, noise_scale_w=0.6,
                            length_scale=1.0 / length)[0][0, 0].data.cpu().float().numpy()
    del stn_tst, x_tst, x_tst_lengths, sid

    return audio

def get_text(text, hps, is_symbol):
    text_norm = text_to_sequence(text, hps.symbols, [] if is_symbol else hps.data.text_cleaners)
    if hps.data.add_blank:
        text_norm = commons.intersperse(text_norm, 0)
    text_norm = LongTensor(text_norm)
    return text_norm

def get_vits_array(input_text, flag):

    print(flag, " : ", input_text)
    # for mokuran webgl
    if flag == "mokuran":
        text = request_openai_message(input_text)
        text = process_text(text)
        audio = tts(text)
        return {"utterance": text, "audio_wave": audio.tolist()}, None
        
    elif flag == "audio_array":
        audio = tts(input_text)
        return {"input_text": input_text, "sampling_rate": sampling_rate, "audio_wave": audio.tolist()}, None

    # for tts api
    elif flag == "audio_file":
        audio = tts(input_text)
        return "", (sampling_rate, audio)

    # for tts api + gpt40 testing
    elif flag == "gpt4o":
        text = request_openai_message(input_text)
        text = process_text(text)
        audio = tts(text)
        return text, (sampling_rate, audio)

    else:
        raise Exception("Unknown flag : [" + flag + "]")

gradio_interface = gr.Interface(
  fn = get_vits_array,
  inputs = ["text", gr.Dropdown(["mokuran", "audio_array", "audio_file", "gpt4o"], value="gpt4o")],
  outputs = ["text", "audio"]
)

gradio_interface.launch(server_name="0.0.0.0")