# -*- coding: utf-8 -*-
"""
讯飞 录音文件转写（标准版）- 本地运行脚本
=================================================
用途：将节目音频转为带时间戳的文本，用于TPS中台三方融合demo。

接口：https://raasr.xfyun.cn/v2/api/upload  /  getResult
签名：signa = base64(HmacSHA1(MD5(appid + ts), secretKey))

特性：
- 支持热词列表（hotWord），提升专业术语识别率
- 支持军事领域参数（pd=mil）
- 支持顺滑+口语规整（输出更接近书面语）
- 默认语种 cn（中文普通话），免费包标配，无需 autodialect

依赖：pip install requests

输出文件（与脚本同目录）：
    asr_result_raw.json     - 讯飞原始返回
    asr_result_full.txt     - 纯文本
    asr_result_timed.txt    - 带时间戳的文本（格式同B稿）
"""

import base64
import hashlib
import hmac
import json
import os
import re
import time
import wave
from urllib.parse import quote

import requests

# ========================================================================
# 配置区 —— 填写凭证 + 音频路径
# ========================================================================

# 老版只用 APPID + SecretKey 两个字段
APP_ID = "84eff996"
SECRET_KEY = "在此填入您的SecretKey"

# 音频文件——可以用之前转好的WAV，也可以直接用mp3（讯飞标准版mp3也支持）
AUDIO_FILE = "现代防空反导大对决音频_16k_mono.wav"

# 自动去除空白
APP_ID = APP_ID.strip()
SECRET_KEY = SECRET_KEY.strip()

# ========================================================================
# 接口配置
# ========================================================================

HOST = "https://raasr.xfyun.cn/v2/api"
UPLOAD_URL = HOST + "/upload"
RESULT_URL = HOST + "/getResult"

# 业务参数
LANGUAGE = "cn"           # 中文普通话
PD = "mil"                # 军事领域优化
ENG_SMOOTHPROC = "true"   # 顺滑（去掉"嗯/那个"）
ENG_COLLOQPROC = "true"   # 口语规整

# 热词列表——节目里出现的关键专业术语（讯飞限制200个热词，每个2-16字）
# 我从A稿提取了下面这批，能显著提升识别准确率
HOT_WORDS = [
    # 武器装备
    "萨德", "爱国者", "标准-6", "铁穹", "塔米尔",
    "大卫投石索", "箭式", "箭-2", "箭-3",
    "霍拉姆沙赫尔-4", "见证者-136", "卡桑", "泥石",
    "C-17", "飞毛腿", "M270",
    # 技术概念
    "高超音速", "弹道导弹", "巡航导弹", "无人机蜂群",
    "中段反导", "饱和攻击", "动能撞击", "诱饵弹",
    "相控阵雷达", "电子战", "外大气层",
    # 事件军方
    "真实承诺-4", "伊斯兰革命卫队", "本·古里安", "拉斐尔",
    "黎巴嫩真主党", "波斯湾",
    # 节目栏目
    "军事科技",
]

# 轮询配置
POLL_INTERVAL_SECONDS = 30
MAX_WAIT_MINUTES = 30


# ========================================================================
# 签名+工具
# ========================================================================


def make_signa(app_id, secret_key, ts):
    """
    讯飞老版签名：signa = base64(HmacSHA1(MD5(appid + ts), secretKey))
    """
    base_string = (app_id + ts).encode("utf-8")
    md5_str = hashlib.md5(base_string).hexdigest()  # 32位小写hex
    mac = hmac.new(
        secret_key.encode("utf-8"),
        md5_str.encode("utf-8"),
        digestmod=hashlib.sha1,
    )
    signa = base64.b64encode(mac.digest()).decode("utf-8")
    return signa


