"""Core module with sarvam client and TTS - play and save functionality"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_client.ipynb.

# %% auto 0
__all__ = ['Sarvam', 'is_installed', 'save', 'play']

# %% ../nbs/00_client.ipynb 2
import base64
import os
import shutil
import subprocess
import wave

import requests
from dotenv import load_dotenv
from typing import Iterator

load_dotenv()

# %% ../nbs/00_client.ipynb 3
class Sarvam:
    def __init__(self, api_key: str = os.getenv("SARVAM_API_KEY")):
        self.api_key = api_key

    def get_api_key(self):
        return self.api_key
    
    def text_to_speech(self, text: str, target_language:str, speaker="meera", pitch=0, pace=1, loudness=1, enable_preprocessing=True, model="bulbul:v1"):
        chunk_size = 500
        chunks = [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)]
        response_chunks = []
        for i, chunk in enumerate(chunks):
            headers = {
                "api-subscription-key": self.api_key,
                "Content-Type": "application/json"
            }

            payload = {
                "inputs": [chunk],
                "target_language_code": target_language,
                "speaker": speaker,
                "pitch": pitch,
                "pace": pace,
                "loudness": loudness,
                "enable_preprocessing": enable_preprocessing,
                "model": model
            }

            url = "https://api.sarvam.ai/text-to-speech"
            response = requests.post(url, headers=headers, json=payload)
            
            if response.status_code == 200:
                response_chunks.append(response.json()["audios"][0])
            else:
                raise Exception(f"Error: {response.status_code} {response.text}")
        return {"chunks": response_chunks}


# %% ../nbs/00_client.ipynb 4
def is_installed(lib_name: str) -> bool:
    lib = shutil.which(lib_name)
    if lib is None:
        return False
    return True

# %% ../nbs/00_client.ipynb 6
def save(audio, filename: str) -> None:
    if isinstance(audio["chunks"], list):
        combined_audio = "".join(audio["chunks"])
        b64_file = base64.b64decode(combined_audio)

        with wave.open(filename, "wb") as wav_file:
            # Set the parameters for the .wav file
            wav_file.setnchannels(1)  # Mono audio
            wav_file.setsampwidth(2)  # 2 bytes per sample
            wav_file.setframerate(22050)  # Sample rate of 22050 Hz
            wav_file.writeframes(b64_file)


# %% ../nbs/00_client.ipynb 7
def play(
    audio,
    notebook: bool = False, 
    use_ffmpeg: bool = True
) -> None:
    audio = "".join(audio["chunks"])
    audio = base64.b64decode(audio)

    if notebook:
        try:
            from IPython.display import Audio, display  # type: ignore
        except ModuleNotFoundError:
            message = (
                "`pip install ipython` required when `notebook=False` "
            )
            raise ValueError(message)

        display(Audio(audio, rate=44100, autoplay=True))
    elif use_ffmpeg:
        if not is_installed("ffplay"):
            message = (
                "ffplay from ffmpeg not found, necessary to play audio. "
                "On mac you can install it with 'brew install ffmpeg'. "
                "On linux and windows you can install it from https://ffmpeg.org/"
            )
            raise ValueError(message)
        
        args = ["ffplay", "-autoexit", "-", "-nodisp"]
        proc = subprocess.Popen(
            args=args,
            stdout=subprocess.PIPE,
            stdin=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        out, err = proc.communicate(input=audio)
        proc.poll()
    else:
        try:
            import io

            import sounddevice as sd  # type: ignore
            import soundfile as sf  # type: ignore
        except ModuleNotFoundError:
            message = (
                "`pip install sounddevice soundfile` required when `use_ffmpeg=False` "
            )
            raise ValueError(message)
        sd.play(*sf.read(io.BytesIO(audio)))
        sd.wait()
