d3ro-voice/apps/desktop/sidecar/main.py

"""
D3RO-VOICE STT Sidecar (FastAPI HTTP 서버)
faster-whisper + CTranslate2만 사용. torch/pyannote 비의존으로 슬림 배포.

사용법:
    python main.py --port 18765 --models-dir <path>

엔드포인트:
    GET  /health          - 헬스체크
    POST /load            - Whisper 모델 로딩
    POST /transcribe      - 오디오 전사 (multipart)
    POST /download        - 모델 다운로드 시작 (백그라운드)
    GET  /download/status - 다운로드 진행률 조회
    POST /download/cancel - 다운로드 취소
    POST /shutdown        - 서버 종료

주: 화자 구분(diarization)은 Phase 15.5에서 LLM 추정 경로가 primary이며,
pyannote 기반 고정밀 화자 구분은 추후 서버 사이드 API로 제공될 예정.
"""

from __future__ import annotations

import argparse
import fnmatch
import logging
import os
import signal
import sys
import threading
import time
from contextlib import asynccontextmanager
from pathlib import Path
from typing import AsyncGenerator

import numpy as np
import uvicorn
from fastapi import FastAPI, File, Form, UploadFile
from fastapi.responses import JSONResponse

# ── 로깅 설정 ──────────────────────────────────────────────

logging.basicConfig(
    level=logging.INFO,
    format="[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    stream=sys.stdout,
)
logger = logging.getLogger("sidecar")

# ── 전역 상태 ──────────────────────────────────────────────

_model: "WhisperModel | None" = None
_model_id: str | None = None
_gpu_available: bool = False
_server: uvicorn.Server | None = None
_models_dir: Path | None = None

# ── 다운로드 상태 (스레드 공유) ────────────────────────────

_download_lock = threading.Lock()
_download_thread: threading.Thread | None = None
_download_cancel = threading.Event()
_download_state: dict = {
    "status": "idle",  # idle | downloading | done | cancelled | error
    "model_id": None,
    "percent": 0,
    "downloaded_bytes": 0,
    "total_bytes": 0,
    "bytes_per_second": 0,
    "message": None,
}

# faster-whisper가 다운로드하는 파일과 동일한 화이트리스트
_DOWNLOAD_PATTERNS = [
    "config.json",
    "preprocessor_config.json",
    "model.bin",
    "tokenizer.json",
    "vocabulary.*",
]

# faster_whisper.utils._MODELS 매핑 실패 시 폴백
_FALLBACK_REPOS = {
    "large-v3-turbo": "mobiuslabsgmbh/faster-whisper-large-v3-turbo",
    "turbo": "mobiuslabsgmbh/faster-whisper-large-v3-turbo",
}

# ── FastAPI 앱 ─────────────────────────────────────────────


@asynccontextmanager
async def lifespan(_app: FastAPI) -> AsyncGenerator[None, None]:
    """서버 시작/종료 생명주기."""
    logger.info("Sidecar 서버 시작")
    _detect_gpu()
    yield
    logger.info("Sidecar 서버 종료")


app = FastAPI(title="D3RO-VOICE STT Sidecar", lifespan=lifespan)


def _detect_gpu() -> None:
    """GPU(CUDA) 사용 가능 여부를 ctranslate2로 감지한다.

    torch 의존 제거를 위해 ctranslate2의 네이티브 CUDA 감지를 사용한다.
    ctranslate2는 faster-whisper의 백엔드이므로 항상 함께 설치된다.
    """
    global _gpu_available
    try:
        import ctranslate2

        cuda_count = ctranslate2.get_cuda_device_count()
        _gpu_available = cuda_count > 0
        if _gpu_available:
            logger.info("GPU 감지: CUDA 디바이스 %d개", cuda_count)
        else:
            logger.info("GPU 미감지, CPU 모드로 동작")
    except Exception as exc:
        _gpu_available = False
        logger.info("GPU 감지 실패, CPU 모드로 동작: %s", exc)


# ── 모델 다운로드 헬퍼 ─────────────────────────────────────


def _resolve_repo(model_id: str) -> str:
    """모델 ID를 HuggingFace repo ID로 변환한다."""
    if "/" in model_id:
        return model_id
    try:
        from faster_whisper.utils import _MODELS

        if model_id in _MODELS:
            return _MODELS[model_id]
    except Exception:
        pass
    if model_id in _FALLBACK_REPOS:
        return _FALLBACK_REPOS[model_id]
    return f"Systran/faster-whisper-{model_id}"


