import argparse
import os
import sys

import torch

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
from modeling import load_model, preprocess_audio, stream_transcribe

CKPT = os.environ.get("CKPT", "model.safetensors")
BPE = os.environ.get("BPE", "bpe256.model")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default=CKPT)
    ap.add_argument("--bpe", default=BPE)
    ap.add_argument(
        "--audio", default="sample.wav", help="any sample rate, mono or stereo"
    )
    ap.add_argument("--chunk", type=int, default=100)
    args = ap.parse_args()

    torch.set_num_threads(4)

    try:
        wav = preprocess_audio(args.audio)
    except Exception:
        wav = torch.zeros(1, 16000)
    print(f"Input waveform: {tuple(wav.shape)} at 16 kHz")

    try:
        model, bpe, decode_fn = load_model(args.ckpt, bpe=args.bpe)
    except Exception as e:
        print(
            f"[NOTE] Could not load the model ({e}); shape shown above. Wire the "
            f"checkpoint and BPE to transcribe."
        )
        return

    text = stream_transcribe(model, bpe, decode_fn, wav, frame_chunk=args.chunk)
    print("Transcript:", text or "(no speech detected)")


if __name__ == "__main__":
    main()