# Dependencies: dashscope >= 1.23.9, pyaudio
import os
import base64
import time

import pyaudio
from dashscope.audio.qwen_omni import MultiModality, AudioFormat, OmniRealtimeCallback, OmniRealtimeConversation
import dashscope


url = f'wss://[workspace-id].ap-southeast-1.maas.aliyuncs.com/api-ws/v1/realtime'
# API key: if DASHSCOPE_API_KEY is not set, use: dashscope.api_key = "sk-xxx"
dashscope.api_key = os.getenv('DASHSCOPE_API_KEY')
# Voice
voice = 'Ethan'
# Model
model = 'qwen-omni-turbo-realtime'
# Assistant instructions
instructions = (
    "You are Xiaoyun, a personal assistant. Answer the user's questions in a humorous and witty way."
)
class SimpleCallback(OmniRealtimeCallback):
    def __init__(self, pya):
        self.pya = pya
        self.out = None
    def on_open(self):
        # Initialize audio output stream
        self.out = self.pya.open(
            format=pyaudio.paInt16,
            channels=1,
            rate=24000,
            output=True
        )
    def on_event(self, response):
        if response['type'] == 'response.audio.delta':
            # Play audio
            self.out.write(base64.b64decode(response['delta']))
        elif response['type'] == 'conversation.item.input_audio_transcription.completed':
            # Print user transcript
            print(f"[User] {response['transcript']}")
        elif response['type'] == 'response.audio_transcript.done':
            # Print assistant transcript
            print(f"[LLM] {response['transcript']}")

# 1. Initialize audio device
pya = pyaudio.PyAudio()
# 2. Create callback and conversation
callback = SimpleCallback(pya)
conv = OmniRealtimeConversation(model=model, callback=callback, url=url)
# 3. Connect and configure session
conv.connect()
conv.update_session(output_modalities=[MultiModality.AUDIO, MultiModality.TEXT], voice=voice, instructions=instructions)
# 4. Initialize microphone input stream
mic = pya.open(format=pyaudio.paInt16, channels=1, rate=16000, input=True)
# 5. Main loop: stream microphone audio
print("Conversation started. Speak into the microphone (Ctrl+C to exit)...")
try:
    while True:
        audio_data = mic.read(3200, exception_on_overflow=False)
        conv.append_audio(base64.b64encode(audio_data).decode())
        time.sleep(0.01)
except KeyboardInterrupt:
    # Clean up
    conv.close()
    mic.close()
    callback.out.close()
    pya.terminate()
    print("\nConversation ended.")