#!/usr/bin/env python3
"""
TTS 语音生成器 - 使用 MiniMax
支持并行生成与时间戳输出
"""

import argparse
import json
import os
import re
import subprocess
import time
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed

import requests


def get_key_file_path():
    repo_root = Path(__file__).resolve().parents[4]
    key_path = repo_root / ".ai" / "skills" / "cgerike-maker" / ".config" / "minimax_api_key.txt"
    if key_path.exists():
        return key_path
    return repo_root / ".trae" / "skills" / "cgerike-maker" / ".config" / "minimax_api_key.txt"


def get_api_key():
    api_key = os.environ.get("MINIMAX_API_KEY")
    if api_key:
        return api_key
    key_path = get_key_file_path()
    if key_path.exists():
        return key_path.read_text(encoding="utf-8").strip()
    key_path.parent.mkdir(parents=True, exist_ok=True)
    api_key = input("请输入 MINIMAX API Key: ").strip()
    if not api_key:
        raise ValueError("MINIMAX API Key 不能为空")
    key_path.write_text(api_key, encoding="utf-8")
    return api_key


def create_speech_task(api_key, text=None, file_id=None):
    url = "https://api.minimax.io/v1/t2a_async_v2"
    payload = {
        "model": "speech-2.8-hd",
        "language_boost": "auto",
        "voice_setting": {
            "voice_id": "", # 请在此填入 MiniMax Voice ID
            "speed": 1,
            "vol": 1,
            "pitch": 0
        },
        "audio_setting": {
            "audio_sample_rate": 32000,
            "bitrate": 128000,
            "format": "mp3",
            "channel": 1
        }
    }
    if text:
        payload["text"] = text
    elif file_id:
        payload["text_file_id"] = file_id
    else:
        raise ValueError("text 或 file_id 必须提供")
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    response = requests.post(url, headers=headers, json=payload)
    response.raise_for_status()
    result = response.json()
    if "task_id" not in result:
        raise ValueError(f"缺少 task_id: {result}")
    return result["task_id"]


def check_task_status(api_key, task_id):
    url = f"https://api.minimax.io/v1/query/t2a_async_query_v2?task_id={task_id}"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "content-type": "application/json"
    }
    response = requests.get(url, headers=headers)
    response.raise_for_status()
    return response.json()


def download_audio_file(api_key, file_id, output_path):
    url = f"https://api.minimax.io/v1/files/retrieve_content?file_id={file_id}"
    headers = {
        "content-type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    response = requests.get(url, headers=headers)
    response.raise_for_status()
    with open(output_path, "wb") as f:
        f.write(response.content)


def split_sentences(text):
    normalized = re.sub(r"\s+", " ", text).strip()
    if not normalized:
        return []
    parts = re.split(r"([。！？!?])", normalized)
    sentences = []
    i = 0
    while i < len(parts):
        seg = parts[i].strip()
        punct = parts[i + 1] if i + 1 < len(parts) else ""
        if seg:
            sentences.append((seg + punct).strip())
        i += 2
    return sentences if sentences else [normalized]


def get_duration(file_path):
    result = subprocess.run(
        ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "csv=p=0", str(file_path)],
        capture_output=True,
        text=True
    )
    return float(result.stdout.strip())


def remux_audio(output_path):
    output_path = Path(output_path)
    temp_file = output_path.with_suffix(output_path.suffix + ".tmp")
    if temp_file.exists():
        temp_file.unlink()
    output_path.replace(temp_file)
    result = subprocess.run(
        ["ffmpeg", "-y", "-i", str(temp_file), "-c:a", "copy", "-map_metadata", "0", str(output_path)],
        capture_output=True,
        text=True
    )
    if result.returncode == 0:
        temp_file.unlink()
    else:
        if output_path.exists():
            output_path.unlink()
        temp_file.replace(output_path)


def generate_sentence_audio(api_key, text, output_path):
    task_id = create_speech_task(api_key, text=text)
    max_retries = 120
    retry_interval = 5
    for _ in range(max_retries):
        status_response = check_task_status(api_key, task_id)
        task_status = status_response.get("status", "").lower()
        if task_status == "success":
            audio_file_id = status_response.get("file_id")
            if not audio_file_id:
                raise ValueError(f"缺少 file_id: {status_response}")
            download_audio_file(api_key, audio_file_id, output_path)
            return output_path
        if task_status == "failed":
            error_msg = status_response.get("error_msg", "Unknown error")
            raise ValueError(f"生成失败: {error_msg}")
        time.sleep(retry_interval)
    raise TimeoutError("音频生成超时")


def concat_audio(part_paths, output_path, temp_dir):
    list_path = Path(temp_dir) / "concat_list.txt"
    with open(list_path, "w", encoding="utf-8") as f:
        for p in part_paths:
            f.write(f"file '{Path(p).as_posix()}'\n")
    result = subprocess.run(
        ["ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", str(list_path), "-c", "copy", str(output_path)],
        capture_output=True,
        text=True
    )
    if result.returncode != 0:
        raise RuntimeError(result.stderr[-1000:])


def generate_tts(text, output_path, with_timestamps=False, max_workers=8):
    api_key = get_api_key()
    sentences = split_sentences(text)
    if not sentences:
        raise ValueError("文本为空")
    output_path = Path(output_path)
    work_dir = output_path.parent
    temp_dir = work_dir / "temp"
    parts_dir = temp_dir / "tts_parts"
    temp_dir.mkdir(exist_ok=True)
    parts_dir.mkdir(exist_ok=True)

    part_paths = [parts_dir / f"{i:03d}.mp3" for i in range(len(sentences))]
    workers = min(max_workers, len(sentences))
    with ThreadPoolExecutor(max_workers=workers) as executor:
        futures = {
            executor.submit(generate_sentence_audio, api_key, sentence, part_path): i
            for i, (sentence, part_path) in enumerate(zip(sentences, part_paths))
        }
        for future in as_completed(futures):
            future.result()

    concat_audio(part_paths, output_path, temp_dir)
    remux_audio(output_path)

    if with_timestamps:
        timestamps = []
        current = 0.0
        for sentence, part_path in zip(sentences, part_paths):
            dur = get_duration(part_path)
            start = round(current, 3)
            end = round(current + dur, 3)
            timestamps.append({
                "text": sentence,
                "start": start,
                "end": end,
                "type": "sentence"
            })
            current = end
        ts_path = output_path.with_suffix(".json")
        with open(ts_path, "w", encoding="utf-8") as f:
            json.dump(timestamps, f, ensure_ascii=False, indent=2)
    return output_path


def main():
    parser = argparse.ArgumentParser(description="MiniMax 语音生成器")
    parser.add_argument("--text", type=str, help="要转换的文本")
    parser.add_argument("--file", type=str, help="文本文件路径")
    parser.add_argument("--output", type=str, required=True, help="输出文件路径")
    parser.add_argument("--timestamps", action="store_true", help="输出时间戳JSON文件")
    parser.add_argument("--max-workers", type=int, default=8, help="并行数量")
    args = parser.parse_args()

    if args.file:
        with open(args.file, "r", encoding="utf-8") as f:
            text = f.read()
    else:
        text = args.text

    if not text:
        parser.print_help()
        return
    generate_tts(text, args.output, args.timestamps, args.max_workers)


if __name__ == "__main__":
    main()