def get_audio_duration_ms(filepath):
    """获取音频时长（毫秒）。WAV用内置，MP3用mutagen。"""
    ext = os.path.splitext(filepath)[1].lower()

    if ext == ".wav":
        with wave.open(filepath, "rb") as wf:
            n_frames = wf.getnframes()
            sample_rate = wf.getframerate()
            duration_ms = int(round(n_frames / sample_rate * 1000))
            print(f"[info] WAV规格: {sample_rate}Hz, 声道={wf.getnchannels()}, 位深={wf.getsampwidth()*8}bit")
            return duration_ms

    if ext == ".mp3":
        try:
            from mutagen.mp3 import MP3
            return int(MP3(filepath).info.length * 1000)
        except ImportError:
            print("[warn] mutagen未安装，时长按0传——讯飞文档说该字段不严格校验")
            return 0

    raise ValueError(f"不支持的音频格式: {ext}")


# ========================================================================
# 上传
# ========================================================================


def upload_audio(filepath):
    """上传音频，返回 orderId"""
    if not os.path.exists(filepath):
        raise FileNotFoundError(f"音频文件不存在: {filepath}")

    file_size = os.path.getsize(filepath)
    file_name = os.path.basename(filepath)
    duration_ms = get_audio_duration_ms(filepath)
    ts = str(int(time.time()))

    signa = make_signa(APP_ID, SECRET_KEY, ts)
    print(f"[info] ts={ts}, signa={signa}")
    print(f"[info] 文件: {file_name}, {file_size/1024/1024:.2f}MB, {duration_ms/1000:.1f}秒")

    # 构建URL参数
    params = {
        "appId": APP_ID,
        "signa": signa,
        "ts": ts,
        "fileSize": str(file_size),
        "fileName": file_name,
        "duration": str(duration_ms),
        "language": LANGUAGE,
        "pd": PD,
        "eng_smoothproc": ENG_SMOOTHPROC,
        "eng_colloqproc": ENG_COLLOQPROC,
    }

    # 热词，用 | 分隔
    if HOT_WORDS:
        hot_word_str = "|".join(HOT_WORDS)
        params["hotWord"] = hot_word_str
        print(f"[info] 热词数量: {len(HOT_WORDS)} 个")

    url_parts = [f"{quote(k, safe='')}={quote(str(v), safe='')}" for k, v in params.items()]
    url = f"{UPLOAD_URL}?{'&'.join(url_parts)}"

    headers = {
        "Content-Type": "application/json",
    }

    with open(filepath, "rb") as f:
        audio_bytes = f.read()

    print(f"[info] 正在上传... ({file_size/1024/1024:.1f}MB)")
    resp = requests.post(url, headers=headers, data=audio_bytes, timeout=300)
    print(f"[info] HTTP状态: {resp.status_code}")

    data = resp.json()
    print(f"[info] 响应: {json.dumps(data, ensure_ascii=False, indent=2)}")

    if data.get("code") != "000000":
        raise RuntimeError(f"上传失败: code={data.get('code')}, desc={data.get('descInfo')}")

    order_id = data["content"]["orderId"]
    estimate_ms = data["content"].get("taskEstimateTime", 0)
    print(f"\n[ok] orderId = {order_id}")
    print(f"[ok] 预估完成时间: {estimate_ms/1000:.1f}秒\n")
    return order_id


# ========================================================================
# 查询结果
# ========================================================================


def query_result(order_id):
    """单次查询"""
    ts = str(int(time.time()))
    signa = make_signa(APP_ID, SECRET_KEY, ts)

    params = {
        "appId": APP_ID,
        "signa": signa,
        "ts": ts,
        "orderId": order_id,
        "resultType": "transfer",
    }
    url_parts = [f"{quote(k, safe='')}={quote(str(v), safe='')}" for k, v in params.items()]
    url = f"{RESULT_URL}?{'&'.join(url_parts)}"

    resp = requests.post(url, timeout=30)
    return resp.json()


