"""
Qwen3-ASR-1.7B vLLM Streaming Server - HuggingFace Space
DashScope-compatible WebSocket protocol with server-side VAD.

Your existing QwenCloudASRSTTService pipecat client works by just changing the URL.

Endpoints:
  GET  /health                    - Health check
  POST /v1/audio/transcriptions   - Batch file transcription
  WS   /v1/realtime               - Streaming ASR (DashScope protocol)
"""

import os
import sys
import json
import time
import base64
import asyncio
import threading
import logging
import tempfile
import uuid
import copy
import concurrent.futures
import numpy as np
import soundfile as sf

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
import uvicorn

# ============================================================
# Logging
# ============================================================
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
log = logging.getLogger("qwen3-asr-vllm")

# ============================================================
# Configuration
# ============================================================
MODEL_ID = os.getenv("MODEL_ID", "Qwen/Qwen3-ASR-1.7B")
GPU_MEMORY_UTILIZATION = float(os.getenv("GPU_MEMORY_UTILIZATION", "0.80"))
STREAMING_MAX_NEW_TOKENS = int(os.getenv("STREAMING_MAX_NEW_TOKENS", "1024"))

# Streaming chunk config — smaller = faster partials, larger = more context
CHUNK_SIZE_SEC = float(os.getenv("CHUNK_SIZE_SEC", "4.0"))
UNFIXED_CHUNK_NUM = int(os.getenv("UNFIXED_CHUNK_NUM", "5"))
UNFIXED_TOKEN_NUM = int(os.getenv("UNFIXED_TOKEN_NUM", "15"))

SAMPLE_RATE = 16000
LANGUAGE = os.getenv("LANGUAGE", "English")

# VAD defaults (can be overridden per-session via session.update)
VAD_THRESHOLD = float(os.getenv("VAD_THRESHOLD", "0.5"))
VAD_MIN_SILENCE_MS = int(os.getenv("VAD_MIN_SILENCE_MS", "800"))
VAD_SPEECH_PAD_MS = int(os.getenv("VAD_SPEECH_PAD_MS", "300"))

# Hallucination filtering
HALLUCINATION_PHRASES = {
    "transcript", "transcription", "thank you", "thanks for watching",
    "you", "bye", "goodbye", "the end", "subtitle", "subtitles",
}

log.info("=" * 60)
log.info("Qwen3-ASR vLLM Streaming Server Config:")
log.info(f"  MODEL_ID              = {MODEL_ID}")
log.info(f"  GPU_MEMORY_UTILIZATION = {GPU_MEMORY_UTILIZATION}")
log.info(f"  STREAMING_MAX_TOKENS  = {STREAMING_MAX_NEW_TOKENS}")
log.info(f"  CHUNK_SIZE_SEC        = {CHUNK_SIZE_SEC}s")
log.info(f"  UNFIXED_CHUNK_NUM     = {UNFIXED_CHUNK_NUM} (revise last {CHUNK_SIZE_SEC * UNFIXED_CHUNK_NUM:.0f}s)")
log.info(f"  UNFIXED_TOKEN_NUM     = {UNFIXED_TOKEN_NUM}")
log.info(f"  LANGUAGE              = {LANGUAGE}")
log.info(f"  VAD_THRESHOLD         = {VAD_THRESHOLD}")
log.info(f"  VAD_MIN_SILENCE_MS    = {VAD_MIN_SILENCE_MS}")
log.info("=" * 60)

# Thread pool for running synchronous model inference off the event loop
_executor = concurrent.futures.ThreadPoolExecutor(max_workers=4)

# ============================================================
# FastAPI App
# ============================================================
app = FastAPI(title="Qwen3-ASR vLLM Streaming", version="2.0.0")

# ============================================================
# ASR Model Loading (singleton, thread-safe)
# ============================================================
_asr_model = None
_asr_lock = threading.Lock()
_model_ready = threading.Event()


def _get_dtype():
    try:
        import torch
        if torch.cuda.is_available():
            cap = torch.cuda.get_device_capability()
            if cap[0] * 10 + cap[1] >= 80:
                return "bfloat16"
    except Exception:
        pass
    return "half"


