"""
InsAIts SDK - Edge/Hybrid Swarm Router
=======================================
Dynamic routing between local, cloud, and edge devices for embeddings
and AI processing with automatic failover based on latency and privacy.

Key Features:
- Node scoring (latency, privacy, availability)
- Automatic cloud-to-local failover if latency > threshold
- Privacy modes: strict (local only), balanced, performance
- Health monitoring for each node type
- Retry logic with exponential backoff

Usage:
    from insa_its.edge_router import EmbeddingRouter, PrivacyMode

    router = EmbeddingRouter(privacy_mode=PrivacyMode.BALANCED)
    embedding = router.get_embedding("text to embed")
"""

import time
import hashlib
import logging
import threading
from enum import Enum
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Callable, Any
from datetime import datetime, timedelta
import numpy as np

logger = logging.getLogger(__name__)


class PrivacyMode(Enum):
    """Privacy levels for routing decisions."""
    STRICT = "strict"       # Local only - no data leaves machine
    BALANCED = "balanced"   # Prefer local, cloud as fallback
    PERFORMANCE = "performance"  # Prefer cloud for speed, local as fallback


class NodeType(Enum):
    """Types of processing nodes."""
    LOCAL = "local"         # sentence-transformers on local machine
    CLOUD = "cloud"         # Cloud API (InsAIts backend)
    EDGE = "edge"           # Edge device (future: Raspberry Pi, etc.)
    SYNTHETIC = "synthetic" # Fallback synthetic embeddings


@dataclass
class NodeHealth:
    """Health status for a processing node."""
    node_type: NodeType
    is_available: bool = True
    last_check: float = field(default_factory=time.time)
    avg_latency_ms: float = 0.0
    success_count: int = 0
    failure_count: int = 0
    consecutive_failures: int = 0
    last_error: Optional[str] = None

    @property
    def success_rate(self) -> float:
        """Calculate success rate (0.0 to 1.0)."""
        total = self.success_count + self.failure_count
        if total == 0:
            return 1.0
        return self.success_count / total

    @property
    def is_healthy(self) -> bool:
        """Check if node is considered healthy."""
        # Unhealthy if 3+ consecutive failures
        if self.consecutive_failures >= 3:
            return False
        # Unhealthy if success rate < 50% with at least 10 attempts
        if self.success_count + self.failure_count >= 10:
            if self.success_rate < 0.5:
                return False
        return self.is_available

    def record_success(self, latency_ms: float):
        """Record a successful operation."""
        self.success_count += 1
        self.consecutive_failures = 0
        self.is_available = True
        self.last_check = time.time()
        # Exponential moving average for latency
        if self.avg_latency_ms == 0:
            self.avg_latency_ms = latency_ms
        else:
            self.avg_latency_ms = 0.8 * self.avg_latency_ms + 0.2 * latency_ms

    def record_failure(self, error: str):
        """Record a failed operation."""
        self.failure_count += 1
        self.consecutive_failures += 1
        self.last_error = error
        self.last_check = time.time()
        # Mark unavailable after 3 consecutive failures
        if self.consecutive_failures >= 3:
            self.is_available = False

    def reset(self):
        """Reset health stats (e.g., after recovery)."""
        self.consecutive_failures = 0
        self.is_available = True
        self.last_error = None


@dataclass
class NodeConfig:
    """Configuration for a processing node."""
    node_type: NodeType
    name: str
    priority: int = 0  # Higher = preferred
    latency_threshold_ms: float = 500.0
    timeout_ms: float = 5000.0
    max_retries: int = 2
    enabled: bool = True
    metadata: Dict[str, Any] = field(default_factory=dict)


# Default node configurations
DEFAULT_NODES: Dict[NodeType, NodeConfig] = {
    NodeType.LOCAL: NodeConfig(
        node_type=NodeType.LOCAL,
        name="Local Embeddings",
        priority=80,  # High priority for privacy
        latency_threshold_ms=100.0,
        timeout_ms=10000.0,  # Model loading can be slow
        max_retries=1,
    ),
    NodeType.CLOUD: NodeConfig(
        node_type=NodeType.CLOUD,
        name="Cloud Embeddings",
        priority=90,  # Highest priority for performance mode
        latency_threshold_ms=500.0,
        timeout_ms=5000.0,
        max_retries=2,
    ),
    NodeType.SYNTHETIC: NodeConfig(
        node_type=NodeType.SYNTHETIC,
        name="Synthetic Embeddings",
        priority=10,  # Last resort
        latency_threshold_ms=50.0,
        timeout_ms=1000.0,
        max_retries=0,  # Never fails
    ),
}