def _local_model_dir(model_id: str) -> Path | None:
    """models_dir 내 다운로드 완료된 모델 디렉토리를 반환한다 (없으면 None)."""
    if _models_dir is None:
        return None
    model_dir = _models_dir / model_id
    if (model_dir / "model.bin").exists():
        return model_dir
    return None


def _set_download_state(**kwargs: object) -> None:
    with _download_lock:
        _download_state.update(kwargs)


def _download_worker(model_id: str) -> None:
    """백그라운드 스레드: HF repo 파일들을 스트리밍 다운로드한다."""
    import requests
    from huggingface_hub import HfApi, hf_hub_url

    try:
        repo_id = _resolve_repo(model_id)
        logger.info("모델 다운로드 시작: %s (repo=%s)", model_id, repo_id)

        api = HfApi()
        info = api.model_info(repo_id, files_metadata=True)
        files = [
            s
            for s in (info.siblings or [])
            if any(fnmatch.fnmatch(s.rfilename, p) for p in _DOWNLOAD_PATTERNS)
        ]
        if not files:
            raise RuntimeError(f"다운로드할 파일이 없습니다: {repo_id}")

        assert _models_dir is not None
        target_dir = _models_dir / model_id
        target_dir.mkdir(parents=True, exist_ok=True)

        total_bytes = sum(s.size or 0 for s in files)
        downloaded = 0
        start_time = time.monotonic()
        _set_download_state(total_bytes=total_bytes, downloaded_bytes=0, percent=0)

        for sibling in files:
            fname = sibling.rfilename
            fsize = sibling.size or 0
            dest = target_dir / fname

            # 멱등: 이미 크기 일치하는 파일은 스킵
            if dest.exists() and fsize > 0 and dest.stat().st_size == fsize:
                downloaded += fsize
                _set_download_state(
                    downloaded_bytes=downloaded,
                    percent=int(downloaded * 100 / total_bytes) if total_bytes else 0,
                )
                logger.info("이미 존재, 스킵: %s", fname)
                continue

            if _download_cancel.is_set():
                raise InterruptedError()

            url = hf_hub_url(repo_id, fname)
            part = dest.with_suffix(dest.suffix + ".part")
            logger.info("다운로드: %s (%.1f MB)", fname, fsize / 1e6)

            with requests.get(url, stream=True, timeout=30) as resp:
                resp.raise_for_status()
                with open(part, "wb") as fh:
                    for chunk in resp.iter_content(chunk_size=1024 * 1024):
                        if _download_cancel.is_set():
                            raise InterruptedError()
                        fh.write(chunk)
                        downloaded += len(chunk)
                        elapsed = time.monotonic() - start_time
                        _set_download_state(
                            downloaded_bytes=downloaded,
                            percent=(
                                int(downloaded * 100 / total_bytes) if total_bytes else 0
                            ),
                            bytes_per_second=int(downloaded / elapsed) if elapsed > 0 else 0,
                        )
            os.replace(part, dest)

        _set_download_state(status="done", percent=100)
        logger.info("모델 다운로드 완료: %s (%.1f MB)", model_id, downloaded / 1e6)

    except InterruptedError:
        _set_download_state(status="cancelled", message="사용자 취소")
        logger.info("모델 다운로드 취소: %s", model_id)
        _cleanup_partial(model_id)
    except Exception as exc:
        _set_download_state(status="error", message=str(exc))
        logger.error("모델 다운로드 실패: %s", exc, exc_info=True)
        _cleanup_partial(model_id)


def _cleanup_partial(model_id: str) -> None:
    """취소/실패 시 .part 잔여 파일 정리."""
    if _models_dir is None:
        return
    target_dir = _models_dir / model_id
    if not target_dir.exists():
        return
    for part in target_dir.glob("*.part"):
        try:
            part.unlink()
        except OSError:
            pass


# ── 엔드포인트 ─────────────────────────────────────────────


@app.get("/health")
async def health() -> JSONResponse:
    """헬스체크. sidecar가 준비되었는지 확인한다."""
    return JSONResponse(
        content={
            "status": "ready" if _model is not None else "no_model",
            "model": _model_id,
            "gpu": _gpu_available,
        }
    )


