import sys
import numpy as np
import sounddevice as sd
from transformers import AutoProcessor, CohereAsrForConditionalGeneration
from transformers.audio_utils import load_audio
from huggingface_hub import hf_hub_download

# Load model
print("Loading model...")
processor = AutoProcessor.from_pretrained("CohereLabs/cohere-transcribe-03-2026")
model = CohereAsrForConditionalGeneration.from_pretrained(
    "CohereLabs/cohere-transcribe-03-2026",
    device_map="auto"
)

def transcribe_audio(audio, language="en"):
    inputs = processor(audio, sampling_rate=16000, return_tensors="pt", language=language)
    inputs.to(model.device, dtype=model.dtype)
    
    outputs = model.generate(**inputs, max_new_tokens=256)
    text = processor.decode(outputs, skip_special_tokens=True)
    return text

def record_audio(duration, samplerate=16000):
    print(f"Recording for {duration} seconds...")
    audio = sd.rec(int(duration * samplerate), samplerate=samplerate, channels=1, dtype='float32')
    sd.wait()
    return audio.flatten()

# Parse arguments
if len(sys.argv) > 1 and sys.argv[1] == "--mic":
    duration = int(sys.argv[2]) if len(sys.argv) > 2 else 5
    try:
        mic_audio = record_audio(duration)
        print("Transcribing...")
        text = transcribe_audio(mic_audio)
        print(f"\nTranscription:\n{text}\n")
    except OSError as e:
        print(f"Microphone error: {e}")
        print("Hint: Run with nix-shell for PortAudio support")
else:
    print("Loading demo audio...")
    audio_file = hf_hub_download(
        repo_id="CohereLabs/cohere-transcribe-03-2026",
        filename="demo/voxpopuli_test_en_demo.wav",
    )
    audio = load_audio(audio_file, sampling_rate=16000)

    print("Transcribing...")
    text = transcribe_audio(audio)
    print(f"\nTranscription:\n{text}\n")