class EdgeRouter:
    """
    Edge/Hybrid Swarm Router for dynamic node selection.

    Routes requests to the best available node based on:
    - Privacy requirements
    - Node health and availability
    - Latency thresholds
    - Success rates
    """

    def __init__(
        self,
        privacy_mode: PrivacyMode = PrivacyMode.BALANCED,
        nodes: Optional[Dict[NodeType, NodeConfig]] = None,
        health_check_interval: float = 60.0,
    ):
        self.privacy_mode = privacy_mode
        self.nodes = nodes or DEFAULT_NODES.copy()
        self.health: Dict[NodeType, NodeHealth] = {
            nt: NodeHealth(node_type=nt) for nt in self.nodes
        }
        self.health_check_interval = health_check_interval
        self._lock = threading.Lock()

        # Apply privacy restrictions
        self._apply_privacy_restrictions()

    def _apply_privacy_restrictions(self):
        """Apply privacy mode restrictions to node availability."""
        if self.privacy_mode == PrivacyMode.STRICT:
            # Disable all cloud/remote nodes
            for node_type in [NodeType.CLOUD, NodeType.EDGE]:
                if node_type in self.nodes:
                    self.nodes[node_type].enabled = False
                    logger.info(f"Disabled {node_type.value} node (strict privacy mode)")

    def get_best_node(self, exclude: Optional[List[NodeType]] = None) -> Optional[NodeType]:
        """
        Get the best available node for processing.

        Args:
            exclude: Node types to exclude from selection

        Returns:
            Best available NodeType or None if all unavailable
        """
        exclude = exclude or []

        # Build candidate list based on privacy mode
        if self.privacy_mode == PrivacyMode.STRICT:
            candidates = [NodeType.LOCAL, NodeType.SYNTHETIC]
        elif self.privacy_mode == PrivacyMode.BALANCED:
            candidates = [NodeType.LOCAL, NodeType.CLOUD, NodeType.SYNTHETIC]
        else:  # PERFORMANCE
            candidates = [NodeType.CLOUD, NodeType.LOCAL, NodeType.SYNTHETIC]

        # Filter by availability and health
        available = []
        for nt in candidates:
            if nt in exclude:
                continue
            if nt not in self.nodes:
                continue
            config = self.nodes[nt]
            health = self.health[nt]

            if not config.enabled:
                continue
            if not health.is_healthy:
                continue

            # Score based on priority and latency
            score = config.priority
            if health.avg_latency_ms > 0:
                # Penalize high latency
                if health.avg_latency_ms > config.latency_threshold_ms:
                    score -= 20
            score += health.success_rate * 10  # Bonus for reliability

            available.append((nt, score))

        if not available:
            return None

        # Sort by score (descending)
        available.sort(key=lambda x: x[1], reverse=True)
        return available[0][0]

    def get_fallback_chain(self) -> List[NodeType]:
        """
        Get ordered list of nodes to try (primary → fallbacks).

        Returns:
            List of NodeTypes in order of preference
        """
        chain = []
        excluded = []

        while True:
            node = self.get_best_node(exclude=excluded)
            if node is None:
                break
            chain.append(node)
            excluded.append(node)

        return chain

    def record_success(self, node_type: NodeType, latency_ms: float):
        """Record successful operation on a node."""
        with self._lock:
            if node_type in self.health:
                self.health[node_type].record_success(latency_ms)

    def record_failure(self, node_type: NodeType, error: str):
        """Record failed operation on a node."""
        with self._lock:
            if node_type in self.health:
                self.health[node_type].record_failure(error)

    def get_health_report(self) -> Dict[str, Any]:
        """Get health report for all nodes."""
        report = {
            "privacy_mode": self.privacy_mode.value,
            "nodes": {},
        }

        for nt, health in self.health.items():
            config = self.nodes.get(nt)
            report["nodes"][nt.value] = {
                "name": config.name if config else nt.value,
                "enabled": config.enabled if config else False,
                "is_healthy": health.is_healthy,
                "is_available": health.is_available,
                "success_rate": f"{health.success_rate:.1%}",
                "avg_latency_ms": f"{health.avg_latency_ms:.1f}",
                "success_count": health.success_count,
                "failure_count": health.failure_count,
                "consecutive_failures": health.consecutive_failures,
                "last_error": health.last_error,
            }

        return report

    def reset_node(self, node_type: NodeType):
        """Reset a node's health status."""
        with self._lock:
            if node_type in self.health:
                self.health[node_type].reset()
                if node_type in self.nodes:
                    self.nodes[node_type].enabled = True


