"""
WebSocket Gateway Client - Real-Time Validation with Persistent Connections.

This module provides a WebSocket client for real-time validation with
<100ms latency target. It maintains persistent connections to the backend
gateway for low-latency pre-execution validation and receives push-based
interventions.

Key Features:
- Persistent WebSocket connection management
- Pre-execution validation hooks
- Real-time intervention receiving
- Reconnection with exponential backoff (2s interval, max 5 attempts)
- Validation result caching (max 100 results)
- Intervention caching (max 50 interventions)
- Keepalive ping every 30s

Built on BaseWebSocketClient for shared connection management.
"""

import asyncio
import hashlib
import json
import logging
import time
from collections import OrderedDict
from collections.abc import Awaitable
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Callable, Dict, List

from ..realtime.base_ws import BaseWebSocketClient, ConnectionState

logger = logging.getLogger("aigie.gateway.websocket")


# Re-export ConnectionState as WebSocketConnectionState for backward compat
WebSocketConnectionState = ConnectionState


@dataclass
class ValidationResult:
    """Result of a tool call validation."""

    request_id: str
    tool_name: str
    decision: str  # "allow", "block", "modify", "delay"
    confidence: float = 1.0
    reason: str | None = None
    modified_args: Dict[str, Any] | None = None
    delay_ms: int | None = None
    signals: List[Dict[str, Any]] = field(default_factory=list)
    latency_ms: float = 0.0
    cached: bool = False
    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))

    def to_dict(self) -> Dict[str, Any]:
        return {
            "request_id": self.request_id,
            "tool_name": self.tool_name,
            "decision": self.decision,
            "confidence": self.confidence,
            "reason": self.reason,
            "modified_args": self.modified_args,
            "delay_ms": self.delay_ms,
            "signals": self.signals,
            "latency_ms": self.latency_ms,
            "cached": self.cached,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ValidationResult":
        return cls(
            request_id=data.get("request_id", ""),
            tool_name=data.get("tool_name", ""),
            decision=data.get("decision", "allow"),
            confidence=data.get("confidence", 1.0),
            reason=data.get("reason"),
            modified_args=data.get("modified_args"),
            delay_ms=data.get("delay_ms"),
            signals=data.get("signals", []),
            latency_ms=data.get("latency_ms", 0.0),
            cached=data.get("cached", False),
        )


@dataclass
class Intervention:
    """An intervention signal pushed from the backend."""

    intervention_id: str
    trace_id: str
    intervention_type: str  # "block", "modify", "warn", "escalate"
    reason: str
    action: Dict[str, Any] | None = None
    severity: str = "medium"  # "low", "medium", "high", "critical"
    metadata: Dict[str, Any] = field(default_factory=dict)
    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    acknowledged: bool = False

    def to_dict(self) -> Dict[str, Any]:
        return {
            "intervention_id": self.intervention_id,
            "trace_id": self.trace_id,
            "intervention_type": self.intervention_type,
            "reason": self.reason,
            "action": self.action,
            "severity": self.severity,
            "metadata": self.metadata,
            "timestamp": self.timestamp.isoformat(),
            "acknowledged": self.acknowledged,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Intervention":
        timestamp = data.get("timestamp")
        if isinstance(timestamp, str):
            try:
                timestamp = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
            except ValueError:
                timestamp = datetime.now(timezone.utc)
        elif not isinstance(timestamp, datetime):
            timestamp = datetime.now(timezone.utc)

        return cls(
            intervention_id=data.get("intervention_id", ""),
            trace_id=data.get("trace_id", ""),
            intervention_type=data.get("intervention_type", "warn"),
            reason=data.get("reason", ""),
            action=data.get("action"),
            severity=data.get("severity", "medium"),
            metadata=data.get("metadata", {}),
            timestamp=timestamp,
            acknowledged=data.get("acknowledged", False),
        )


@dataclass
class WebSocketMetrics:
    """Metrics for WebSocket gateway client operations."""

    connections_attempted: int = 0
    connections_successful: int = 0
    connections_failed: int = 0
    reconnections: int = 0
    validations_sent: int = 0
    validations_received: int = 0
    validations_cached: int = 0
    validations_timed_out: int = 0
    validations_blocked: int = 0
    interventions_received: int = 0
    interventions_acknowledged: int = 0
    messages_sent: int = 0
    messages_received: int = 0
    total_latency_ms: float = 0.0
    _latency_samples: List[float] = field(default_factory=list)

    def record_validation(self, latency_ms: float, decision: str, cached: bool = False):
        """Record a validation request."""
        self.validations_received += 1
        self.total_latency_ms += latency_ms
        self._latency_samples.append(latency_ms)

        # Keep only last 1000 samples
        if len(self._latency_samples) > 1000:
            self._latency_samples.pop(0)

        if cached:
            self.validations_cached += 1
        if decision == "block":
            self.validations_blocked += 1

    @property
    def avg_latency_ms(self) -> float:
        if not self._latency_samples:
            return 0.0
        return sum(self._latency_samples) / len(self._latency_samples)

    @property
    def p50_latency_ms(self) -> float:
        if not self._latency_samples:
            return 0.0
        sorted_samples = sorted(self._latency_samples)
        idx = len(sorted_samples) // 2
        return sorted_samples[idx]

    @property
    def p95_latency_ms(self) -> float:
        if not self._latency_samples:
            return 0.0
        sorted_samples = sorted(self._latency_samples)
        idx = int(len(sorted_samples) * 0.95)
        return sorted_samples[min(idx, len(sorted_samples) - 1)]

    @property
    def p99_latency_ms(self) -> float:
        if not self._latency_samples:
            return 0.0
        sorted_samples = sorted(self._latency_samples)
        idx = int(len(sorted_samples) * 0.99)
        return sorted_samples[min(idx, len(sorted_samples) - 1)]

    def to_dict(self) -> Dict[str, Any]:
        return {
            "connections_attempted": self.connections_attempted,
            "connections_successful": self.connections_successful,
            "connections_failed": self.connections_failed,
            "reconnections": self.reconnections,
            "validations_sent": self.validations_sent,
            "validations_received": self.validations_received,
            "validations_cached": self.validations_cached,
            "validations_timed_out": self.validations_timed_out,
            "validations_blocked": self.validations_blocked,
            "interventions_received": self.interventions_received,
            "interventions_acknowledged": self.interventions_acknowledged,
            "messages_sent": self.messages_sent,
            "messages_received": self.messages_received,
            "avg_latency_ms": self.avg_latency_ms,
            "p50_latency_ms": self.p50_latency_ms,
            "p95_latency_ms": self.p95_latency_ms,
            "p99_latency_ms": self.p99_latency_ms,
        }


class LRUCache:
    """Simple LRU cache for validation results."""

    def __init__(self, max_size: int = 100):
        self._max_size = max_size
        self._cache: OrderedDict[str, Any] = OrderedDict()

    def get(self, key: str) -> Any | None:
        if key in self._cache:
            # Move to end (most recently used)
            self._cache.move_to_end(key)
            return self._cache[key]
        return None

    def put(self, key: str, value: Any) -> None:
        if key in self._cache:
            self._cache.move_to_end(key)
        else:
            if len(self._cache) >= self._max_size:
                # Remove oldest item
                self._cache.popitem(last=False)
        self._cache[key] = value

    def clear(self) -> None:
        self._cache.clear()

    def __len__(self) -> int:
        return len(self._cache)


class GatewayWebSocketClient(BaseWebSocketClient):
    """
    Real-time WebSocket validation client with <100ms latency target.

    Extends BaseWebSocketClient with gateway-specific features:
    - Pre-execution validation for tool calls
    - Validation result caching (LRU)
    - Real-time intervention receiving via push
    - Intervention caching and acknowledgment
    - Trace subscriptions
    """

    # Configuration defaults
    DEFAULT_RECONNECT_INTERVAL_SEC = 2.0
    DEFAULT_MAX_RECONNECT_ATTEMPTS = 5
    DEFAULT_VALIDATION_TIMEOUT_MS = 500
    DEFAULT_CONNECT_TIMEOUT_SEC = 10
    DEFAULT_PING_INTERVAL_SEC = 30
    DEFAULT_VALIDATION_CACHE_SIZE = 100
    DEFAULT_INTERVENTION_CACHE_SIZE = 50

    def __init__(
        self,
        api_url: str,
        api_key: str,
        *,
        validation_timeout_ms: float = DEFAULT_VALIDATION_TIMEOUT_MS,
        connect_timeout_sec: float = DEFAULT_CONNECT_TIMEOUT_SEC,
        ping_interval_sec: float = DEFAULT_PING_INTERVAL_SEC,
        reconnect_interval_sec: float = DEFAULT_RECONNECT_INTERVAL_SEC,
        max_reconnect_attempts: int = DEFAULT_MAX_RECONNECT_ATTEMPTS,
        validation_cache_size: int = DEFAULT_VALIDATION_CACHE_SIZE,
        intervention_cache_size: int = DEFAULT_INTERVENTION_CACHE_SIZE,
        enable_validation_cache: bool = True,
        enable_interventions: bool = True,
        fallback_on_error: str = "allow",  # "allow" or "block"
        agent_name: str | None = None,
        framework: str | None = None,
    ):
        super().__init__(
            api_url=api_url,
            api_key=api_key,
            ws_path="/v1/gateway/ws",
            auto_reconnect=True,
            reconnect_interval=reconnect_interval_sec,
            max_reconnect_attempts=max_reconnect_attempts,
            ping_interval=ping_interval_sec,
            connect_timeout=connect_timeout_sec,
        )

        self._validation_timeout_ms = validation_timeout_ms
        self._enable_validation_cache = enable_validation_cache
        self._enable_interventions = enable_interventions
        self._fallback_decision = fallback_on_error

        # Agent identification
        self._agent_name = agent_name
        self._framework = framework

        # Connection metadata
        self._connection_id: str | None = None

        # Request counter
        self._request_counter = 0

        # Caches
        self._validation_cache = LRUCache(max_size=validation_cache_size)
        self._intervention_cache = LRUCache(max_size=intervention_cache_size)

        # Intervention callbacks
        self._intervention_callbacks: List[Callable[[Intervention], Awaitable[None]]] = []

        # Pattern update callbacks (triggered when backend promotes new patterns)
        self._pattern_update_callbacks: List[Callable] = []

        # Trace subscriptions
        self._subscribed_traces: set = set()

        # Gateway-specific metrics (richer than base stats)
        self._metrics = WebSocketMetrics()

    # ─── Override URL building for gateway endpoint ──────────────────

    def _get_ws_url(self) -> str:
        """Build gateway WebSocket URL with non-sensitive query params."""
        import os
        from urllib.parse import urlencode

        from .. import __version__

        params = {
            "client_type": "sdk",
            "client_version": __version__,
        }
        if self._agent_name:
            params["agent_name"] = self._agent_name
        if self._framework:
            params["framework"] = self._framework
        # Pod identification for Kubernetes deployments
        hostname = os.getenv("HOSTNAME")
        if hostname:
            params["pod_name"] = hostname
        # Read namespace from K8s service account (always available in-cluster)
        namespace = os.getenv("POD_NAMESPACE")
        if not namespace:
            try:
                with open("/var/run/secrets/kubernetes.io/serviceaccount/namespace") as f:
                    namespace = f.read().strip()
            except OSError:
                pass
        if namespace:
            params["pod_namespace"] = namespace
        return f"{self._ws_base_url}?{urlencode(params)}"

    def _get_connect_headers(self) -> Dict[str, str]:
        """Send API key and client metadata as WebSocket handshake headers."""
        return super()._get_connect_headers()

    # ─── Properties ──────────────────────────────────────────────────

    @property
    def connection_id(self) -> str | None:
        return self._connection_id

    @property
    def metrics(self) -> WebSocketMetrics:
        return self._metrics

    @property
    def state(self) -> WebSocketConnectionState:
        return self._state

    # ─── Post-connect: read welcome message ──────────────────────────

    async def _on_connected(self) -> None:
        """Read welcome message and resubscribe to traces."""
        if self._websocket:
            try:
                welcome_msg = await asyncio.wait_for(self._websocket.recv(), timeout=5.0)
                welcome_data = json.loads(welcome_msg)
                self._connection_id = welcome_data.get("connection_id")
                self._metrics.messages_received += 1
                logger.info(f"Connected to gateway (connection_id: {self._connection_id})")
            except Exception as e:
                logger.warning(f"Failed to read welcome message: {e}")

        # Resubscribe to previously subscribed traces
        for trace_id in self._subscribed_traces:
            await self._send_subscribe(trace_id)

    # ─── Message dispatch ────────────────────────────────────────────

    async def _on_message(self, data: Dict[str, Any]) -> None:
        """Handle incoming gateway messages."""
        msg_type = data.get("type")

        if msg_type == "validation_result":
            payload = data.get("payload", {})
            request_id = payload.get("request_id")
            if request_id:
                self._resolve_request(request_id, payload)

        elif msg_type == "intervention":
            if self._enable_interventions:
                self._metrics.interventions_received += 1
                payload = data.get("payload", {})
                try:
                    intervention = Intervention.from_dict(payload)
                    self._intervention_cache.put(intervention.intervention_id, intervention)
                    for callback in self._intervention_callbacks:
                        asyncio.create_task(self._safe_callback(callback, intervention))
                except Exception as e:
                    logger.warning(f"Failed to parse intervention: {e}")

        elif msg_type == "pong":
            pass

        elif msg_type == "ping":
            await self._send(
                {
                    "type": "pong",
                    "timestamp": datetime.now(timezone.utc).isoformat(),
                }
            )

        elif msg_type == "error":
            logger.warning(f"Gateway error: {data.get('error')}")

        elif msg_type == "mode_change":
            payload = data.get("payload", {})
            logger.info(f"Mode change notification: {payload}")

        elif msg_type == "pattern_updated":
            # Backend promoted new patterns — notify pattern cache
            for callback in self._pattern_update_callbacks:
                asyncio.create_task(self._safe_callback(callback, data.get("payload", data)))

    # ─── Validation ──────────────────────────────────────────────────

    async def validate_tool_call(
        self,
        tool_name: str,
        args: Dict[str, Any],
        *,
        trace_id: str | None = None,
        span_id: str | None = None,
        timeout_ms: float | None = None,
        use_cache: bool = True,
    ) -> ValidationResult:
        """
        Validate a tool call before execution. Target latency: <100ms.

        Args:
            tool_name: Name of the tool being called
            args: Tool call arguments
            trace_id: Optional trace ID for context
            span_id: Optional span ID for context
            timeout_ms: Request timeout in milliseconds
            use_cache: Whether to use cached results

        Returns:
            ValidationResult with decision and any modifications
        """
        start_time = time.perf_counter()
        timeout = (timeout_ms or self._validation_timeout_ms) / 1000.0

        # Generate request ID
        self._request_counter += 1
        request_id = f"val_{self._request_counter}_{int(time.time() * 1000)}"

        # Check cache first
        if use_cache and self._enable_validation_cache:
            cache_key = self._make_cache_key(tool_name, args)
            cached = self._validation_cache.get(cache_key)
            if cached:
                latency_ms = (time.perf_counter() - start_time) * 1000
                cached.latency_ms = latency_ms
                cached.cached = True
                self._metrics.record_validation(latency_ms, cached.decision, cached=True)
                return cached

        # Check connection
        if not self.is_connected:
            return self._make_fallback_result(
                request_id=request_id,
                tool_name=tool_name,
                reason="Not connected to gateway",
                start_time=start_time,
            )

        self._metrics.validations_sent += 1

        message = {
            "type": "validate",
            "request_id": request_id,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "payload": {
                "tool_name": tool_name,
                "arguments": args,
                "trace_id": trace_id,
                "span_id": span_id,
            },
        }

        response_data = await self._send_and_wait(request_id, message, timeout)

        if response_data is None:
            self._metrics.validations_timed_out += 1
            return self._make_fallback_result(
                request_id=request_id,
                tool_name=tool_name,
                reason=f"Validation timeout ({timeout * 1000:.0f}ms)",
                start_time=start_time,
            )

        latency_ms = (time.perf_counter() - start_time) * 1000
        result = ValidationResult(
            request_id=request_id,
            tool_name=tool_name,
            decision=response_data.get("decision", "allow"),
            confidence=response_data.get("confidence", 1.0),
            reason=response_data.get("reason"),
            modified_args=response_data.get("modified_args"),
            delay_ms=response_data.get("delay_ms"),
            signals=response_data.get("signals", []),
            latency_ms=latency_ms,
            cached=False,
        )

        # Cache the result
        if self._enable_validation_cache:
            cache_key = self._make_cache_key(tool_name, args)
            self._validation_cache.put(cache_key, result)

        self._metrics.record_validation(latency_ms, result.decision)
        return result

    # ─── Interventions ───────────────────────────────────────────────

    async def receive_intervention(self) -> Intervention | None:
        """Return the oldest unacknowledged intervention, or None."""
        for key in list(self._intervention_cache._cache.keys()):
            intervention = self._intervention_cache.get(key)
            if intervention and not intervention.acknowledged:
                return intervention
        return None

    def on_intervention(self, callback: Callable[[Intervention], Awaitable[None]]) -> None:
        """Register callback for intervention signals."""
        self._intervention_callbacks.append(callback)

    def on_pattern_update(self, callback: Callable) -> None:
        """Register callback for pattern cache refresh notifications."""
        self._pattern_update_callbacks.append(callback)

    def remove_intervention_callback(
        self, callback: Callable[[Intervention], Awaitable[None]]
    ) -> bool:
        """Remove an intervention callback."""
        try:
            self._intervention_callbacks.remove(callback)
            return True
        except ValueError:
            return False

    async def acknowledge_intervention(self, intervention_id: str) -> bool:
        """Acknowledge an intervention."""
        intervention = self._intervention_cache.get(intervention_id)
        if intervention:
            intervention.acknowledged = True
            self._metrics.interventions_acknowledged += 1

            if self.is_connected:
                await self._send(
                    {
                        "type": "acknowledge_intervention",
                        "timestamp": datetime.now(timezone.utc).isoformat(),
                        "payload": {"intervention_id": intervention_id},
                    }
                )

            return True
        return False

    # ─── Trace subscriptions ─────────────────────────────────────────

    def subscribe_to_trace(self, trace_id: str) -> None:
        """Subscribe to interventions for a specific trace."""
        self._subscribed_traces.add(trace_id)
        if self.is_connected:
            try:
                loop = asyncio.get_running_loop()
                loop.create_task(self._send_subscribe(trace_id))
            except RuntimeError:
                pass  # No running event loop; subscription recorded, will be sent on reconnect

    def unsubscribe_from_trace(self, trace_id: str) -> None:
        """Unsubscribe from trace interventions."""
        self._subscribed_traces.discard(trace_id)
        if self.is_connected:
            try:
                loop = asyncio.get_running_loop()
                loop.create_task(self._send_unsubscribe(trace_id))
            except RuntimeError:
                pass  # No running event loop; will be cleaned up on reconnect

    async def subscribe_trace(self, trace_id: str) -> bool:
        """Subscribe to trace-specific fixes via WebSocket."""
        try:
            self._subscribed_traces.add(trace_id)
            await self._send(
                {
                    "type": "subscribe",
                    "payload": {"trace_id": trace_id},
                }
            )
            return True
        except Exception as e:
            logger.warning(f"Failed to subscribe to trace {trace_id}: {e}")
            return False

    async def unsubscribe_trace(self, trace_id: str) -> bool:
        """Unsubscribe from trace-specific fixes."""
        try:
            self._subscribed_traces.discard(trace_id)
            await self._send(
                {
                    "type": "unsubscribe",
                    "payload": {"trace_id": trace_id},
                }
            )
            return True
        except Exception as e:
            logger.warning(f"Failed to unsubscribe from trace {trace_id}: {e}")
            return False

    async def _send_subscribe(self, trace_id: str) -> None:
        await self._send(
            {
                "type": "subscribe",
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "payload": {"trace_id": trace_id},
            }
        )

    async def _send_unsubscribe(self, trace_id: str) -> None:
        await self._send(
            {
                "type": "unsubscribe",
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "payload": {"trace_id": trace_id},
            }
        )

    # ─── Helpers ─────────────────────────────────────────────────────

    def _make_cache_key(self, tool_name: str, args: Dict[str, Any]) -> str:
        """Create a cache key for validation results."""
        content = json.dumps({"tool": tool_name, "args": args}, sort_keys=True)
        return hashlib.md5(content.encode()).hexdigest()

    def _make_fallback_result(
        self,
        request_id: str,
        tool_name: str,
        reason: str,
        start_time: float,
    ) -> ValidationResult:
        """Create a fallback validation result."""
        latency_ms = (time.perf_counter() - start_time) * 1000
        return ValidationResult(
            request_id=request_id,
            tool_name=tool_name,
            decision=self._fallback_decision,
            confidence=0.0,
            reason=reason,
            latency_ms=latency_ms,
            cached=False,
        )

    async def _safe_callback(
        self,
        callback: Callable[[Intervention], Awaitable[None]],
        intervention: Intervention,
    ) -> None:
        """Safely call a callback function."""
        try:
            await callback(intervention)
        except Exception as e:
            logger.warning(f"Intervention callback error: {e}")

    def clear_validation_cache(self) -> None:
        """Clear the validation result cache."""
        self._validation_cache.clear()

    def clear_intervention_cache(self) -> None:
        """Clear the intervention cache."""
        self._intervention_cache.clear()

    # ─── Stats ───────────────────────────────────────────────────────

    def get_stats(self) -> Dict[str, Any]:
        """Get client statistics."""
        # Sync base stats into gateway metrics
        self._metrics.connections_attempted = self._stats["connections_attempted"]
        self._metrics.connections_successful = self._stats["connections_successful"]
        self._metrics.connections_failed = self._stats["connections_failed"]
        self._metrics.reconnections = self._stats["reconnections"]
        self._metrics.messages_sent = self._stats["messages_sent"]
        self._metrics.messages_received = self._stats["messages_received"]

        return {
            **self._metrics.to_dict(),
            "state": self._state.value,
            "connection_id": self._connection_id,
            "pending_requests": len(self._pending_requests),
            "validation_cache_size": len(self._validation_cache),
            "intervention_cache_size": len(self._intervention_cache),
            "subscribed_traces": len(self._subscribed_traces),
            "reconnect_attempts": self._reconnect_attempts,
            "last_connect_time": self._last_connect_time.isoformat()
            if self._last_connect_time
            else None,
        }

    def reset_stats(self) -> None:
        """Reset statistics."""
        self._metrics = WebSocketMetrics()

    async def connect(self, url: str | None = None) -> bool:
        """
        Establish WebSocket connection to gateway.

        Args:
            url: Optional override URL for connection (backward compat)
        """
        if url:
            # Temporarily override the URL for this connection
            original = self._ws_base_url
            ws_url = url.replace("http://", "ws://").replace("https://", "wss://")
            self._ws_base_url = ws_url
            result = await super().connect()
            if not result:
                self._ws_base_url = original
            return result
        return await super().connect()
