#!/usr/bin/env python3
"""STT Server persistente (frigate.vpn:8883).
ASR: whisper.cpp (Vulkan) | Diarization opzionale: pyannote Community-1 (CPU)

Endpoints:
  GET  /health            → {"status": "ok", "model": ...}
  GET  /glossary          → glossario personale (se GLOSSARY_FILE configurato)
  PUT  /glossary          → salva glossario (se GLOSSARY_FILE configurato)
  POST /transcribe        → multipart: file=<audio> + form:
        lang, model, diarize, prompt (initial_prompt personale),
        glossary (JSON string con persone/termini → normalizzazione
        deterministica alias→canonical; il testo originale resta in 'raw')
      diarize=false (default): ASR only → {"text": ..., "raw_text": ...}
      diarize=true:            ASR + diarization → segments con speaker

Uso: uvicorn server:app --host 0.0.0.0 --port 8883
"""
import asyncio, json, os, tempfile, time
from pathlib import Path

from fastapi import FastAPI, File, Form, UploadFile
from fastapi.responses import JSONResponse

import transcribe

app = FastAPI(title="STT Server", version="1.1")

# Semaforo: whisper-cli Vulkan e pyannote non sono thread-safe per uso concorrente
_lock = asyncio.Lock()

# Cache dei modelli già caricati (pyannote pipeline è costosa da caricare)
_pipeline_cache = {}

# Glossario server-side condiviso (opzionale): env GLOSSARY_FILE
_glossary_file = os.environ.get("GLOSSARY_FILE") or ""


def _asr_only(audio_path: str, lang: str, model: str, prompt: str | None) -> list[dict]:
    """Path veloce: solo ASR (whisper.cpp Vulkan), niente diarization."""
    with tempfile.TemporaryDirectory() as tmp:
        segs = transcribe.asr(audio_path, model, lang, tmp, prompt=prompt)
    return segs


def _full_pipeline(audio_path: str, lang: str, model: str, prompt: str | None) -> list[dict]:
    """Path completo: ASR + diarization pyannote."""
    with tempfile.TemporaryDirectory() as tmp:
        segs = transcribe.asr(audio_path, model, lang, tmp, prompt=prompt)
        turns = transcribe.diarize(audio_path, tmp)
        segs = transcribe.merge(segs, turns)
    return segs


@app.get("/health")
async def health():
    return {"status": "ok", "service": "stt", "model": "large-v3-turbo",
            "version": "1.1"}


@app.get("/glossary")
async def get_glossary():
    if not _glossary_file or not os.path.exists(_glossary_file):
        return JSONResponse({"error": "glossario non configurato"}, status_code=404)
    with open(_glossary_file, encoding="utf-8") as f:
        return JSONResponse(json.load(f))


@app.put("/glossary")
async def put_glossary(body: dict):
    if not _glossary_file:
        return JSONResponse({"error": "glossario non configurato (GLOSSARY_FILE)"},
                            status_code=404)
    with open(_glossary_file, "w", encoding="utf-8") as f:
        json.dump(body, f, ensure_ascii=False, indent=2)
    return {"status": "ok", "entries": sum(len(body.get(g) or []) for g in
            ("people", "organizations", "terms", "custom"))}


@app.post("/transcribe")
async def transcribe_endpoint(
    file: UploadFile = File(...),
    lang: str = Form("it"),
    model: str = Form("large-v3-turbo"),
    diarize: bool = Form(False),
    prompt: str | None = Form(None),
    glossary: str | None = Form(None),
):
    t0 = time.time()
    # Salva l'upload in un file temporaneo
    suffix = Path(file.filename or "audio.wav").suffix or ".wav"
    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
        tmp.write(await file.read())
        audio_path = tmp.name
    flat_map = transcribe.load_glossary(glossary) if glossary else {}
    try:
        async with _lock:
            loop = asyncio.get_running_loop()
            if diarize:
                segs = await loop.run_in_executor(None, _full_pipeline, audio_path, lang, model, prompt)
            else:
                segs = await loop.run_in_executor(None, _asr_only, audio_path, lang, model, prompt)
    finally:
        os.unlink(audio_path)

    raw_text = " ".join(s["text"] for s in segs).strip()
    if flat_map:
        segs = transcribe.apply_glossary(segs, flat_map)
    text = " ".join(s["text"] for s in segs).strip()
    elapsed = round(time.time() - t0, 2)
    resp = {
        "text": text,
        "segments": segs,
        "diarize": diarize,
        "lang": lang,
        "model": model,
        "elapsed_s": elapsed,
    }
    # raw_text solo quando c'è stato un intervento (prompt/glossario)
    if flat_map or prompt:
        resp["raw_text"] = raw_text
    return JSONResponse(resp)


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", "8883")))