"""diora's built-in Whisper: a small OpenAI-compatible transcription server.

POST /v1/audio/transcriptions (multipart: file, model, language, prompt, response_format)
returns {text} or, with response_format=verbose_json, {language, duration, text, segments}.
The `model` field is ignored; WHISPER_MODEL decides. The model is loaded on the first request
and freed again after WHISPER_IDLE seconds without work, so it holds no memory while unused.
One transcription at a time; parallel requests wait (the server has few cores to share).
See docs/transcription.md.
"""
import os
import subprocess
import threading
import time

import numpy as np
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from faster_whisper import WhisperModel

MODEL = os.environ.get('WHISPER_MODEL', 'large-v3-turbo')
THREADS = int(os.environ.get('WHISPER_THREADS', '4'))
IDLE = int(os.environ.get('WHISPER_IDLE', '600'))
CACHE = os.environ.get('WHISPER_CACHE', '/models')

app = FastAPI()
work = threading.Lock()      # one transcription at a time
state = {'model': None, 'last': 0.0}


def model():
    if state['model'] is None:
        state['model'] = WhisperModel(MODEL, device='cpu', compute_type='int8', cpu_threads=THREADS,
                                      download_root=CACHE)
    return state['model']


def unloader():
    while True:
        time.sleep(30)
        if state['model'] is not None and time.time() - state['last'] > IDLE and work.acquire(blocking=False):
            try:
                state['model'] = None
            finally:
                work.release()


threading.Thread(target=unloader, daemon=True).start()


def decode(data: bytes) -> np.ndarray:
    p = subprocess.run(['ffmpeg', '-nostdin', '-loglevel', 'error', '-i', 'pipe:0', '-ac', '1', '-ar', '16000',
                        '-f', 'f32le', 'pipe:1'], input=data, capture_output=True)
    if p.returncode != 0:
        raise HTTPException(400, 'not audio: ' + p.stderr.decode('utf-8', 'replace')[-200:])
    return np.frombuffer(p.stdout, np.float32)


@app.get('/health')
def health():
    return {'ok': True, 'model': MODEL, 'loaded': state['model'] is not None}


@app.post('/v1/audio/transcriptions')
def transcriptions(file: UploadFile = File(...), language: str = Form(''), prompt: str = Form(''),
                   response_format: str = Form('json')):
    audio = decode(file.file.read())
    t0 = time.time()
    with work:
        segs, info = model().transcribe(audio, language=language or None, initial_prompt=prompt or None,
                                        beam_size=5, vad_filter=True)
        segs = [{'id': i, 'start': round(s.start, 2), 'end': round(s.end, 2), 'text': s.text}
                for i, s in enumerate(segs)]
        state['last'] = time.time()
    dur = len(audio) / 16000
    print(f'{dur:.0f} s audio in {time.time() - t0:.0f} s', flush=True)
    text = ''.join(s['text'] for s in segs).strip()
    if response_format == 'verbose_json':
        return {'task': 'transcribe', 'language': info.language, 'duration': dur, 'text': text, 'segments': segs}
    return {'text': text}