#!/usr/bin/env python3
"""
语音备忘录转文字脚本
使用 AI Builder Space 的语音识别接口 (grok-transcription) 将音频转写为文本，
并将结果保存到 transcripts/ 目录下同名的 .txt 文件中。
"""

import os
import sys
import argparse
from pathlib import Path
import requests

API_URL = "https://space.ai-builders.com/backend/v1/audio/grok-transcription"
SUPPORTED_EXTENSIONS = {".wav", ".mp3", ".m4a", ".aac", ".ogg", ".flac", ".mp4", ".mkv", ".webm"}


def transcribe_audio(audio_path: Path, api_token: str, language: str = None) -> str:
    """调用 AI Builder Space 语音识别接口转写单个音频文件。"""
    headers = {
        "Authorization": f"Bearer {api_token}",
    }
    data = {
        "simple": "true",
    }
    if language:
        data["language"] = language

    with open(audio_path, "rb") as f:
        files = {
            "audio_file": (audio_path.name, f),
        }
        response = requests.post(API_URL, headers=headers, files=files, data=data, timeout=60)

    if response.status_code != 200:
        raise RuntimeError(f"转写失败 ({response.status_code}): {response.text}")

    result = response.json()
    return result.get("text", "")


def process_memos(memos_dir: Path, output_dir: Path, api_token: str, language: str = None):
    """批量转写 memos 目录下的所有音频文件并存为 txt。"""
    if not memos_dir.exists():
        print(f"错误: 目录 {memos_dir} 不存在", file=sys.stderr)
        sys.exit(1)

    output_dir.mkdir(parents=True, exist_ok=True)

    # 扫描音频文件并按文件名排序
    audio_files = sorted(
        [p for p in memos_dir.iterdir() if p.is_file() and p.suffix.lower() in SUPPORTED_EXTENSIONS]
    )

    if not audio_files:
        print(f"在 {memos_dir} 中未找到支持的音频文件。")
        return

    print(f"找到 {len(audio_files)} 个音频文件，开始转写...")
    success_count = 0

    for idx, audio_file in enumerate(audio_files, 1):
        target_file = output_dir / f"{audio_file.stem}.txt"
        print(f"[{idx}/{len(audio_files)}] 正在处理: {audio_file.name} -> {target_file.name} ...", end=" ", flush=True)

        try:
            transcript = transcribe_audio(audio_file, api_token, language=language)
            with open(target_file, "w", encoding="utf-8") as out:
                out.write(transcript.strip() + "\n")
            print("完成")
            success_count += 1
        except Exception as e:
            print(f"失败: {e}", file=sys.stderr)

    print(f"\n全部处理完毕: 成功 {success_count}/{len(audio_files)} 个文件，保存在 {output_dir}/")


def main():
    parser = argparse.ArgumentParser(description="AI Builder Space 语音备忘录转写工具")
    parser.add_argument(
        "--memos-dir",
        type=Path,
        default=Path("memos"),
        help="音频文件所在目录（默认: memos）",
    )
    parser.add_argument(
        "--transcripts-dir",
        type=Path,
        default=Path("transcripts"),
        help="输出文本保存目录（默认: transcripts）",
    )
    parser.add_argument(
        "--language",
        type=str,
        default=None,
        help="语言提示代码（如 zh），留空则自动识别",
    )

    args = parser.parse_args()

    # 从环境变量读取 API key，严禁硬编码或打印
    api_token = os.environ.get("AI_BUILDER_TOKEN")
    if not api_token:
        print("错误: 环境变量 AI_BUILDER_TOKEN 未设置，无法调用转写接口。", file=sys.stderr)
        sys.exit(1)

    process_memos(args.memos_dir, args.transcripts_dir, api_token, language=args.language)


if __name__ == "__main__":
    main()