"""
BackendConnector - WebSocket connection for real-time backend consultation.

Provides bidirectional communication with the Aigie backend for:
- Real-time interception decisions
- Leveraging historical data and patterns
- Receiving proactive fixes and recommendations

Built on BaseWebSocketClient for shared connection management.
"""

import asyncio
import logging
import time
import uuid
from collections.abc import Awaitable
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional

from .base_ws import BaseWebSocketClient

if TYPE_CHECKING:
    from ..interceptor.protocols import (
        InterceptionContext,
        PostCallResult,
        PreCallResult,
    )

logger = logging.getLogger("aigie.realtime")


@dataclass
class ConsultationRequest:
    """Request for backend consultation."""

    request_id: str
    """Unique request ID for correlation."""

    consultation_type: str
    """Type: 'pre_call', 'post_call', 'drift_check', 'fix_request'."""

    context: Dict[str, Any]
    """Serialized interception context."""

    urgency: str = "normal"
    """Urgency level: 'low', 'normal', 'high', 'critical'."""

    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))


@dataclass
class ConsultationResponse:
    """Response from backend consultation."""

    request_id: str
    """Correlated request ID."""

    decision: str
    """Decision: 'allow', 'block', 'modify', 'retry'."""

    reason: str | None = None
    """Reason for the decision."""

    fixes: List[Dict[str, Any]] = field(default_factory=list)
    """List of fix actions to apply."""

    modified_request: Dict[str, Any] | None = None
    """Modified request parameters."""

    modified_response: Dict[str, Any] | None = None
    """Modified response content."""

    confidence: float = 0.0
    """Confidence score for this decision."""

    latency_ms: float = 0.0
    """Backend processing latency."""

    metadata: Dict[str, Any] = field(default_factory=dict)
    """Additional metadata from backend."""


