import os
import gradio as gr
import requests

API_KEY = os.environ.get("SILICONFLOW_API_KEY", "")

ASR_URL = "https://api.siliconflow.cn/v1/audio/transcriptions"
CHAT_URL = "https://api.siliconflow.cn/v1/chat/completions"
ASR_MODEL = "FunAudioLLM/SenseVoiceSmall"
MT_MODEL = "tencent/Hunyuan-MT-7B"

LANG_MAP = {
    "英语 → 中文": {"src": "English", "tgt": "Chinese", "src_tag": "EN", "tgt_tag": "中"},
    "汉语 → 英文": {"src": "Chinese", "tgt": "English", "src_tag": "中", "tgt_tag": "EN"},
    "中文 → 法语": {"src": "Chinese", "tgt": "French", "src_tag": "中", "tgt_tag": "FR"},
    "法语 → 中文": {"src": "French", "tgt": "Chinese", "src_tag": "FR", "tgt_tag": "中"},
}


def transcribe(audio_path):
    with open(audio_path, "rb") as f:
        files = {"file": (os.path.basename(audio_path), f, "audio/wav")}
        data = {"model": ASR_MODEL}
        headers = {"Authorization": "Bearer " + API_KEY}
        r = requests.post(ASR_URL, headers=headers, data=data, files=files, timeout=60)
    r.raise_for_status()
    text = r.json().get("text", "").strip()
    text = text.replace("<|", "").replace("|>", "")  # 简单清理特殊标记
    return text


def translate(text, src, tgt):
    headers = {"Authorization": "Bearer " + API_KEY, "Content-Type": "application/json"}
    payload = {
        "model": MT_MODEL,
        "messages": [
            {
                "role": "system",
                "content": f"You are a professional translator. Translate the user text from {src} to {tgt}. "
                           f"Output ONLY the translated text, no explanations, no quotes, no pinyin.",
            },
            {"role": "user", "content": text},
        ],
        "temperature": 0.3,
        "max_tokens": 512,
    }
    r = requests.post(CHAT_URL, headers=headers, json=payload, timeout=60)
    r.raise_for_status()
    return r.json()["choices"][0]["message"]["content"].strip()


def process(audio_path, mode_label, history):
    history = history or []

    if not API_KEY:
        history.append({"role": "assistant", "content": "⚠️ 服务器未配置 SILICONFLOW_API_KEY，请检查 Space 的 Settings -> Repository secrets"})
        return history, None

    if audio_path is None:
        return history, None

    cfg = LANG_MAP[mode_label]

    try:
        text = transcribe(audio_path)
        if not text:
            return history, None
        trans = translate(text, cfg["src"], cfg["tgt"])
    except Exception as e:
        history.append({"role": "assistant", "content": f"⚠️ 出错了：{e}"})
        return history, None

    history.append({"role": "user", "content": f"[{cfg['src_tag']}] {text}"})
    history.append({"role": "assistant", "content": f"[{cfg['tgt_tag']}] {trans}"})
    return history, None


with gr.Blocks(title="传声筒") as demo:
    gr.Markdown("## 🎙️ 传声筒 · 实时语音传译\n点麦克风开始录音，说一句完整的话，再点一次停止，自动识别+翻译。")

    mode = gr.Radio(
        choices=list(LANG_MAP.keys()),
        value="英语 → 中文",
        label="翻译方向",
    )

    audio_in = gr.Audio(
        sources=["microphone"],
        type="filepath",
        label="点击录音，再点一次停止",
    )

    chatbot = gr.Chatbot(label="传译记录", height=420)

    audio_in.stop_recording(
        fn=process,
        inputs=[audio_in, mode, chatbot],
        outputs=[chatbot, audio_in],
    )

demo.launch()