def get_asr_model():
    global _asr_model
    if _asr_model is None:
        with _asr_lock:
            if _asr_model is None:
                log.info(f"Loading {MODEL_ID}...")
                start = time.time()
                dtype = _get_dtype()

                from qwen_asr import Qwen3ASRModel
                _asr_model = Qwen3ASRModel.LLM(
                    model=MODEL_ID,
                    gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
                    dtype=dtype,
                    max_new_tokens=STREAMING_MAX_NEW_TOKENS,
                )
                log.info(f"Model loaded in {time.time() - start:.1f}s (dtype={dtype})")
                _model_ready.set()
    return _asr_model


# ============================================================
# Silero VAD (per-connection instances via deep copy)
# ============================================================
_vad_model_template = None
_vad_utils = None
_vad_lock = threading.Lock()


def _ensure_vad_loaded():
    """Load the VAD model template once (thread-safe)."""
    global _vad_model_template, _vad_utils
    if _vad_model_template is None:
        with _vad_lock:
            if _vad_model_template is None:
                import torch
                _vad_model_template, _vad_utils = torch.hub.load(
                    "snakers4/silero-vad", "silero_vad",
                    trust_repo=True, verbose=False,
                )
                log.info("Silero VAD model loaded")


def create_vad(threshold=VAD_THRESHOLD, min_silence_ms=VAD_MIN_SILENCE_MS,
               speech_pad_ms=VAD_SPEECH_PAD_MS):
    """Create a new VAD iterator for a WebSocket connection."""
    try:
        _ensure_vad_loaded()
        model_copy = copy.deepcopy(_vad_model_template)
        VADIterator = _vad_utils[3]
        return VADIterator(
            model_copy,
            threshold=threshold,
            sampling_rate=SAMPLE_RATE,
            min_silence_duration_ms=min_silence_ms,
            speech_pad_ms=speech_pad_ms,
        )
    except Exception as e:
        log.warning(f"Silero VAD failed ({e}), using RMS fallback")
        return RMSVad(threshold=0.01, silence_frames=int(min_silence_ms / 20))


class RMSVad:
    """Simple energy-based VAD fallback."""

    def __init__(self, threshold=0.01, silence_frames=30, speech_frames=3):
        self.threshold = threshold
        self.silence_frames = silence_frames
        self.speech_frames = speech_frames
        self.is_speaking = False
        self._silent_count = 0
        self._speech_count = 0

    def __call__(self, audio_chunk):
        import torch
        if isinstance(audio_chunk, np.ndarray):
            audio_chunk = torch.from_numpy(audio_chunk)
        rms = float(torch.sqrt(torch.mean(audio_chunk.float() ** 2)))

        if rms > self.threshold:
            self._speech_count += 1
            self._silent_count = 0
            if not self.is_speaking and self._speech_count >= self.speech_frames:
                self.is_speaking = True
                return {"start": 0}
        else:
            self._silent_count += 1
            self._speech_count = 0
            if self.is_speaking and self._silent_count >= self.silence_frames:
                self.is_speaking = False
                return {"end": 0}
        return None

    def reset_states(self):
        self.is_speaking = False
        self._silent_count = 0
        self._speech_count = 0


# ============================================================
# Streaming Session (per-utterance ASR state)
# ============================================================
def _is_hallucination(text: str) -> bool:
    return text.lower().strip().rstrip(".!?,;:") in HALLUCINATION_PHRASES


class UtteranceSession:
    """Manages vLLM streaming state for a single utterance."""

    def __init__(self, model):
        self.model = model
        self.state = model.init_streaming_state(
            unfixed_chunk_num=UNFIXED_CHUNK_NUM,
            unfixed_token_num=UNFIXED_TOKEN_NUM,
            chunk_size_sec=CHUNK_SIZE_SEC,
            language=LANGUAGE,
        )
        self.is_finalized = False
        self.last_text = ""
        self.total_audio_sec = 0.0

    def feed(self, audio: np.ndarray) -> dict:
        """Feed audio chunk, return {text, delta, is_final}."""
        if self.is_finalized:
            return {"text": self.last_text, "delta": "", "is_final": True}

        audio = np.asarray(audio, dtype=np.float32)
        self.total_audio_sec += len(audio) / SAMPLE_RATE

        self.model.streaming_transcribe(audio, self.state)
        current = self.state.text or ""

        if current and _is_hallucination(current):
            current = self.last_text

        # Compute delta (new text since last update)
        delta = ""
        if current and current != self.last_text:
            if current.startswith(self.last_text):
                delta = current[len(self.last_text):]
            else:
                delta = current  # Full revision occurred

        self.last_text = current
        return {"text": current, "delta": delta, "is_final": False}

    def finalize(self) -> dict:
        """Finalize utterance, return final transcript."""
        if self.is_finalized:
            return {"text": self.last_text, "is_final": True}

        fallback = self.state.text or ""
        self.model.finish_streaming_transcribe(self.state)
        text = self.state.text or fallback

        if text and _is_hallucination(text):
            text = ""

        self.is_finalized = True
        self.last_text = text
        log.info(f"Utterance finalized: {self.total_audio_sec:.1f}s audio -> '{text[:100]}'")
        return {"text": text, "is_final": True}


