import logging
from typing import AsyncGenerator, Iterator, Union
import os
import io

from pydub import AudioSegment

from achatbot.common.interface import ITts
from achatbot.common.session import Session
from achatbot.common.types import GTTS_SYNTHESIS_FILE, GTTSArgs, RECORDS_DIR
from .base import BaseTTS


class GTTS(BaseTTS, ITts):
    TAG = "tts_g"

    @classmethod
    def get_args(cls, **kwargs) -> dict:
        return {**GTTSArgs().__dict__, **kwargs}

    def __init__(self, **args) -> None:
        self.args = GTTSArgs(**args)
        self.file_path = os.path.join(RECORDS_DIR, GTTS_SYNTHESIS_FILE)

    async def _inference(self, session: Session, text: str) -> AsyncGenerator[bytes, None]:
        from gtts import gTTS
        with io.BytesIO() as f:
            tts = gTTS(
                text=text,
                lang=self.args.language,
                tld=self.args.tld,
                slow=self.args.slow)
            tts.write_to_fp(f)
            f.seek(0)

            audio: AudioSegment = AudioSegment.from_mp3(f)
            if self.args.speed_increase != 1.0:
                audio = audio.speedup(
                    playback_speed=self.args.speed_increase,
                    chunk_size=self.args.chunk_size,
                    crossfade=self.args.crossfade_lenght)
            audio_resampled = audio.set_frame_rate(
                22050).set_channels(
                1).set_sample_width(2)  # 16bit sample_width 16/8=2
            audio_data = audio_resampled.raw_data
            yield audio_data

    def get_voices(self):
        from gtts.lang import tts_langs
        voices = []
        languages = tts_langs()
        tlds = ['com', 'com.au', 'co.uk', 'us', 'ca', 'co.in', 'ie', 'co.za']

        for lang in languages.keys():
            for tld in tlds:
                voices.append(GTTSArgs(language=lang, tld=tld))

        return voices
