import numpy as np
from transformers import AutoProcessor, CohereAsrForConditionalGeneration
from transformers.audio_utils import load_audio

MODEL_ID = "CohereLabs/cohere-transcribe-03-2026"
SAMPLE_RATE = 16000


def load_model():
    print("Loading model...")
    processor = AutoProcessor.from_pretrained(MODEL_ID)
    model = CohereAsrForConditionalGeneration.from_pretrained(
        MODEL_ID, device_map="auto"
    )
    return processor, model


def transcribe_audio(processor, model, audio, language="en"):
    inputs = processor(audio, sampling_rate=SAMPLE_RATE, return_tensors="pt", language=language)
    inputs.to(model.device, dtype=model.dtype)
    outputs = model.generate(**inputs, max_new_tokens=256)
    texts = processor.decode(outputs, skip_special_tokens=True)
    return " ".join(texts) if isinstance(texts, list) else texts


def record_audio(duration, device=None, normalize=False):
    import sounddevice as sd
    from .vad import resolve_input_rate, resample_to_target

    print(f"Recording for {duration} seconds...")
    rate = resolve_input_rate(device)
    audio = sd.rec(int(duration * rate), samplerate=rate, channels=1, dtype="float32", device=device)
    sd.wait()
    audio = resample_to_target(audio.flatten(), rate)
    if normalize:
        from .compressor import Compressor
        audio = Compressor().process(audio)
    return audio
