import gradio as gr
from gradio_webrtc import WebRTC, ReplyOnPause, AdditionalOutputs
import anthropic
from pyht import Client as PyHtClient, TTSOptions
import dataclasses
import os
import numpy as np
from huggingface_hub import InferenceClient
import io
from pydub import AudioSegment
from dotenv import load_dotenv

load_dotenv()

account_sid = os.environ.get("TWILIO_ACCOUNT_SID")
auth_token = os.environ.get("TWILIO_AUTH_TOKEN")

if account_sid and auth_token:
    from twilio.rest import Client
    client = Client(account_sid, auth_token)

    token = client.tokens.create()

    rtc_configuration = {
        "iceServers": token.ice_servers,
        "iceTransportPolicy": "relay",
    }
else:
    rtc_configuration = None


@dataclasses.dataclass
class Clients:
    claude: anthropic.Anthropic
    play_ht: PyHtClient
    hf: InferenceClient


tts_options = TTSOptions(voice="s3://voice-cloning-zero-shot/775ae416-49bb-4fb6-bd45-740f205d20a1/jennifersaad/manifest.json",
                         sample_rate=24000)


def aggregate_chunks(chunks_iterator):
    leftover = b''  # Store incomplete bytes between chunks
    
    for chunk in chunks_iterator:
        # Combine with any leftover bytes from previous chunk
        current_bytes = leftover + chunk
        
        # Calculate complete samples
        n_complete_samples = len(current_bytes) // 2  # int16 = 2 bytes
        bytes_to_process = n_complete_samples * 2
        
        # Split into complete samples and leftover
        to_process = current_bytes[:bytes_to_process]
        leftover = current_bytes[bytes_to_process:]
        
        if to_process:  # Only yield if we have complete samples
            audio_array = np.frombuffer(to_process, dtype=np.int16).reshape(1, -1)
            yield audio_array


def audio_to_bytes(audio: tuple[int, np.ndarray]) -> bytes:
    audio_buffer = io.BytesIO()
    segment = AudioSegment(
        audio[1].tobytes(),
        frame_rate=audio[0],
        sample_width=audio[1].dtype.itemsize,
        channels=1,
    )
    segment.export(audio_buffer, format="mp3")
    return audio_buffer.getvalue()


def set_api_key(claude_key, play_ht_username, play_ht_key):
    claude_client = anthropic.Anthropic(api_key=claude_key)
    play_ht_client = PyHtClient(user_id=play_ht_username, api_key=play_ht_key)
    gr.Info("Successfully set API keys.", duration=3)
    return Clients(claude=claude_client, play_ht=play_ht_client,
                   hf=InferenceClient()), gr.skip()


def response(audio: tuple[int, np.ndarray], conversation_llm_format: list[dict],
             chatbot: list[dict], client_state: Clients):
    if not client_state:
        raise gr.Error("Please set your API keys first.")
    prompt = client_state.hf.automatic_speech_recognition(audio_to_bytes(audio)).text
    conversation_llm_format.append({"role": "user", "content": prompt})
    response = client_state.claude.messages.create(
        model="claude-3-5-haiku-20241022",
        max_tokens=512,
        messages=conversation_llm_format,
    )
    response_text = " ".join(block.text for block in response.content if getattr(block, "type", None) == "text")
    conversation_llm_format.append({"role": "assistant", "content": response_text})
    chatbot.append({"role": "user", "content": prompt})
    chatbot.append({"role": "assistant", "content": response_text})
    yield AdditionalOutputs(conversation_llm_format, chatbot)
    iterator = client_state.play_ht.tts(response_text, options=tts_options, voice_engine="Play3.0")
    for chunk in aggregate_chunks(iterator):
        audio_array = np.frombuffer(chunk, dtype=np.int16).reshape(1, -1)
        yield (24000, audio_array, "mono")


with gr.Blocks() as demo:
    with gr.Group():
        with gr.Row():
            chatbot = gr.Chatbot(label="Conversation", type="messages")
        with gr.Row(equal_height=True):
            with gr.Column(scale=1):
                with gr.Row():
                    claude_key = gr.Textbox(type="password", value=os.getenv("ANTHROPIC_API_KEY"),
                                            label="Enter your Anthropic API Key")
                    play_ht_username = gr.Textbox(type="password",
                                                value=os.getenv("PLAY_HT_USER_ID"),
                                                label="Enter your PlayHt Username")
                    play_ht_key = gr.Textbox(type="password",
                                            value=os.getenv("PLAY_HT_API_KEY"),
                                            label="Enter your PlayHt API Key")
                with gr.Row():
                    set_key_button = gr.Button("Set Keys", variant="primary")
            with gr.Column(scale=5):
                audio = WebRTC(modality="audio", mode="send-receive",
                                label="Audio Stream",
                                rtc_configuration=rtc_configuration)

    client_state = gr.State(None)
    conversation_llm_format = gr.State([])

    set_key_button.click(set_api_key, inputs=[claude_key, play_ht_username, play_ht_key],
                         outputs=[client_state, set_key_button])
    audio.stream(
        ReplyOnPause(response),
        inputs=[audio, conversation_llm_format, chatbot, client_state],
        outputs=[audio]
    )
    audio.on_additional_outputs(lambda l, g: (l, g), outputs=[conversation_llm_format, chatbot])


if __name__ == "__main__":
    demo.launch()