speech/app/asr/router.py

from __future__ import annotations

import io
import json
import struct
import wave
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

import httpx
from fastapi import APIRouter, File, Form, HTTPException, UploadFile, WebSocket, WebSocketDisconnect
from fastapi.responses import FileResponse

from asr.config import (
    DEFAULT_BACKEND,
    DEFAULT_LANGUAGE,
    DEFAULT_MODEL,
    FASTER_WHISPER_URL,
    QWEN3_URL,
    RESULT_DIR,
    UPLOAD_DIR,
    ensure_runtime_dirs,
)

router = APIRouter()
ws_router = APIRouter()

BACKENDS = {
    "faster-whisper": FASTER_WHISPER_URL,
    "qwen3": QWEN3_URL,
}

REALTIME_SAMPLE_RATE = 16000
REALTIME_CHUNK_SECONDS = 3


# ─── Config ────────────────────────────────────────────────────────────────────

@router.get("/config")
def config() -> Dict[str, Any]:
    return {
        "default_backend": DEFAULT_BACKEND,
        "default_model": DEFAULT_MODEL,
        "default_language": DEFAULT_LANGUAGE,
        "backends": {
            "faster-whisper": {
                "models": ["tiny", "base", "small", "medium", "large-v3", "large-v2", "turbo"],
            },
            "qwen3": {
                "models": ["Qwen/Qwen3-ASR-2B", "Qwen/Qwen3-ASR-8B"],
            },
        },
    }


# ─── History ──────────────────────────────────────────────────────────────────

@router.get("/history")
def list_history() -> List[Dict[str, Any]]:
    ensure_runtime_dirs()
    records = []
    for p in sorted(RESULT_DIR.glob("*.json"), reverse=True):
        try:
            data = json.loads(p.read_text(encoding="utf-8"))
            upload_path = data.get("upload_file", "")
            upload_name = Path(upload_path).name if upload_path else ""
            records.append({
                "id": p.stem,
                "filename": data.get("_filename", upload_name),
                "backend": data.get("gateway_backend", data.get("backend", "")),
                "model": data.get("model", ""),
                "language": data.get("language", ""),
                "duration": data.get("duration"),
                "text_preview": (data.get("text", "") or "")[:120],
                "upload_file": upload_name,
                "diarized": data.get("diarized", False),
            })
        except Exception:
            continue
    return records


@router.get("/history/{record_id}")
def get_history(record_id: str) -> Dict[str, Any]:
    safe_id = Path(record_id).name
    p = RESULT_DIR / f"{safe_id}.json"
    if not p.exists():
        raise HTTPException(status_code=404, detail="Record not found")
    return json.loads(p.read_text(encoding="utf-8"))


@router.delete("/history/{record_id}")
def delete_history(record_id: str) -> Dict[str, str]:
    safe_id = Path(record_id).name
    p = RESULT_DIR / f"{safe_id}.json"
    if not p.exists():
        raise HTTPException(status_code=404, detail="Record not found")

    # Also remove upload file if referenced
    try:
        data = json.loads(p.read_text(encoding="utf-8"))
        upload_path = data.get("upload_file", "")
        if upload_path:
            up = Path(upload_path)
            if up.exists() and up.is_relative_to(UPLOAD_DIR):
                up.unlink(missing_ok=True)
    except Exception:
        pass

    p.unlink(missing_ok=True)
    return {"status": "deleted", "id": safe_id}


@router.get("/uploads/{filename}")
def serve_upload(filename: str) -> FileResponse:
    safe_name = Path(filename).name
    p = UPLOAD_DIR / safe_name
    if not p.exists():
        raise HTTPException(status_code=404, detail="File not found")
    return FileResponse(p)


# ─── File Transcription ───────────────────────────────────────────────────────