def poll_until_done(order_id):
    """轮询直到完成"""
    print(f"[info] 开始轮询（每 {POLL_INTERVAL_SECONDS} 秒一次）...")
    start_time = time.time()
    while True:
        elapsed = time.time() - start_time
        if elapsed > MAX_WAIT_MINUTES * 60:
            raise TimeoutError(f"超过 {MAX_WAIT_MINUTES} 分钟未完成")

        data = query_result(order_id)
        order_info = data.get("content", {}).get("orderInfo", {})
        status = order_info.get("status")
        fail_type = order_info.get("failType", 0)
        elapsed_min = elapsed / 60
        print(f"[poll] 已等 {elapsed_min:.1f} 分钟, status={status}, failType={fail_type}")

        if status == 4:
            print(f"\n[ok] 转写完成！耗时 {elapsed_min:.1f} 分钟")
            return data
        if status == -1:
            raise RuntimeError(f"转写失败: failType={fail_type}, 数据: {data}")

        time.sleep(POLL_INTERVAL_SECONDS)


# ========================================================================
# 结果解析
# ========================================================================


def parse_order_result(order_result_str):
    """
    解析嵌套JSON，返回 [(sentence_start_ms, sentence_end_ms, text), ...]
    """
    if not order_result_str:
        return []

    cleaned = re.sub(r"\\\\", r"\\", order_result_str)
    outer = json.loads(cleaned)

    sentences = []
    for item in outer.get("lattice", []):
        inner_str = item.get("json_1best", "")
        if not inner_str:
            continue
        inner = json.loads(inner_str)
        st = inner.get("st", {})
        bg = int(st.get("bg", 0))
        ed = int(st.get("ed", 0))

        words = []
        for rt in st.get("rt", []):
            for ws in rt.get("ws", []):
                for cw in ws.get("cw", []):
                    w = cw.get("w", "").strip()
                    wp = cw.get("wp", "n")
                    if w and wp != "g":
                        words.append(w)
        sentence = "".join(words).strip()
        if sentence:
            sentences.append((bg, ed, sentence))

    return sentences


def format_timestamp(ms):
    total_sec = ms // 1000
    return f"{total_sec // 60}m{total_sec % 60}s"


def write_outputs(sentences, raw_data):
    with open("asr_result_raw.json", "w", encoding="utf-8") as f:
        json.dump(raw_data, f, ensure_ascii=False, indent=2)
    print(f"[ok] 已写 asr_result_raw.json")

    timed_lines = [f"[{format_timestamp(bg)}] {text}" for bg, _, text in sentences]
    with open("asr_result_timed.txt", "w", encoding="utf-8") as f:
        f.write("\n".join(timed_lines))
    print(f"[ok] 已写 asr_result_timed.txt ({len(timed_lines)} 句)")

    full_text = "".join(t for _, _, t in sentences)
    with open("asr_result_full.txt", "w", encoding="utf-8") as f:
        f.write(full_text)
    print(f"[ok] 已写 asr_result_full.txt ({len(full_text)} 字)")


# ========================================================================
# 入口
# ========================================================================


def main():
    print("=" * 60)
    print("讯飞 录音文件转写（标准版） - 本地脚本")
    print("=" * 60)

    if SECRET_KEY.startswith("在此填入"):
        print("[error] 请先填写 SECRET_KEY")
        return

    print(f"[diag] APP_ID = {APP_ID} (长度 {len(APP_ID)})")
    print(f"[diag] SECRET_KEY 长度 {len(SECRET_KEY)}, 首4位={SECRET_KEY[:4]}, 末4位={SECRET_KEY[-4:]}")
    print(f"[diag] AUDIO_FILE = {AUDIO_FILE}")
    print()

    # 1. 上传
    order_id = upload_audio(AUDIO_FILE)

    # 2. 轮询
    final_data = poll_until_done(order_id)

    # 3. 解析+写文件
    order_result_str = final_data["content"]["orderResult"]
    sentences = parse_order_result(order_result_str)
    write_outputs(sentences, final_data)

    print("\n" + "=" * 60)
    print(f"完成！共识别 {len(sentences)} 句")
    print(f"请把 asr_result_timed.txt 发给 Claude 进行三方融合")
    print("=" * 60)


if __name__ == "__main__":
    main()