class EmbeddingRouter:
    """
    Smart embedding router with automatic failover.

    Routes embedding requests to the best available backend:
    - Local (sentence-transformers)
    - Cloud (InsAIts API)
    - Synthetic (fallback hashing)

    Features:
    - Automatic failover on errors or high latency
    - Privacy mode enforcement
    - Health tracking and recovery
    - Caching for repeated texts
    """

    def __init__(
        self,
        privacy_mode: PrivacyMode = PrivacyMode.BALANCED,
        api_key: Optional[str] = None,
        api_endpoint: Optional[str] = None,
        cache_size: int = 2000,
        latency_threshold_ms: float = 500.0,
    ):
        self.router = EdgeRouter(privacy_mode=privacy_mode)
        self.api_key = api_key
        self.api_endpoint = api_endpoint or "https://insaitsapi-production.up.railway.app/api/embeddings/generate"
        self.latency_threshold_ms = latency_threshold_ms

        # Import embedding functions
        from .embeddings import (
            get_local_embedding,
            get_synthetic_embedding,
            cache as embedding_cache,
            LOCAL_MODEL_AVAILABLE,
        )
        self._get_local = get_local_embedding
        self._get_synthetic = get_synthetic_embedding
        self._cache = embedding_cache
        self._local_available = LOCAL_MODEL_AVAILABLE

        # Update local node health based on model availability
        if not self._local_available:
            self.router.health[NodeType.LOCAL].is_available = False
            logger.info("Local embeddings unavailable (model not loaded)")

    def get_embedding(
        self,
        text: str,
        force_node: Optional[NodeType] = None,
    ) -> np.ndarray:
        """
        Get embedding for text with automatic routing and failover.

        Args:
            text: Text to embed
            force_node: Force specific node (bypasses routing)

        Returns:
            Embedding vector (384 dimensions)
        """
        # Check cache first
        cached = self._cache.get(text)
        if cached is not None:
            return cached

        # Get fallback chain
        if force_node:
            chain = [force_node]
        else:
            chain = self.router.get_fallback_chain()

        if not chain:
            # No nodes available, force synthetic
            chain = [NodeType.SYNTHETIC]

        last_error = None

        for node_type in chain:
            try:
                start_time = time.time()

                if node_type == NodeType.LOCAL:
                    embedding = self._get_local_embedding(text)
                elif node_type == NodeType.CLOUD:
                    embedding = self._get_cloud_embedding(text)
                elif node_type == NodeType.SYNTHETIC:
                    embedding = self._get_synthetic(text)
                else:
                    continue

                latency_ms = (time.time() - start_time) * 1000
                self.router.record_success(node_type, latency_ms)

                # Check latency threshold (trigger failover next time)
                if latency_ms > self.latency_threshold_ms:
                    logger.warning(
                        f"{node_type.value} latency {latency_ms:.0f}ms > "
                        f"threshold {self.latency_threshold_ms:.0f}ms"
                    )

                # Cache result
                self._cache.set(text, embedding)
                return embedding

            except Exception as e:
                last_error = str(e)
                self.router.record_failure(node_type, last_error)
                logger.warning(f"{node_type.value} failed: {last_error}")
                continue

        # All nodes failed, return synthetic as absolute fallback
        logger.error(f"All nodes failed, using synthetic fallback. Last error: {last_error}")
        return self._get_synthetic(text)

    def _get_local_embedding(self, text: str) -> np.ndarray:
        """Get embedding from local model."""
        if not self._local_available:
            raise RuntimeError("Local model not available")
        return self._get_local(text)

    def _get_cloud_embedding(self, text: str) -> np.ndarray:
        """Get embedding from cloud API."""
        if not self.api_key:
            raise RuntimeError("API key required for cloud embeddings")

        import requests

        response = requests.post(
            self.api_endpoint,
            json={"api_key": self.api_key, "text": text},
            timeout=5.0,
        )

        if response.status_code != 200:
            raise RuntimeError(f"Cloud API error: {response.status_code}")

        data = response.json()
        if "embedding" not in data:
            raise RuntimeError("Invalid cloud response: missing embedding")

        return np.array(data["embedding"], dtype=np.float32)

    def get_batch_embeddings(
        self,
        texts: List[str],
        force_node: Optional[NodeType] = None,
    ) -> List[np.ndarray]:
        """
        Get embeddings for multiple texts.

        Args:
            texts: List of texts to embed
            force_node: Force specific node

        Returns:
            List of embedding vectors
        """
        return [self.get_embedding(t, force_node) for t in texts]

    def get_routing_info(self) -> Dict[str, Any]:
        """Get current routing configuration and health."""
        return {
            "privacy_mode": self.router.privacy_mode.value,
            "latency_threshold_ms": self.latency_threshold_ms,
            "local_available": self._local_available,
            "cloud_configured": bool(self.api_key),
            "health": self.router.get_health_report(),
            "fallback_chain": [n.value for n in self.router.get_fallback_chain()],
        }

    def set_privacy_mode(self, mode: PrivacyMode):
        """Change privacy mode at runtime."""
        self.router.privacy_mode = mode
        self.router._apply_privacy_restrictions()
        logger.info(f"Privacy mode changed to: {mode.value}")


# Convenience functions
def create_router(
    privacy_mode: str = "balanced",
    api_key: Optional[str] = None,
) -> EmbeddingRouter:
    """
    Create an EmbeddingRouter with specified privacy mode.

    Args:
        privacy_mode: "strict", "balanced", or "performance"
        api_key: Optional API key for cloud embeddings

    Returns:
        Configured EmbeddingRouter
    """
    mode = PrivacyMode(privacy_mode.lower())
    return EmbeddingRouter(privacy_mode=mode, api_key=api_key)


def get_privacy_modes() -> List[str]:
    """Get list of available privacy modes."""
    return [m.value for m in PrivacyMode]
