# 依赖：dashscope >= 1.23.9，pyaudio
import os
import requests
import base64
import pathlib
import time
import pyaudio
from dashscope.audio.qwen_omni import MultiModality, OmniRealtimeCallback, OmniRealtimeConversation
import dashscope

# ======= 常量配置 =======
DEFAULT_TARGET_MODEL = "qwen3.5-omni-plus-realtime"  # 声音复刻、实时对话要使用相同的模型
DEFAULT_PREFERRED_NAME = "guanyu"
DEFAULT_AUDIO_MIME_TYPE = "audio/mpeg"
VOICE_FILE_PATH = "voice.mp3"  # 用于声音复刻的本地音频文件的相对路径


def create_voice(file_path: str,
                 target_model: str = DEFAULT_TARGET_MODEL,
                 preferred_name: str = DEFAULT_PREFERRED_NAME,
                 audio_mime_type: str = DEFAULT_AUDIO_MIME_TYPE) -> str:
  """
  创建音色，并返回 voice 参数
  """
  # 若没有配置环境变量，请用千问云 API Key 将下行替换为：api_key = "sk-xxx"
  api_key = os.getenv("DASHSCOPE_API_KEY")

  file_path_obj = pathlib.Path(file_path)
  if not file_path_obj.exists():
    raise FileNotFoundError(f"音频文件不存在: {file_path}")

  base64_str = base64.b64encode(file_path_obj.read_bytes()).decode()
  data_uri = f"data:{audio_mime_type};base64,{base64_str}"

  url = "https://dashscope.aliyuncs.com/api/v1/services/audio/tts/customization"
  payload = {
    "model": "qwen-voice-enrollment",
    "input": {
      "action": "create",
      "target_model": target_model,
      "preferred_name": preferred_name,
      "audio": {"data": data_uri}
    }
  }
  headers = {
    "Authorization": f"Bearer {api_key}",
    "Content-Type": "application/json"
  }

  resp = requests.post(url, json=payload, headers=headers)
  if resp.status_code != 200:
    raise RuntimeError(f"创建 voice 失败: {resp.status_code}, {resp.text}")

  try:
    return resp.json()["output"]["voice"]
  except (KeyError, ValueError) as e:
    raise RuntimeError(f"解析 voice 响应失败: {e}")


class SimpleCallback(OmniRealtimeCallback):
  def __init__(self, pya):
    self.pya = pya
    self.out = None
  def on_open(self):
    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':
      self.out.write(base64.b64decode(response['delta']))
    elif response['type'] == 'conversation.item.input_audio_transcription.completed':
      print(f"[User] {response['transcript']}")
    elif response['type'] == 'response.audio_transcript.done':
      print(f"[LLM] {response['transcript']}")


if __name__ == '__main__':
  # 若没有配置环境变量，请用千问云 API Key 将下行替换为：dashscope.api_key = "sk-xxx"
  dashscope.api_key = os.getenv("DASHSCOPE_API_KEY")
  url = "wss://dashscope.aliyuncs.com/api-ws/v1/realtime"

  # 1. 声音复刻：创建专属音色
  voice = create_voice(VOICE_FILE_PATH)
  print(f"声音复刻完成，音色: {voice}")

  # 2. 使用复刻音色进行实时对话
  pya = pyaudio.PyAudio()
  callback = SimpleCallback(pya)
  conv = OmniRealtimeConversation(model=DEFAULT_TARGET_MODEL, callback=callback, url=url)
  conv.connect()
  conv.update_session(
    output_modalities=[MultiModality.AUDIO, MultiModality.TEXT],
    voice=voice  # 使用复刻生成的专属音色
  )
  mic = pya.open(format=pyaudio.paInt16, channels=1, rate=16000, input=True)
  print("对话已开始，对着麦克风说话 (Ctrl+C 退出)...")
  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:
    conv.close()
    mic.close()
    callback.out.close()
    pya.terminate()
    print("\n对话结束")