#!/usr/bin/env python3
"""Batch ASR + Eapil-Voice clone and export generated speech by source name."""

from __future__ import annotations

import argparse
import json
import mimetypes
import os
import re
import subprocess
import sys
import time
from pathlib import Path
from typing import Any

import requests


DEFAULT_BASE_URL = "https://ades.playones.com"
KEY_HELP = "missing API key: set OPENAI_API_KEY, pass --token directly, or get a key from https://kinfo.playones.com"
AUDIO_EXTENSIONS = {".mp3", ".wav", ".m4a", ".flac", ".ogg", ".aac", ".wma"}
END_PUNCTUATION = "。！？.!?~～"
SAFE_NAME_RE = re.compile(r'[<>:"/\\|?*\x00-\x1f]')


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base-url", default=os.getenv("QWEN3_VOICE_API_BASE_URL", DEFAULT_BASE_URL))
    parser.add_argument("--token", default=os.getenv("OPENAI_API_KEY"), required=False)
    parser.add_argument("--input-dir", required=True, help="Directory of source alert audio files")
    parser.add_argument("--reference-audio", required=True, help="Voice clone reference audio")
    parser.add_argument("--output-dir", required=True, help="Directory for generated files")
    parser.add_argument("--format", choices=["wav", "mp3"], default="wav")
    parser.add_argument("--gain-db", type=float, default=0.0, help="Output gain applied during export")
    parser.add_argument("--provider", choices=["qwen", "voxcpm2"], default="qwen")
    parser.add_argument("--asr-language", default="zh")
    parser.add_argument("--clone-language", default="auto")
    parser.add_argument("--timeout", type=int, default=900)
    parser.add_argument("--retries", type=int, default=3)
    parser.add_argument("--overwrite", action="store_true")
    return parser.parse_args()


def log(message: str) -> None:
    print(message, flush=True)


def bearer_token(token: str | None) -> str:
    if not token:
        raise SystemExit(KEY_HELP)
    token = token.strip()
    if token.lower().startswith("bearer "):
        token = token[7:].strip()
    if not token:
        raise SystemExit(KEY_HELP)
    return token


def request_headers(token: str | None) -> dict[str, str]:
    return {"Authorization": f"Bearer {bearer_token(token)}"}


def audio_files(input_dir: Path) -> list[Path]:
    return sorted(
        [path for path in input_dir.iterdir() if path.is_file() and path.suffix.lower() in AUDIO_EXTENSIONS],
        key=lambda path: path.name.casefold(),
    )


def safe_stem(text: str) -> str:
    cleaned = SAFE_NAME_RE.sub("_", text).strip().strip(".")
    return cleaned or "untitled"


def normalize_text(raw: str, fallback: str) -> str:
    text = " ".join((raw or "").replace("\u3000", " ").split()).strip() or fallback.strip()
    text = text.translate(str.maketrans({",": "，", ":": "：", ";": "；", "(": "（", ")": "）"}))
    if text.startswith("哇") and len(text) > 1 and text[1] not in "，,。！？!?":
        text = "哇，" + text[1:]
    if text.startswith("拜托") and "，" not in text[:4] and len(text) > 2:
        text = "拜托，" + text[2:]
    if text and text[-1] not in END_PUNCTUATION:
        text += "。"
    return text


def request_with_retries(
    method: str,
    url: str,
    *,
    retries: int,
    timeout: int,
    **kwargs: Any,
) -> requests.Response:
    last_error: Exception | None = None
    for attempt in range(1, retries + 1):
        try:
            response = requests.request(method, url, timeout=timeout, **kwargs)
            if response.status_code < 500:
                return response
            last_error = RuntimeError(f"HTTP {response.status_code}: {response.text[:500]}")
        except Exception as exc:  # noqa: BLE001
            last_error = exc
        wait = min(30, 2**attempt)
        log(f"retry attempt={attempt}/{retries} wait={wait}s error={last_error}")
        time.sleep(wait)
    if last_error:
        raise last_error
    raise RuntimeError("request failed")


def transcribe(args: argparse.Namespace, source: Path, headers: dict[str, str]) -> dict[str, Any]:
    mime = mimetypes.guess_type(source.name)[0] or "application/octet-stream"
    with source.open("rb") as handle:
        response = request_with_retries(
            "POST",
            f"{args.base_url.rstrip('/')}/v1/asr/transcribe",
            retries=args.retries,
            timeout=args.timeout,
            headers=headers,
            files={"file": (source.name, handle, mime)},
            data={"language": args.asr_language, "temperature": "0"},
        )
    if response.status_code != 200:
        raise RuntimeError(f"ASR HTTP {response.status_code}: {response.text[:1000]}")
    return response.json()