class BackendConnector(BaseWebSocketClient):
    """
    WebSocket connector for real-time backend communication.

    Extends BaseWebSocketClient with consultation-specific logic:
    - Pre/post call consultation with request/response correlation
    - Push-based fix and alert callbacks
    - Event type subscriptions
    """

    def __init__(
        self,
        api_url: str,
        api_key: str,
        auto_reconnect: bool = True,
        reconnect_interval: float = 5.0,
        max_reconnect_attempts: int = 10,
        ping_interval: float = 30.0,
        request_timeout: float = 0.5,  # 500ms default
    ):
        super().__init__(
            api_url=api_url,
            api_key=api_key,
            ws_path="/v1/gateway/ws",
            auto_reconnect=auto_reconnect,
            reconnect_interval=reconnect_interval,
            max_reconnect_attempts=max_reconnect_attempts,
            ping_interval=ping_interval,
        )
        self._request_timeout = request_timeout

        # Callbacks for push events
        self._fix_callbacks: List[Callable[[Dict[str, Any]], Awaitable[None]]] = []
        self._alert_callbacks: List[Callable[[Dict[str, Any]], Awaitable[None]]] = []
        self._pattern_update_callbacks: List[Callable[[Dict[str, Any]], Awaitable[None]]] = []

        # Consultation-specific stats
        self._stats["consultations"] = 0
        self._stats["successful_responses"] = 0
        self._stats["total_latency_ms"] = 0.0

    async def _on_connected(self) -> None:
        """Read welcome message and subscribe to relevant event types."""
        # Read welcome message from gateway
        if self._websocket:
            try:
                import json as _json

                welcome_msg = await asyncio.wait_for(self._websocket.recv(), timeout=5.0)
                welcome_data = _json.loads(welcome_msg)
                logger.info(
                    f"Connected to gateway (connection_id: {welcome_data.get('connection_id')})"
                )
            except Exception as e:
                logger.debug(f"Failed to read welcome message: {e}")

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

        if msg_type == "consultation_response":
            request_id = data.get("request_id")
            if request_id:
                self._resolve_request(request_id, data)

        elif msg_type in ("fix_push", "intervention"):
            payload = data.get("payload", data.get("data", data))
            for callback in self._fix_callbacks:
                asyncio.create_task(callback(payload))

        elif msg_type == "alert":
            for callback in self._alert_callbacks:
                asyncio.create_task(callback(data))

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

        elif msg_type in ("pong", "ping"):
            pass  # Keepalive response

    async def consult_pre_call(
        self,
        ctx: "InterceptionContext",
        timeout: float | None = None,
    ) -> Optional["PreCallResult"]:
        """
        Consult backend for pre-call decision.

        Args:
            ctx: Interception context
            timeout: Request timeout (default: self._request_timeout)

        Returns:
            PreCallResult if successful, None on timeout/error
        """
        from ..interceptor.protocols import InterceptionDecision, PreCallResult

        if not self.is_connected:
            return None

        request = ConsultationRequest(
            request_id=str(uuid.uuid4()),
            consultation_type="pre_call",
            context=self._serialize_context(ctx),
        )

        response = await self._send_consultation(request, timeout or self._request_timeout)

        if response is None:
            return None

        # Convert response to PreCallResult
        decision_map = {
            "allow": InterceptionDecision.ALLOW,
            "block": InterceptionDecision.BLOCK,
            "modify": InterceptionDecision.MODIFY,
        }

        decision = decision_map.get(response.decision, InterceptionDecision.ALLOW)

        if decision == InterceptionDecision.BLOCK:
            return PreCallResult.block(
                reason=response.reason or "Blocked by backend",
                hook_name="backend",
                latency_ms=response.latency_ms,
            )

        if decision == InterceptionDecision.MODIFY and response.modified_request:
            return PreCallResult.modify(
                messages=response.modified_request.get("messages"),
                kwargs=response.modified_request.get("kwargs"),
                reason=response.reason,
                hook_name="backend",
                latency_ms=response.latency_ms,
            )

        return PreCallResult.allow(hook_name="backend", latency_ms=response.latency_ms)

    async def consult_post_call(
        self,
        ctx: "InterceptionContext",
        timeout: float | None = None,
    ) -> Optional["PostCallResult"]:
        """
        Consult backend for post-call decision and fixes.

        Args:
            ctx: Interception context with response
            timeout: Request timeout

        Returns:
            PostCallResult if successful, None on timeout/error
        """
        from ..interceptor.protocols import (
            FixAction,
            FixActionType,
            PostCallResult,
        )

        if not self.is_connected:
            return None

        request = ConsultationRequest(
            request_id=str(uuid.uuid4()),
            consultation_type="post_call",
            context=self._serialize_context(ctx),
            urgency="high" if ctx.error else "normal",
        )

        response = await self._send_consultation(request, timeout or self._request_timeout)

        if response is None:
            return None

        # Convert fixes
        fixes = []
        for fix_data in response.fixes:
            action_type = FixActionType.MODIFY_RESPONSE
            if fix_data.get("action_type") == "retry":
                action_type = FixActionType.RETRY
            elif fix_data.get("action_type") == "fallback":
                action_type = FixActionType.FALLBACK

            fixes.append(
                FixAction(
                    action_type=action_type,
                    parameters=fix_data.get("parameters", {}),
                    confidence=fix_data.get("confidence", 0.8),
                    source="backend",
                    reason=fix_data.get("reason"),
                )
            )

        # Check if retry is needed
        if response.decision == "retry":
            return PostCallResult.retry(
                reason=response.reason or "Backend requested retry",
                retry_kwargs=response.modified_request,
                hook_name="backend",
                latency_ms=response.latency_ms,
            )

        # Return modification result
        if response.modified_response or fixes:
            return PostCallResult.modify(
                response=response.modified_response,
                content=response.modified_response.get("content")
                if response.modified_response
                else None,
                reason=response.reason,
                fixes=fixes,
                hook_name="backend",
                latency_ms=response.latency_ms,
            )

        return PostCallResult.allow(hook_name="backend", latency_ms=response.latency_ms)

    def on_fix_push(self, callback: Callable[[Dict[str, Any]], Awaitable[None]]) -> None:
        """Register callback for backend-pushed fixes."""
        self._fix_callbacks.append(callback)

    def on_alert(self, callback: Callable[[Dict[str, Any]], Awaitable[None]]) -> None:
        """Register callback for backend alerts."""
        self._alert_callbacks.append(callback)

    def on_pattern_update(self, callback: Callable[[Dict[str, Any]], Awaitable[None]]) -> None:
        """Register callback for pattern cache refresh notifications."""
        self._pattern_update_callbacks.append(callback)

    async def _send_consultation(
        self,
        request: ConsultationRequest,
        timeout: float,
    ) -> ConsultationResponse | None:
        """Send consultation request and wait for response."""
        self._stats["consultations"] += 1
        start_time = time.perf_counter()

        message = {
            "type": "consultation",
            "timestamp": request.timestamp.isoformat(),
            "payload": {
                "request_id": request.request_id,
                "consultation_type": request.consultation_type,
                "context": request.context,
                "urgency": request.urgency,
            },
        }

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

        if response_data is None:
            return None

        latency = (time.perf_counter() - start_time) * 1000
        self._stats["successful_responses"] += 1
        self._stats["total_latency_ms"] += latency

        return ConsultationResponse(
            request_id=request.request_id,
            decision=response_data.get("decision", "allow"),
            reason=response_data.get("reason"),
            fixes=response_data.get("fixes", []),
            modified_request=response_data.get("modified_request"),
            modified_response=response_data.get("modified_response"),
            confidence=response_data.get("confidence", 0.0),
            latency_ms=latency,
            metadata=response_data.get("metadata", {}),
        )

    def _serialize_context(self, ctx: "InterceptionContext") -> Dict[str, Any]:
        """Serialize interception context for transmission."""
        return {
            "provider": ctx.provider,
            "model": ctx.model,
            "messages": ctx.messages[-5:] if ctx.messages else [],  # Last 5 messages
            "trace_id": ctx.trace_id,
            "span_id": ctx.span_id,
            "estimated_cost": ctx.estimated_cost,
            "accumulated_cost": ctx.accumulated_cost,
            "drift_score": ctx.drift_score,
            "context_hash": ctx.context_hash,
            "error": str(ctx.error) if ctx.error else None,
            "error_type": ctx.error_type,
            "user_id": ctx.user_id,
            "session_id": ctx.session_id,
            "response_content": ctx.response_content[:1000] if ctx.response_content else None,
            "actual_cost": ctx.actual_cost,
            "response_time_ms": ctx.response_time_ms,
        }

    def get_stats(self) -> Dict[str, Any]:
        """Get connector statistics."""
        stats = super().get_stats()
        stats["avg_latency_ms"] = self._stats["total_latency_ms"] / max(
            self._stats["successful_responses"], 1
        )
        return stats