@app.post("/load")
async def load_model(body: dict) -> JSONResponse:  # noqa: ANN001
    """Whisper 모델을 로딩한다.

    Request body:
        { "model_id": "large-v3" }  -- tiny, base, small, medium, large-v3

    Returns:
        { "status": "loaded", "model_id": "large-v3", "load_time_ms": 1234 }
    """
    global _model, _model_id

    model_id: str = body.get("model_id", "large-v3-turbo")
    logger.info("모델 로딩 시작: %s", model_id)

    start_time = time.monotonic()

    try:
        from faster_whisper import WhisperModel

        device = "cuda" if _gpu_available else "cpu"
        compute_type = "float16" if _gpu_available else "int8"

        # /download로 미리 받아둔 로컬 디렉토리가 있으면 우선 사용.
        # 없으면 faster-whisper의 HF 자동 다운로드 경로로 폴백.
        local_dir = _local_model_dir(model_id)
        model_source = str(local_dir) if local_dir else model_id
        if local_dir:
            logger.info("로컬 모델 디렉토리 사용: %s", local_dir)

        _model = WhisperModel(
            model_source,
            device=device,
            compute_type=compute_type,
        )
        _model_id = model_id

        load_time_ms = int((time.monotonic() - start_time) * 1000)
        logger.info(
            "모델 로딩 완료: %s (device=%s, compute=%s, %dms)",
            model_id,
            device,
            compute_type,
            load_time_ms,
        )

        return JSONResponse(
            content={
                "status": "loaded",
                "model_id": model_id,
                "load_time_ms": load_time_ms,
            }
        )

    except Exception as exc:
        logger.error("모델 로딩 실패: %s", exc)
        return JSONResponse(
            status_code=500,
            content={"status": "error", "message": str(exc)},
        )


@app.post("/transcribe")
async def transcribe(
    audio: UploadFile = File(...),
    language: str = Form("auto"),
    vad_filter: str = Form("true"),
    initial_prompt: str = Form(""),
) -> JSONResponse:
    """오디오 파일을 전사한다.

    Multipart form:
        audio         - PCM16 16kHz mono 바이너리 파일
        language      - 언어 코드 ('auto', 'ko', 'en', ...)
        vad_filter    - VAD 필터 활성화 ('true' / 'false')
        initial_prompt - 초기 프롬프트 (컨텍스트 힌트)
    """
    if _model is None:
        return JSONResponse(
            status_code=503,
            content={"status": "error", "message": "모델이 로딩되지 않았습니다"},
        )

    start_time = time.monotonic()

    try:
        pcm_bytes = await audio.read()

        if len(pcm_bytes) == 0:
            return JSONResponse(
                status_code=400,
                content={"status": "error", "message": "오디오 데이터가 비어있습니다"},
            )

        audio_array = (
            np.frombuffer(pcm_bytes, dtype=np.int16).astype(np.float32) / 32768.0
        )

        sample_rate = 16000
        audio_duration = len(audio_array) / sample_rate

        logger.info(
            "전사 시작: %.1f초 오디오, language=%s, vad=%s",
            audio_duration,
            language,
            vad_filter,
        )

        transcribe_kwargs: dict = {
            "vad_filter": vad_filter.lower() == "true",
            "beam_size": 5,
        }

        if language != "auto":
            transcribe_kwargs["language"] = language

        if initial_prompt:
            transcribe_kwargs["initial_prompt"] = initial_prompt

        # VAD가 전체 오디오를 제거하면 max() 에러 발생 → VAD 없이 재시도
        try:
            segments_iter, info = _model.transcribe(audio_array, **transcribe_kwargs)
        except ValueError as ve:
            if "empty sequence" in str(ve) and transcribe_kwargs.get("vad_filter"):
                logger.warning("VAD가 전체 오디오를 제거함 → VAD 없이 재시도")
                transcribe_kwargs["vad_filter"] = False
                segments_iter, info = _model.transcribe(audio_array, **transcribe_kwargs)
            else:
                raise

        segments_list: list[dict] = []
        full_text_parts: list[str] = []

        for segment in segments_iter:
            seg_dict = {
                "text": segment.text.strip(),
                "start": round(segment.start, 3),
                "end": round(segment.end, 3),
                "avg_logprob": round(segment.avg_logprob, 4),
            }
            segments_list.append(seg_dict)
            full_text_parts.append(segment.text.strip())

        full_text = " ".join(full_text_parts).strip()
        processing_time = int((time.monotonic() - start_time) * 1000)

        detected_language = info.language if info.language else "unknown"

        logger.info(
            "전사 완료: '%s' (lang=%s, %.1f초, %dms)",
            full_text[:80],
            detected_language,
            audio_duration,
            processing_time,
        )

        return JSONResponse(
            content={
                "text": full_text,
                "segments": segments_list,
                "language": detected_language,
                "duration": round(audio_duration, 3),
                "processing_time": processing_time,
            }
        )

    except Exception as exc:
        logger.error("전사 실패: %s", exc, exc_info=True)
        return JSONResponse(
            status_code=500,
            content={"status": "error", "message": str(exc)},
        )


