from __future__ import annotations

import argparse
import base64
import json
import urllib.request
import wave
from pathlib import Path


ENDPOINT = "https://openspeech.bytedance.com/api/v3/tts/unidirectional/sse"


def load_config(path: Path) -> dict:
    with path.open("r", encoding="utf-8") as f:
        return json.load(f)


def save_config(path: Path, cfg: dict) -> None:
    tmp = path.with_suffix(".tmp")
    with tmp.open("w", encoding="utf-8") as f:
        json.dump(cfg, f, ensure_ascii=False, indent=2)
        f.write("\n")
    tmp.replace(path)


def fetch_pcm(api_key: str, voice_type: str, resource_id: str, text: str) -> bytes:
    payload = json.dumps({
        "user": {"uid": "ivox-clone"},
        "namespace": "BidirectionalTTS",
        "req_params": {
            "text": text,
            "speaker": voice_type,
            "audio_params": {"format": "pcm", "sample_rate": 48000},
        },
    }).encode()
    req = urllib.request.Request(ENDPOINT, data=payload, method="POST")
    req.add_header("Content-Type", "application/json")
    req.add_header("X-Api-Key", api_key)
    req.add_header("X-Api-Resource-Id", resource_id)
    req.add_header("X-Api-Request-Id", "ivox-clone")

    pcm = bytearray()
    with urllib.request.urlopen(req, timeout=60) as resp:
        for line_bytes in resp:
            line = line_bytes.decode("utf-8", errors="replace").strip()
            if not line.startswith("data:"):
                continue
            payload = line[5:].strip()
            if not payload or payload == "[DONE]":
                continue
            try:
                event = json.loads(payload)
            except json.JSONDecodeError:
                continue
            code = event.get("code", 0)
            if code not in (0, 20000000):
                raise SystemExit(f"volcengine error: {event.get('message', code)}")
            for key in ("data", "audio", "audio_data"):
                value = event.get(key)
                if isinstance(value, str) and value:
                    pcm.extend(base64.b64decode(value))
    return bytes(pcm[:48000 * 2 * 8])


def write_wav(path: Path, pcm: bytes) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with wave.open(str(path), "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(48000)
        wf.writeframes(pcm)


def main() -> None:
    parser = argparse.ArgumentParser(prog="ivox voice clone volcengine")
    parser.add_argument("--config", required=True)
    parser.add_argument("--api-key", default="")
    parser.add_argument("--voice-type", required=True)
    parser.add_argument("--resource-id", required=True)
    parser.add_argument("--name", required=True)
    parser.add_argument("--text", required=True)
    parser.add_argument("--description", default="")
    args = parser.parse_args()

    config_path = Path(args.config).expanduser().resolve()
    cfg = load_config(config_path)
    api_key = args.api_key or cfg.get("cloneSources", {}).get("volcengine", {}).get("apiKey", "")
    if not api_key:
        raise SystemExit("--api-key is required for volcengine clone")

    ref_name = f"ref_{args.name}.wav"
    ref_rel = f"voices/{ref_name}"
    ref_path = config_path.parent / ref_rel
    pcm = fetch_pcm(api_key, args.voice_type, args.resource_id, args.text)
    if not pcm:
        raise SystemExit("volcengine returned no audio")
    write_wav(ref_path, pcm)

    voices = [v for v in cfg.get("voices", []) if v.get("id") != args.name]
    voices.append({
        "id": args.name,
        "name": args.name,
        "refAudio": ref_rel,
        "refText": args.text,
        "description": args.description or f"cloned from {args.voice_type}",
    })
    cfg["voices"] = voices
    cfg.setdefault("sourceVoices", {})
    save_config(config_path, cfg)
    print(f"cloned voice {args.name}: {ref_rel}")


if __name__ == "__main__":
    main()
