#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch

from speech_bridge_gemma.qwen3_tts_tokenizer_smoke import decode_qwen3_codes, load_qwen3_codec, save_audio


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", required=True)
    parser.add_argument("--out-dir", required=True)
    parser.add_argument("--codec-model", default="Qwen/Qwen3-TTS-Tokenizer-12Hz")
    parser.add_argument("--codec-device", default="cuda")
    return parser.parse_args()


def load_codes(path: str):
    codes = torch.load(path, map_location="cpu")
    if isinstance(codes, dict):
        codes = codes.get("codes") or codes.get("audio_codes")
    return torch.as_tensor(codes).detach().cpu().long()


def main() -> int:
    args = parse_args()
    out_dir = Path(args.out_dir)
    manifest = json.loads(Path(args.manifest).read_text(encoding="utf-8"))
    codec = load_qwen3_codec(args.codec_model, args.codec_device)
    asr_items = []
    for item in manifest["items"]:
        codes_path = item.get("out_codes")
        if not codes_path:
            continue
        codes = load_codes(codes_path)
        wav, sr = decode_qwen3_codes(codec, codes)
        wav_path = out_dir / f"{Path(codes_path).stem.replace('_codes', '')}.wav"
        save_audio(wav_path, wav, sr)
        item["prediction"] = str(wav_path)
        item["sample_rate"] = int(sr)
        asr_items.append({"id": item["id"], "wav": str(wav_path), "expected": item["answer"]})
        print(json.dumps({"event": "decoded", "id": item["id"], "wav": str(wav_path), "frames": int(codes.shape[-1])}, ensure_ascii=False), flush=True)
    (out_dir / "decoded_manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    (out_dir / "asr_items.json").write_text(json.dumps(asr_items, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(json.dumps({"event": "done", "decoded": len(asr_items)}, ensure_ascii=False), flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())