# ============================================================
# Realtime Connection Manager
# ============================================================
class RealtimeConnection:
    """
    Per-WebSocket-connection state.
    Manages VAD + streaming ASR sessions across multiple utterances.
    """

    # Silero VAD requires at least 512 samples (32ms @ 16kHz)
    MIN_VAD_SAMPLES = 512

    def __init__(self, asr_model, vad):
        self.asr_model = asr_model
        self.vad = vad
        self.utterance = None  # Current UtteranceSession
        self.audio_buffer = np.array([], dtype=np.float32)
        self.vad_buffer = np.array([], dtype=np.float32)
        self.is_speaking = False
        self.chunk_samples = int(CHUNK_SIZE_SEC * SAMPLE_RATE)

    def process_audio(self, audio_f32: np.ndarray) -> list[dict]:
        """
        Process audio chunk through VAD + ASR.
        Returns list of DashScope-compatible events to send to client.
        """
        import torch
        events = []

        # Buffer audio for VAD — Silero requires EXACTLY 512 samples at 16kHz
        self.vad_buffer = np.concatenate([self.vad_buffer, audio_f32])

        # Process in exact 512-sample windows (Silero rejects any other size)
        while len(self.vad_buffer) >= self.MIN_VAD_SAMPLES:
            vad_frame = self.vad_buffer[:self.MIN_VAD_SAMPLES]
            self.vad_buffer = self.vad_buffer[self.MIN_VAD_SAMPLES:]

            vad_result = self.vad(torch.from_numpy(vad_frame))

            # Speech started
            if vad_result and "start" in vad_result and not self.is_speaking:
                self.is_speaking = True
                self.utterance = UtteranceSession(self.asr_model)
                self.audio_buffer = np.array([], dtype=np.float32)
                events.append({"type": "input_audio_buffer.speech_started"})

            # Accumulate audio during speech
            if self.is_speaking and self.utterance:
                self.audio_buffer = np.concatenate([self.audio_buffer, vad_frame])

                # Feed chunks to ASR model when we have enough
                while len(self.audio_buffer) >= self.chunk_samples:
                    chunk = self.audio_buffer[:self.chunk_samples]
                    self.audio_buffer = self.audio_buffer[self.chunk_samples:]

                    result = self.utterance.feed(chunk)

                    if result["delta"]:
                        events.append({
                            "type": "conversation.item.input_audio_transcription.delta",
                            "delta": result["delta"],
                        })
                    if result["text"]:
                        events.append({
                            "type": "conversation.item.input_audio_transcription.text",
                            "text": result["text"],
                        })

            # Speech ended
            if vad_result and "end" in vad_result and self.is_speaking:
                events.extend(self._finalize_utterance())

        return events

    def _finalize_utterance(self) -> list[dict]:
        """Finalize current utterance and return events."""
        events = []

        if not self.utterance:
            return events

        # Feed remaining buffered audio
        if len(self.audio_buffer) > 0:
            result = self.utterance.feed(self.audio_buffer)
            self.audio_buffer = np.array([], dtype=np.float32)
            if result["text"]:
                events.append({
                    "type": "conversation.item.input_audio_transcription.text",
                    "text": result["text"],
                })

        # Finalize
        result = self.utterance.finalize()

        events.append({"type": "input_audio_buffer.speech_stopped"})
        events.append({
            "type": "conversation.item.input_audio_transcription.completed",
            "transcript": result["text"],
        })

        self.utterance = None
        self.is_speaking = False
        return events

    def commit(self) -> list[dict]:
        """Force finalize current utterance (manual commit)."""
        if self.utterance and not self.utterance.is_finalized:
            return self._finalize_utterance()
        return []

    def clear(self) -> list[dict]:
        """Clear audio buffer and discard current utterance."""
        self.audio_buffer = np.array([], dtype=np.float32)
        if self.utterance and not self.utterance.is_finalized:
            self.utterance.is_finalized = True
        self.utterance = None
        self.is_speaking = False
        return [{"type": "input_audio_buffer.cleared"}]