@app.post("/download")
async def download_model(body: dict) -> JSONResponse:  # noqa: ANN001
    """모델 다운로드를 백그라운드로 시작한다.

    Request body:
        { "model_id": "large-v3-turbo" }

    Returns:
        { "status": "started" } 또는 이미 완료된 경우 { "status": "done" }
    """
    global _download_thread

    if _models_dir is None:
        return JSONResponse(
            status_code=500,
            content={"status": "error", "message": "models-dir가 설정되지 않았습니다"},
        )

    model_id: str = body.get("model_id", "large-v3-turbo")

    # 이미 다운로드 완료된 모델이면 즉시 done
    if _local_model_dir(model_id) is not None:
        _set_download_state(status="done", model_id=model_id, percent=100)
        return JSONResponse(content={"status": "done"})

    if _download_thread is not None and _download_thread.is_alive():
        return JSONResponse(
            status_code=409,
            content={"status": "error", "message": "이미 다운로드가 진행 중입니다"},
        )

    _download_cancel.clear()
    _set_download_state(
        status="downloading",
        model_id=model_id,
        percent=0,
        downloaded_bytes=0,
        total_bytes=0,
        bytes_per_second=0,
        message=None,
    )
    _download_thread = threading.Thread(
        target=_download_worker, args=(model_id,), daemon=True
    )
    _download_thread.start()

    return JSONResponse(content={"status": "started"})


@app.get("/download/status")
async def download_status() -> JSONResponse:
    """현재 다운로드 상태를 반환한다."""
    with _download_lock:
        return JSONResponse(content=dict(_download_state))


@app.post("/download/cancel")
async def download_cancel() -> JSONResponse:
    """진행 중인 다운로드를 취소한다."""
    if _download_thread is not None and _download_thread.is_alive():
        _download_cancel.set()
        return JSONResponse(content={"status": "cancelling"})
    return JSONResponse(content={"status": "idle"})


@app.post("/shutdown")
async def shutdown() -> JSONResponse:
    """서버를 graceful하게 종료한다."""
    logger.info("종료 요청 수신")

    if _server is not None:
        _server.should_exit = True

    return JSONResponse(content={"status": "shutting_down"})


# ── 메인 ───────────────────────────────────────────────────


def main() -> None:
    """CLI 진입점."""
    global _server, _models_dir

    parser = argparse.ArgumentParser(description="D3RO-VOICE STT Sidecar")
    parser.add_argument(
        "--port",
        type=int,
        default=18765,
        help="HTTP 서버 포트 (기본: 18765)",
    )
    parser.add_argument(
        "--host",
        type=str,
        default="127.0.0.1",
        help="HTTP 서버 호스트 (기본: 127.0.0.1)",
    )
    parser.add_argument(
        "--models-dir",
        type=str,
        default=None,
        help="사전 다운로드 모델 저장 디렉토리 (미지정 시 HF 캐시만 사용)",
    )
    args = parser.parse_args()

    if args.models_dir:
        _models_dir = Path(args.models_dir)
        _models_dir.mkdir(parents=True, exist_ok=True)

    logger.info("D3RO-VOICE STT Sidecar 시작 (port=%d)", args.port)

    def signal_handler(signum: int, _frame: object) -> None:
        sig_name = signal.Signals(signum).name
        logger.info("시그널 수신: %s, 종료 시작", sig_name)
        if _server is not None:
            _server.should_exit = True

    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    config = uvicorn.Config(
        app=app,
        host=args.host,
        port=args.port,
        log_level="warning",
        access_log=False,
    )
    _server = uvicorn.Server(config)
    _server.run()

    logger.info("Sidecar 서버 종료 완료")


if __name__ == "__main__":
    main()