@router.post("/transcribe")
async def transcribe(
    file: UploadFile = File(...),
    backend: str = Form(DEFAULT_BACKEND),
    model: str = Form(DEFAULT_MODEL),
    custom_model_path: Optional[str] = Form(None),
    language: Optional[str] = Form(None),
    task: str = Form("transcribe"),
    # faster-whisper options
    beam_size: int = Form(5),
    temperature: float = Form(0.0),
    word_timestamps: bool = Form(False),
    diarize: bool = Form(False),
    num_speakers: Optional[int] = Form(None),
    min_speakers: Optional[int] = Form(None),
    max_speakers: Optional[int] = Form(None),
    no_repeat_ngram_size: int = Form(0),
    repetition_penalty: float = Form(1.0),
    compression_ratio_threshold: float = Form(2.4),
    log_prob_threshold: float = Form(-1.0),
    no_speech_threshold: float = Form(0.6),
    condition_on_previous_text: bool = Form(True),
) -> Dict[str, Any]:
    if backend not in BACKENDS:
        raise HTTPException(status_code=400, detail=f"Unsupported backend: {backend}")

    ensure_runtime_dirs()
    original_filename = file.filename or "upload.bin"
    suffix = Path(original_filename).suffix or ".bin"
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
    saved_upload = UPLOAD_DIR / f"{timestamp}{suffix}"

    content = await file.read()
    saved_upload.write_bytes(content)

    try:
        worker_url = BACKENDS[backend]
        effective_language = language or DEFAULT_LANGUAGE

        if backend == "qwen3":
            data: Dict[str, str] = {
                "model": model,
                "language": effective_language,
                "task": task,
            }
        else:
            # faster-whisper
            data = {
                "model": model,
                "custom_model_path": custom_model_path or "",
                "language": effective_language,
                "task": task,
                "beam_size": str(beam_size),
                "temperature": str(temperature),
                "word_timestamps": str(word_timestamps).lower(),
                "diarize": str(diarize).lower(),
                "no_repeat_ngram_size": str(no_repeat_ngram_size),
                "repetition_penalty": str(repetition_penalty),
                "compression_ratio_threshold": str(compression_ratio_threshold),
                "log_prob_threshold": str(log_prob_threshold),
                "no_speech_threshold": str(no_speech_threshold),
                "condition_on_previous_text": str(condition_on_previous_text).lower(),
            }
            if num_speakers is not None:
                data["num_speakers"] = str(num_speakers)
            if min_speakers is not None:
                data["min_speakers"] = str(min_speakers)
            if max_speakers is not None:
                data["max_speakers"] = str(max_speakers)

        files_payload = {
            "file": (original_filename, content, file.content_type or "application/octet-stream")
        }
        async with httpx.AsyncClient(timeout=httpx.Timeout(3600.0, connect=60.0)) as client:
            response = await client.post(f"{worker_url}/transcribe", data=data, files=files_payload)

        if response.status_code >= 400:
            raise HTTPException(status_code=response.status_code, detail=f"Backend error: {response.text}")

        payload = response.json()
        payload["gateway_backend"] = backend
        payload["_filename"] = original_filename
        payload["_timestamp"] = timestamp
        payload["upload_file"] = str(saved_upload)

        result_path = RESULT_DIR / f"{timestamp}.json"
        result_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
        payload["result_file"] = str(result_path)
        return payload

    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"Gateway error: {exc}") from exc


# ─── Real-time WebSocket (faster-whisper only) ────────────────────────────────