# ============================================================
# HTTP Endpoints
# ============================================================
@app.get("/health")
def health():
    return {
        "status": "ok",
        "model": MODEL_ID,
        "model_status": "loaded" if _asr_model else "loading",
        "vad": "silero" if _vad_model_template else "rms_fallback",
        "streaming_config": {
            "chunk_size_sec": CHUNK_SIZE_SEC,
            "unfixed_chunk_num": UNFIXED_CHUNK_NUM,
            "unfixed_token_num": UNFIXED_TOKEN_NUM,
        },
    }


@app.get("/")
def root():
    return {
        "service": "Qwen3-ASR vLLM Streaming Server",
        "model": MODEL_ID,
        "protocol": "DashScope-compatible (OpenAI realtime=v1)",
        "endpoints": {
            "/health": "GET - Health check",
            "/v1/audio/transcriptions": "POST - Batch file transcription",
            "/v1/realtime": "WebSocket - Streaming with server-side VAD",
        },
        "websocket_protocol": {
            "url": "wss://YOUR-SPACE.hf.space/v1/realtime",
            "input_format": "PCM int16, 16kHz, mono, base64-encoded",
            "events_in": [
                "session.update",
                "input_audio_buffer.append",
                "input_audio_buffer.commit",
                "input_audio_buffer.clear",
            ],
            "events_out": [
                "session.created",
                "session.updated",
                "input_audio_buffer.speech_started",
                "input_audio_buffer.speech_stopped",
                "conversation.item.input_audio_transcription.delta",
                "conversation.item.input_audio_transcription.text",
                "conversation.item.input_audio_transcription.completed",
                "error",
            ],
        },
    }


@app.post("/v1/audio/transcriptions")
async def transcribe(file: UploadFile = File(...)):
    """Batch transcription (OpenAI-compatible)."""
    if not _model_ready.is_set():
        raise HTTPException(503, "Model still loading")

    suffix = os.path.splitext(file.filename or "")[1] or ".wav"
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
        content = await file.read()
        tmp.write(content)
        tmp_path = tmp.name

    try:
        model = get_asr_model()
        result = model.transcribe([tmp_path], language=[LANGUAGE])
        text = result[0].text if result else ""
        if text and _is_hallucination(text):
            text = ""
        return {"text": text, "model": MODEL_ID, "language": LANGUAGE}
    except Exception as e:
        log.error(f"Transcription error: {e}")
        raise HTTPException(500, str(e))
    finally:
        try:
            os.unlink(tmp_path)
        except OSError:
            pass


