validation/app/websocket/handler.py

import asyncio
import logging
import json
import time
import base64
import numpy as np
import torch
from typing import Dict, Optional
from fastapi import WebSocket, WebSocketDisconnect

from ..config import settings
from .protocol import (
    MessageType,
    ClientMessage,
    ClientAudioChunk,
    ClientConfig,
    ClientPing,
    ClientReset,
    AuditLogEntry,
)
from .manager import ConnectionManager
from ..model.inference import InferenceEngine
from ..utils.timing import timestamp_ms


class WebSocketHandler:
    """
    Handles WebSocket message routing with per-connection inference engines.
    """

    def __init__(
        self,
        manager: ConnectionManager,
        engine_factory,
        max_chunk_size_bytes: int = 64 * 1024,
        storage=None,
        alerts=None,
        verification=None,
    ):
        self.manager = manager
        self._engine_factory = engine_factory
        self._engines: Dict[str, InferenceEngine] = {}
        self._max_chunk_size_bytes = max_chunk_size_bytes
        self.storage = storage
        self.alerts = alerts
        self.verification = verification
        self._alert_cb_set = False
        self._loop = None
        self._logger = logging.getLogger(__name__)

    async def handle_connection(self, websocket: WebSocket, connection_id: str) -> None:
        """Main entry point for a WebSocket connection."""
        try:
            await self.manager.connect(websocket, connection_id)
            engine = self._engine_factory(connection_id)
            self._engines[connection_id] = engine
            self._loop = asyncio.get_running_loop()

            if self.alerts is not None and not self._alert_cb_set:
                self.alerts.set_push_callback(self._push_alert)
                self._alert_cb_set = True

            # Send ready message with model info
            await self._send_ready(connection_id, engine)

            # Main message loop
            async for raw_message in websocket.iter_text():
                connection_id = websocket.headers.get("X-Connection-ID") or connection_id
                await self._handle_text(connection_id, engine, raw_message)

        except WebSocketDisconnect:
            self._logger.info(f"Client disconnected: {connection_id}")
        except Exception as e:
            self._logger.error(f"Unhandled error in connection {connection_id}: {e}", exc_info=True)
            try:
                await self.manager.send_json(connection_id, {
                    "type": MessageType.ERROR.value,
                    "code": "INTERNAL_ERROR",
                    "message": str(e),
                })
            except Exception:
                pass
        finally:
            await self._cleanup(connection_id)

    async def _handle_text(self, connection_id: str, engine: InferenceEngine, raw: str) -> None:
        """Route a single text message based on its type."""
        try:
            data = json.loads(raw)
        except json.JSONDecodeError:
            await self._send_error(connection_id, "INVALID_JSON", "Malformed JSON payload")
            return

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

        try:
            if msg_type == MessageType.AUDIO_CHUNK.value:
                await self._handle_audio_chunk(connection_id, engine, data)
            elif msg_type == MessageType.CONFIG.value:
                await self._handle_config(connection_id, engine, data)
            elif msg_type == MessageType.PING.value:
                await self._handle_ping(connection_id, data)
            elif msg_type == MessageType.RESET.value:
                await self._handle_reset(connection_id, engine)
            else:
                await self._send_error(connection_id, "UNKNOWN_MESSAGE_TYPE", f"Unknown type: {msg_type}")
        except Exception as e:
            self._logger.error(f"Error handling {msg_type} from {connection_id}: {e}")
            await self._send_error(connection_id, "HANDLER_ERROR", str(e))

    async def _handle_audio_chunk(
        self, connection_id: str, engine: InferenceEngine, data: dict
    ) -> None:
        """Process an audio chunk through the inference engine."""
        t_decode0 = time.perf_counter()
        samples = data.get("data")
        if samples is None or (isinstance(samples, (list, dict, str)) and len(samples) == 0):
            await self._send_error(connection_id, "INVALID_AUDIO_DATA", "Missing or empty audio data")
            return

        fmt = data.get("format", "int16")
        try:
            if isinstance(samples, str):
                raw_bytes = base64.b64decode(samples)
                dtype = np.float32 if fmt == "float32" else np.int16
                audio = np.frombuffer(raw_bytes, dtype=dtype)
            elif isinstance(samples, list):
                dtype = np.float32 if fmt == "float32" else np.int16
                audio = np.array(samples, dtype=dtype)
            elif isinstance(samples, dict):
                dtype = np.float32 if fmt == "float32" else np.int16
                audio = np.array(list(samples.values()), dtype=dtype)
            else:
                await self._send_error(connection_id, "INVALID_AUDIO_DATA", "Unsupported audio data format")
                return
        except Exception as e:
            await self._send_error(connection_id, "INVALID_AUDIO_DATA", f"Error parsing audio: {e}")
            return

        if len(audio) * 2 > self._max_chunk_size_bytes:
            await self._send_error(connection_id, "CHUNK_TOO_LARGE", f"Chunk exceeds {self._max_chunk_size_bytes} bytes")
            return

        decode_ms = (time.perf_counter() - t_decode0) * 1000.0

        sample_rate = data.get("sample_rate", 16000)
        timestamp = data.get("timestamp_ms")
        chunk_index = data.get("chunk_index")
        fmt = data.get("format", "int16")

        self.manager.touch(connection_id, n_chunks=1, n_bytes=len(audio.tobytes()))

        # Update chunk index tracking if provided by client
        if chunk_index is not None:
            setattr(engine, "_last_client_chunk_index", chunk_index)

        # Process through inference engine (streaming pipeline)
        result = engine.process_stream_chunk(audio, sample_rate)

        if result is not None:
            # Build detection result message (includes multi-detector fields)
            if result.probabilities is not None:
                probs_dict = {
                    settings.class_labels[i]: round(float(p), 4)
                    for i, p in enumerate(result.probabilities)
                }
            else:
                probs_dict = None

            latency_stages = dict(result.latency_stages or {})
            latency_stages["decode_ms"] = round(decode_ms, 2)

            detection_msg = {
                "type": MessageType.DETECTION_RESULT.value,
                "chunk_index": result.chunk_index,
                "timestamp_ms": result.timestamp_ms,
                "window_id": result.window_id,
                "status": result.status,
                "probabilities": probs_dict,
                "predicted_class": result.predicted_class,
                "predicted_label": result.predicted_label,
                "confidence": round(float(result.confidence), 4),
                "is_bonafide": bool(result.is_bonafide),
                "threat_score": (None if result.threat_score is None
                                 else round(float(result.threat_score), 2)),
                "authenticity_score": (None if result.authenticity_score is None
                                       else round(float(result.authenticity_score), 2)),
                "inference_ms": round(float(result.inference_ms), 2),
                "ema_synthetic_probability": (None if result.ema_synthetic_probability is None
                                              else round(float(result.ema_synthetic_probability), 4)),
                "tensor_stats": result.tensor_stats,
                "latency_stages": latency_stages,
                "vad": {
                    "is_speech": bool(result.vad_result["is_speech"]),
                    "energy": round(float(result.vad_result["energy"]), 6),
                    "snr": round(float(result.vad_result["snr"]), 2),
                    "confidence": round(float(result.vad_result["confidence"]), 4),
                },
                "audio_hash": result.audio_hash,
                "replay_score": round(float(result.replay_score), 4),
                "replay_high_confidence": bool(result.replay_high_confidence),
                "speaker_match": None if result.speaker_match is None else round(float(result.speaker_match), 4),
                "speaker_similarity": None if result.speaker_similarity is None else round(float(result.speaker_similarity), 4),
                "matched_speaker": result.matched_speaker,
                "fused_synthetic_probability": round(float(result.fused_synthetic_probability), 4),
                "calibrated_synthetic_probability": round(float(result.calibrated_synthetic_probability), 4),
                "calibrated_speaker_match": None if result.calibrated_speaker_match is None else round(float(result.calibrated_speaker_match), 4),
                "detector_consistency": round(float(result.detector_consistency), 4),
                "risk_level": result.risk_level,
                "risk_score": round(float(result.risk_score), 2),
                "risk_action": result.risk_action,
                "decision": result.decision,
                "behavioral_risk": (None if result.behavioral_risk is None
                                    else round(float(result.behavioral_risk), 2)),
                "context_risk": (None if result.context_risk is None
                                 else round(float(result.context_risk), 2)),
                "risk_escalated": bool(result.risk_escalated),
                "detections": result.detections,
            }
            await self.manager.send_json(connection_id, detection_msg)

            # Non-speech / low-quality statuses are informational only (no alerts)
            if result.status != "detected":
                return

            # Send threat alert if threat score is elevated
            if result.threat_score is not None and result.threat_score > 50:
                alert_msg = {
                    "type": MessageType.THREAT_ALERT.value,
                    "threat_score": round(result.threat_score, 2),
                    "authenticity_score": round(result.authenticity_score, 2),
                    "predicted_label": result.predicted_label,
                    "confidence": round(result.confidence, 4),
                    "timestamp_ms": result.timestamp_ms,
                    "chunk_index": result.chunk_index,
                    "audio_hash": result.audio_hash,
                }
                await self.manager.send_json(connection_id, alert_msg)

            # Send risk alert on HIGH / CRITICAL levels
            if result.risk_level in ("HIGH", "CRITICAL"):
                risk_msg = {
                    "type": MessageType.RISK_ALERT.value,
                    "risk_level": result.risk_level,
                    "risk_score": round(float(result.risk_score), 2),
                    "action": result.risk_action,
                    "decision": result.decision,
                    "reasons": list(result.risk_reasons),
                    "calibrated_synthetic_probability": round(float(result.calibrated_synthetic_probability), 4),
                    "speaker_match": None if result.speaker_match is None else round(float(result.speaker_match), 4),
                    "replay_probability": round(float(result.replay_score), 4),
                    "consecutive_level": 2 if result.risk_escalated else 1,
                    "escalated": bool(result.risk_escalated),
                    "timestamp_ms": result.timestamp_ms,
                    "chunk_index": result.chunk_index,
                    "audio_hash": result.audio_hash,
                }
                await self.manager.send_json(connection_id, risk_msg)
                self._logger.warning(
                    f"RISK_ALERT [{result.risk_level}] for {connection_id}: "
                    f"action={result.risk_action}, risk={result.risk_score:.1f}"
                )

                # Create real in-app alert (+ simulated SMS/email), pushed on WS.
                if self.alerts is not None:
                    alert_type = "fraud_team" if result.risk_level == "CRITICAL" else "supervisor"
                    severity = "critical" if result.risk_level == "CRITICAL" else "warning"
                    reasons = ", ".join(result.risk_reasons) or "no reasons"
                    self.alerts.send(
                        alert_type=alert_type,
                        severity=severity,
                        message=(f"risk={result.risk_score:.0f} level={result.risk_level} "
                                 f"decision={result.decision} reasons={reasons}"),
                        call_id=connection_id,
                        connection_id=connection_id,
                    )

            # Proactive verification workflow from risk decision.
            if result.decision in ("VERIFY", "BLOCK") and self.verification is not None:
                try:
                    flow = self.verification.begin(
                        connection_id=connection_id,
                        call_id=connection_id,
                        risk_level=result.risk_level,
                        reason=" | ".join(result.risk_reasons),
                    )
                    await self.manager.send_json(connection_id, {
                        "type": MessageType.VERIFY_REQUIRED.value,
                        "flow_id": flow.flow_id,
                        "call_id": flow.call_id,
                        "state": flow.state,
                        "risk_level": flow.risk_level,
                        "reason": flow.reason,
                        "attempts": flow.attempts,
                        "prototype": True,
                        "timestamp_ms": result.timestamp_ms,
                        "message": ("Verification requested (SIMULATED OTP workflow for the demo); "
                                    "submit a code via POST /api/verification/submit."),
                    })
                except Exception as e:  # verification must never break the audio loop
                    self._logger.warning(f"verification begin failed for {connection_id}: {e}")

            # Send kill-switch message if triggered
            if result.kill_switch_triggered:
                kill_msg = {
                    "type": MessageType.KILL_SWITCH.value,
                    "action": "TERMINATE_CALL",
                    "reason": result.kill_switch_reason,
                    "threat_score": round(result.threat_score, 2),
                    "consecutive_detections": 2,
                    "timestamp_ms": result.timestamp_ms,
                    "audio_hash": result.audio_hash,
                }
                await self.manager.send_json(connection_id, kill_msg)
                self._logger.warning(f"KILL_SWITCH triggered for {connection_id}: {result.kill_switch_reason}")

    async def _handle_config(self, connection_id: str, engine: InferenceEngine, data: dict) -> None:
        """Update configuration for this connection's engine."""
        threat_threshold = data.get("threat_threshold")
        consecutive_windows = data.get("consecutive_windows")
        enable_vad = data.get("enable_vad")
        enable_kill_switch = data.get("enable_kill_switch")

        if threat_threshold is not None:
            engine.kill_switch.threshold = float(threat_threshold)
        if consecutive_windows is not None:
            engine.kill_switch.consecutive_windows = int(consecutive_windows)
        if enable_vad is not None:
            engine.processor.vad.enabled = not enable_vad  # placeholder, adjust as needed
        if enable_kill_switch is not None:
            engine.kill_switch.enabled = bool(enable_kill_switch)

        self._logger.info(f"Config updated for {connection_id}: threshold={threat_threshold}")

        # Send updated stats
        await self._send_stats(connection_id, engine)

    async def _handle_ping(self, connection_id: str, data: dict) -> None:
        """Respond to client ping with server timestamp."""
        client_ts = data.get("timestamp_ms", 0)
        await self.manager.send_json(connection_id, {
            "type": MessageType.PONG.value,
            "timestamp_ms": client_ts,
            "server_time_ms": timestamp_ms(),
        })

    async def _handle_reset(self, connection_id: str, engine: InferenceEngine) -> None:
        """Reset engine state for a fresh stream."""
        engine.reset()
        await self.manager.send_json(connection_id, {
            "type": MessageType.STATS.value,
            "total_chunks": 0,
            "avg_inference_ms": 0.0,
            "kill_switch_triggers": 0,
            "errors": 0,
            "buffer_fill_ratio": 0.0,
            "uptime_ms": 0,
        })

    async def _send_ready(self, connection_id: str, engine: InferenceEngine) -> None:
        """Send ready message with model info."""
        model_obj = getattr(engine, "model", None)
        model_info = None
        if model_obj is not None and hasattr(model_obj, "get_model_info"):
            try:
                model_info = model_obj.get_model_info()
            except Exception:
                model_info = None
        if model_info is None:
            model_info = {
                "architecture": "VoiceGuardRawNet",
                "quantized": settings.quantized,
                "num_classes": settings.num_classes,
                "class_labels": settings.class_labels,
                "input_shape": [1, 1, 24000],
                "sample_rate": settings.sample_rate,
            }
        config = {
            "sample_rate": settings.sample_rate,
            "chunk_duration_ms": settings.chunk_duration_ms,
            "buffer_duration_s": settings.buffer_duration_s,
            "window_duration_s": settings.window_duration_s,
            "window_step_duration_s": settings.window_step_duration_s,
            "threat_threshold": engine.kill_switch.threshold,
            "consecutive_windows": engine.kill_switch.consecutive_windows,
            "quantized": settings.quantized,
            "speaker_enrolled": engine.speaker_verifier.enrolled,
            "storage_root": str(self.storage.root) if self.storage is not None else None,
        }
        await self.manager.send_json(connection_id, {
            "type": MessageType.READY.value,
            "model_info": model_info,
            "config": config,
        })

    async def _send_stats(self, connection_id: str, engine: InferenceEngine) -> None:
        """Send current engine stats."""
        stats = engine.get_stats()
        await self.manager.send_json(connection_id, {
            "type": MessageType.STATS.value,
            "total_chunks": stats.total_chunks,
            "avg_inference_ms": round(stats.avg_inference_ms, 2),
            "kill_switch_triggers": stats.kill_switch_triggers,
            "errors": stats.errors,
            "buffer_fill_ratio": round(engine.buffer.fill_ratio, 4),
            "uptime_ms": 0,
        })

    async def _send_error(self, connection_id: str, code: str, message: str) -> None:
        """Send error message to connection."""
        await self.manager.send_json(connection_id, {
            "type": MessageType.ERROR.value,
            "code": code,
            "message": message,
        })

    def _push_alert(self, connection_id: str, alert_data: dict) -> None:
        """Sync callback from AlertService -> async WS push (thread-safe)."""
        loop = self._loop
        if loop is None or loop.is_closed():
            return
        msg = dict(alert_data)
        msg["type"] = MessageType.ALERT.value
        try:
            asyncio.run_coroutine_threadsafe(
                self.manager.send_json(connection_id, msg), loop
            )
        except Exception:  # pragma: no cover - push must never break alerts
            pass

    def _get_quantized_layers(self, engine: InferenceEngine) -> list:
        """Safely get list of quantized layer names."""
        try:
            model = engine.model.model if hasattr(engine.model, "model") else None
            if model is None:
                return []
            return [
                name for name, module in model.named_modules()
                if hasattr(module, 'weight') and hasattr(module.weight, 'dtype') and module.weight.dtype == torch.qint8
            ]
        except Exception:
            return []

    async def _cleanup(self, connection_id: str) -> None:
        """Cleanup connection resources."""
        if connection_id in self._engines:
            self._engines[connection_id].reset()
            del self._engines[connection_id]
        if self.verification is not None:
            try:
                self.verification.close(connection_id)
            except Exception:
                pass
        await self.manager.disconnect(connection_id)