import json, os, sys, tempfile, wave

WAV_HEADER_SIZE = 44

WHISPER_CACHE = os.path.expanduser("~/.cache/whisper")

def write_wav(path: str, raw_pcm: bytes, sample_rate: int = 16000):
    with wave.open(path, 'wb') as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(sample_rate)
        w.writeframes(raw_pcm)

def transcribe(wav_path: str, model_size: str = "large-v3-turbo", language: str | None = None) -> dict:
    try:
        import whisper
        model = whisper.load_model(model_size, download_root=WHISPER_CACHE)
        opts = {}
        if language:
            opts["language"] = language
        result = model.transcribe(wav_path, **opts)
        return {"success": True, "text": result["text"].strip(), "language": result.get("language", "")}
    except ImportError:
        pass

    try:
        from faster_whisper import WhisperModel
        model_path = model_size
        if not os.path.isdir(model_path):
            local_path = os.path.expanduser(f"~/.cache/huggingface/hub/models--Systran--faster-whisper-{model_size}")
            if os.path.isdir(local_path):
                model_path = local_path
        model = WhisperModel(model_path, device="cpu", compute_type="int8")
        opts = {"beam_size": 1}
        if language:
            opts["language"] = language
        segments, info = model.transcribe(wav_path, **opts)
        text = " ".join(seg.text for seg in segments)
        return {"success": True, "text": text.strip(), "language": info.language}
    except ImportError:
        pass

    return {"success": False, "error": "Neither openai-whisper nor faster-whisper is installed. Run: pip install openai-whisper"}

def main():
    if len(sys.argv) < 2:
        print(json.dumps({"success": False, "error": "Usage: transcribe.py <wav_path> [--model <size>] [--language <code>]"}))
        sys.exit(1)

    wav_path = sys.argv[1]
    model_size = "large-v3-turbo"
    language = None

    i = 2
    while i < len(sys.argv):
        if sys.argv[i] == "--model" and i + 1 < len(sys.argv):
            model_size = sys.argv[i + 1]
            i += 2
        elif sys.argv[i] == "--language" and i + 1 < len(sys.argv):
            language = sys.argv[i + 1]
            i += 2
        else:
            i += 1

    if sys.argv[1] == "--stdin-pcm":
        sample_rate = int(sys.argv[2]) if len(sys.argv) > 2 else 16000
        raw = sys.stdin.buffer.read()
        tmp = tempfile.mktemp(suffix=".wav")
        write_wav(tmp, raw, sample_rate)
        wav_path = tmp
        result = transcribe(wav_path, model_size, language)
        os.unlink(tmp)
    else:
        result = transcribe(wav_path, model_size, language)

    print(json.dumps(result))

if __name__ == "__main__":
    main()