# ============================================================
# WebSocket Streaming (DashScope-compatible protocol)
# ============================================================
@app.websocket("/v1/realtime")
async def realtime_stream(ws: WebSocket):
    await ws.accept()
    client = ws.client
    log.info(f"WebSocket connected: {client}")

    # Wait for model
    if not _model_ready.is_set():
        await ws.send_json({"type": "error", "error": {"message": "Model loading..."}})
        loop = asyncio.get_event_loop()
        ready = await loop.run_in_executor(None, _model_ready.wait, 300)
        if not ready:
            await ws.close(1011, "Model load timeout")
            return

    # Send session.created
    session_id = f"sess_{uuid.uuid4().hex[:12]}"
    await ws.send_json({
        "type": "session.created",
        "session": {"id": session_id, "model": MODEL_ID},
    })

    model = get_asr_model()

    # Create per-connection VAD (deep copy of shared template)
    vad = create_vad()

    # Create connection manager
    conn = RealtimeConnection(model, vad)

    try:
        while True:
            message = await ws.receive()

            if message["type"] == "websocket.disconnect":
                break

            # --- JSON messages ---
            if "text" in message and message["text"]:
                try:
                    data = json.loads(message["text"])
                except json.JSONDecodeError:
                    await ws.send_json({
                        "type": "error",
                        "error": {"message": "Invalid JSON"},
                    })
                    continue

                event_type = data.get("type", "")

                if event_type == "session.update":
                    # Accept VAD config from client
                    session_cfg = data.get("session", {})
                    td = session_cfg.get("turn_detection", {})
                    if td.get("type") == "server_vad":
                        threshold = td.get("threshold", VAD_THRESHOLD)
                        silence_ms = td.get("silence_duration_ms", VAD_MIN_SILENCE_MS)
                        # Recreate VAD with client-specified params
                        vad = create_vad(
                            threshold=threshold,
                            min_silence_ms=silence_ms,
                        )
                        conn.vad = vad
                        log.info(f"VAD updated: threshold={threshold}, silence={silence_ms}ms")

                    await ws.send_json({
                        "type": "session.updated",
                        "session": {"id": session_id},
                    })

                elif event_type == "input_audio_buffer.append":
                    audio_b64 = data.get("audio", "")
                    if not audio_b64:
                        continue

                    raw = base64.b64decode(audio_b64)
                    chunk_int16 = np.frombuffer(raw, dtype=np.int16)
                    chunk_f32 = chunk_int16.astype(np.float32) / 32768.0

                    # Run VAD + ASR in thread pool (avoid blocking event loop)
                    events = await asyncio.get_event_loop().run_in_executor(
                        _executor, conn.process_audio, chunk_f32,
                    )
                    for ev in events:
                        await ws.send_json(ev)

                elif event_type == "input_audio_buffer.commit":
                    events = await asyncio.get_event_loop().run_in_executor(
                        _executor, conn.commit,
                    )
                    for ev in events:
                        await ws.send_json(ev)

                elif event_type == "input_audio_buffer.clear":
                    events = conn.clear()
                    for ev in events:
                        await ws.send_json(ev)

            # --- Binary audio (alternative to base64 JSON) ---
            elif "bytes" in message and message["bytes"]:
                raw = message["bytes"]
                chunk_int16 = np.frombuffer(raw, dtype=np.int16)
                chunk_f32 = chunk_int16.astype(np.float32) / 32768.0

                events = await asyncio.get_event_loop().run_in_executor(
                    _executor, conn.process_audio, chunk_f32,
                )
                for ev in events:
                    await ws.send_json(ev)

    except WebSocketDisconnect:
        log.info(f"WebSocket disconnected: {client}")
    except Exception as e:
        log.error(f"WebSocket error: {e}", exc_info=True)
    finally:
        # Finalize any in-progress utterance
        if conn.utterance and not conn.utterance.is_finalized:
            try:
                events = await asyncio.get_event_loop().run_in_executor(
                    _executor, conn.commit,
                )
                for ev in events:
                    try:
                        await ws.send_json(ev)
                    except Exception:
                        pass
            except Exception:
                pass
        log.info(f"WebSocket session ended: {client}")


# ============================================================
# Main
# ============================================================
if __name__ == "__main__":
    log.info("=" * 60)
    log.info("Qwen3-ASR vLLM Streaming Server")
    log.info("=" * 60)

    if not os.environ.get("OMP_NUM_THREADS"):
        os.environ["OMP_NUM_THREADS"] = "4"

    # Log system info
    try:
        import torch
        log.info(f"Python: {sys.version.split()[0]}")
        log.info(f"PyTorch: {torch.__version__}")
        log.info(f"CUDA: {torch.cuda.is_available()}")
        if torch.cuda.is_available():
            log.info(f"CUDA version: {torch.version.cuda}")
            for i in range(torch.cuda.device_count()):
                props = torch.cuda.get_device_properties(i)
                mem = getattr(props, "total_memory", 0) or getattr(props, "total_mem", 0)
                log.info(f"GPU {i}: {props.name} ({mem / (1024**3):.1f} GB)")
    except Exception as e:
        log.warning(f"System info: {e}")

    # Pre-load ASR model (blocking)
    log.info("Loading ASR model...")
    try:
        get_asr_model()
        log.info("ASR model ready")
    except Exception as e:
        log.error(f"ASR model load failed: {e}", exc_info=True)

    # Pre-load VAD model template
    log.info("Loading Silero VAD...")
    try:
        _ensure_vad_loaded()
        log.info("Silero VAD ready")
    except Exception as e:
        log.warning(f"Silero VAD failed, RMS fallback will be used: {e}")

    log.info(f"Starting server on 0.0.0.0:7860")
    log.info(f"  Health:    http://0.0.0.0:7860/health")
    log.info(f"  Batch:     http://0.0.0.0:7860/v1/audio/transcriptions")
    log.info(f"  Streaming: ws://0.0.0.0:7860/v1/realtime")

    uvicorn.run(app, host="0.0.0.0", port=7860)