def clone(args: argparse.Namespace, reference_audio: Path, ref_text: str, text: str, headers: dict[str, str]) -> bytes:
    mime = mimetypes.guess_type(reference_audio.name)[0] or "audio/wav"
    with reference_audio.open("rb") as handle:
        response = request_with_retries(
            "POST",
            f"{args.base_url.rstrip('/')}/v1/voice/clone",
            retries=args.retries,
            timeout=args.timeout,
            headers=headers,
            files={"file": (reference_audio.name, handle, mime)},
            data={
                "text": text,
                "ref_text": ref_text,
                "language": args.clone_language,
                "provider": args.provider,
            },
        )
    if response.status_code != 200:
        raise RuntimeError(f"CLONE HTTP {response.status_code}: {response.text[:1000]}")
    return response.content


def export_audio(raw_wav: Path, output: Path, gain_db: float, output_format: str) -> None:
    output.parent.mkdir(parents=True, exist_ok=True)
    if output_format == "wav":
        command = [
            "ffmpeg",
            "-hide_banner",
            "-loglevel",
            "error",
            "-y",
            "-i",
            str(raw_wav),
            "-filter:a",
            f"volume={gain_db:.6f}dB",
            str(output),
        ]
    else:
        command = [
            "ffmpeg",
            "-hide_banner",
            "-loglevel",
            "error",
            "-y",
            "-i",
            str(raw_wav),
            "-filter:a",
            f"volume={gain_db:.6f}dB",
            "-codec:a",
            "libmp3lame",
            "-q:a",
            "2",
            str(output),
        ]
    result = subprocess.run(command, capture_output=True, text=True, encoding="utf-8", errors="replace")
    if result.returncode != 0:
        raise RuntimeError(result.stderr[-2000:] or result.stdout[-2000:] or "ffmpeg failed")


def main() -> int:
    args = parse_args()
    if not args.token:
        log(KEY_HELP)
        return 2

    input_dir = Path(args.input_dir)
    reference_audio = Path(args.reference_audio)
    output_dir = Path(args.output_dir)
    raw_dir = output_dir / "_raw_wav"
    manifest_dir = output_dir / "_manifests"
    output_dir.mkdir(parents=True, exist_ok=True)
    raw_dir.mkdir(parents=True, exist_ok=True)
    manifest_dir.mkdir(parents=True, exist_ok=True)

    headers = request_headers(args.token)
    files = audio_files(input_dir)
    ref_text = reference_audio.stem
    summary: list[dict[str, Any]] = []

    log(f"base_url={args.base_url.rstrip('/')}")
    log(f"input_files={len(files)} output_dir={output_dir} format={args.format} gain_db={args.gain_db}")

    for index, source in enumerate(files, 1):
        output = output_dir / f"{safe_stem(source.stem)}.{args.format}"
        raw_output = raw_dir / f"{safe_stem(source.stem)}.wav"
        try:
            if output.exists() and not args.overwrite:
                record = {"source": str(source), "output": str(output), "status": "skipped", "reason": "exists"}
            else:
                asr = transcribe(args, source, headers)
                text = normalize_text(str(asr.get("text") or ""), source.stem)
                raw_output.write_bytes(clone(args, reference_audio, ref_text, text, headers))
                export_audio(raw_output, output, args.gain_db, args.format)
                record = {
                    "source": str(source),
                    "output": str(output),
                    "raw_wav": str(raw_output),
                    "status": "ok",
                    "text": text,
                    "asr": asr,
                }
                (manifest_dir / f"{safe_stem(source.stem)}.json").write_text(
                    json.dumps(record, ensure_ascii=False, indent=2),
                    encoding="utf-8",
                )
            log(f"{index}/{len(files)} {record['status'].upper()} {source.name}")
        except Exception as exc:  # noqa: BLE001
            record = {"source": str(source), "status": "error", "error": str(exc)}
            log(f"{index}/{len(files)} ERROR {source.name}: {exc}")
        summary.append(record)

    summary_path = output_dir / "conversion_summary.json"
    summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    errors = sum(1 for record in summary if record.get("status") == "error")
    log(f"done files={len(files)} errors={errors} summary={summary_path}")
    return 1 if errors else 0


if __name__ == "__main__":
    sys.exit(main())