def _pcm_float32_to_wav(pcm_bytes: bytes, sample_rate: int = 16000) -> bytes:
    n = len(pcm_bytes) // 4
    try:
        import numpy as np
        arr = np.frombuffer(pcm_bytes, dtype=np.float32)
        arr_i16 = np.clip(arr * 32767.0, -32768, 32767).astype(np.int16)
        raw16 = arr_i16.tobytes()
    except ImportError:
        raw16 = b"".join(
            struct.pack("<h", max(-32768, min(32767, int(struct.unpack("<f", pcm_bytes[i*4:(i+1)*4])[0] * 32767))))
            for i in range(n)
        )
    buf = io.BytesIO()
    with wave.open(buf, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(sample_rate)
        wf.writeframes(raw16)
    return buf.getvalue()


async def _send_chunk(wav_bytes: bytes, *, model: str, language: str, beam_size: int,
                      no_repeat_ngram_size: int, repetition_penalty: float,
                      compression_ratio_threshold: float, log_prob_threshold: float,
                      no_speech_threshold: float, condition_on_previous_text: bool) -> Dict[str, Any]:
    data = {
        "model": model,
        "language": language,
        "task": "transcribe",
        "beam_size": str(beam_size),
        "temperature": "0",
        "word_timestamps": "false",
        "diarize": "false",
        "no_repeat_ngram_size": str(no_repeat_ngram_size),
        "repetition_penalty": str(repetition_penalty),
        "compression_ratio_threshold": str(compression_ratio_threshold),
        "log_prob_threshold": str(log_prob_threshold),
        "no_speech_threshold": str(no_speech_threshold),
        "condition_on_previous_text": str(condition_on_previous_text).lower(),
    }
    files_payload = {"file": ("chunk.wav", wav_bytes, "audio/wav")}
    async with httpx.AsyncClient(timeout=httpx.Timeout(300.0, connect=30.0)) as client:
        resp = await client.post(f"{FASTER_WHISPER_URL}/transcribe", data=data, files=files_payload)
    resp.raise_for_status()
    return resp.json()


@ws_router.websocket("/ws/realtime")
async def realtime_ws(websocket: WebSocket) -> None:
    if not websocket.session.get("user"):
        await websocket.close(code=4401)
        return
    await websocket.accept()

    audio_buf = bytearray()
    partial_texts: List[str] = []
    all_segments: List[Dict[str, Any]] = []
    time_offset = 0.0

    try:
        cfg = json.loads(await websocket.receive_text())

        model = cfg.get("model", DEFAULT_MODEL)
        language = cfg.get("language") or DEFAULT_LANGUAGE
        beam_size = int(cfg.get("beam_size", 5))
        no_repeat_ngram_size = int(cfg.get("no_repeat_ngram_size", 0))
        repetition_penalty = float(cfg.get("repetition_penalty", 1.0))
        compression_ratio_threshold = float(cfg.get("compression_ratio_threshold", 2.4))
        log_prob_threshold = float(cfg.get("log_prob_threshold", -1.0))
        no_speech_threshold = float(cfg.get("no_speech_threshold", 0.6))
        condition_on_previous_text = bool(cfg.get("condition_on_previous_text", True))
        chunk_seconds = int(cfg.get("chunk_seconds", REALTIME_CHUNK_SECONDS))
        chunk_bytes = chunk_seconds * REALTIME_SAMPLE_RATE * 4

        await websocket.send_json({"type": "ready"})

        async def process_buffer(buf: bytearray) -> None:
            nonlocal time_offset
            if len(buf) < REALTIME_SAMPLE_RATE * 4 * 0.3:
                return
            wav = _pcm_float32_to_wav(bytes(buf))
            result = await _send_chunk(
                wav, model=model, language=language, beam_size=beam_size,
                no_repeat_ngram_size=no_repeat_ngram_size,
                repetition_penalty=repetition_penalty,
                compression_ratio_threshold=compression_ratio_threshold,
                log_prob_threshold=log_prob_threshold,
                no_speech_threshold=no_speech_threshold,
                condition_on_previous_text=condition_on_previous_text,
            )
            text = result.get("text", "").strip()
            segs = result.get("segments", [])
            for seg in segs:
                seg["start"] = round(seg.get("start", 0) + time_offset, 3)
                seg["end"] = round(seg.get("end", 0) + time_offset, 3)
            time_offset += len(buf) / 4 / REALTIME_SAMPLE_RATE
            if text:
                partial_texts.append(text)
                all_segments.extend(segs)
            await websocket.send_json({
                "type": "partial",
                "text": text,
                "segments": segs,
                "full_text": " ".join(x for x in partial_texts if x),
            })

        while True:
            data = await websocket.receive()
            if data.get("type") == "websocket.disconnect":
                break
            if "bytes" in data:
                audio_buf.extend(data["bytes"])
                while len(audio_buf) >= chunk_bytes:
                    await process_buffer(bytearray(audio_buf[:chunk_bytes]))
                    audio_buf = audio_buf[chunk_bytes:]
            elif "text" in data:
                msg = json.loads(data["text"])
                if msg.get("cmd") == "stop":
                    if audio_buf:
                        await process_buffer(audio_buf)
                    await websocket.send_json({
                        "type": "final",
                        "text": " ".join(x for x in partial_texts if x),
                        "segments": all_segments,
                        "diarized": False,
                    })
                    break

    except WebSocketDisconnect:
        pass
    except Exception as e:
        try:
            await websocket.send_json({"type": "error", "detail": str(e)})
        except Exception:
            pass