import io
import os

import gradio as gr
import requests
from pydub import AudioSegment

from TTSs.base_tts import Base_TTS


class REECHO_TTS(Base_TTS):
    available_roles = {}

    def get_name(self):
        return 'REECHO 睿声'

    def get_reecho_ai_name(self, token):
        url = 'https://v1.reecho.cn/api/tts/voice'
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {token}'
        }
        r = requests.get(url, headers=headers).json()

        if r['status'] == 200:
            role_dict = {}
            for i in r['data']:
                info_data = [i['id'], i['metadata']['promptMP3StorageUrl'], b'']

                if i['metadata']['description'] == '':
                    role_dict['{}'.format(i['name'])] = info_data
                else:
                    role_dict['{}——{}'.format(i['name'], i['metadata']['description'])] = info_data
            return role_dict

        else:
            return {'获取角色列表失败': ['', '', b''], str(r): ['', '', b'']}

    def get_example_audio(self, role):
        try:
            if self.available_roles[role][2] != b'':
                return self.available_roles[role][2]

            audio_url = self.available_roles[role][1]
            audio_data = requests.get(audio_url).content
            self.available_roles[role][2] = audio_data
            return audio_data
        except Exception as e:
            return None

    def refresh_roles(self, reecho_api_key):
        self.available_roles = self.get_reecho_ai_name(reecho_api_key)

    def _get_config_page(self):
        with gr.Group(visible=False) as config_reecho_ai:
            gr.Markdown('控制台在此：https://dash.reecho.ai/')
            with gr.Row():
                reecho_api_key = gr.Textbox(label="reecho的API Key（默认为环境变量值）",
                                            placeholder="请输入reecho的API Key",
                                            type="password",
                                            interactive=True,
                                            value=os.environ.get('REECHO_API_KEY', ''),
                                            scale=3)
                btn_refush = gr.Button(value="刷新角色列表")

            with gr.Row():
                randomness = gr.Slider(minimum=0, maximum=100, value=97, step=1, label="多样性",
                                       interactive=True)
                stability_boost = gr.Slider(minimum=0, maximum=100, value=0, step=1, label="稳定性过度",
                                            interactive=True)

            self.refresh_roles(reecho_api_key.value)

            role = gr.Dropdown(choices=self.available_roles.keys(),
                               value=list(self.available_roles.keys())[0],
                               label="角色", interactive=True)

            example_audio = gr.Audio(label="该角色示例音频", interactive=False,
                                     value=self.get_example_audio(role.value))

            role.change(self.get_example_audio,
                        inputs=[role], outputs=[example_audio])

            btn_refush.click(lambda x: (self.refresh_roles(x),
                                        gr.Dropdown(choices=self.available_roles.keys(),
                                                    value=list(self.available_roles.keys())[0],
                                                    label="角色", interactive=True)
                                        )[-1],
                             inputs=[reecho_api_key],
                             outputs=[role])

        inputs = [
            reecho_api_key, role, randomness, stability_boost
        ]

        return config_reecho_ai, inputs

    def _generate(self, text, reecho_api_key, role, randomness, stability_boost):
        url = 'https://v1.reecho.cn/api/tts/simple-generate'

        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {reecho_api_key}'
        }

        body = {
            'model': 'reecho-neural-voice-001',
            'randomness': randomness,
            'stability_boost': stability_boost,
            'voiceId': self.available_roles[role][0],
            'text': text
        }

        r = requests.post(url, headers=headers, json=body)
        if r.status_code != 200:
            raise Exception(f'请求失败，HTTP状态码为{r.status_code}\n\n{r.__dict__}')
        r = r.json()
        mp3_data = requests.get(r['data']['audio']).content

        original_audio = AudioSegment.from_file(io.BytesIO(mp3_data), format="mp3")
        return original_audio