# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
High Availability Module - Phase 4

Provides reliability features for production deployments:
- Health check endpoints
- Redis session store (shared state)
- Leader election (active/standby)
- Automatic failover
- Systemd watchdog integration

Architecture:
    ┌─────────────┐        ┌─────────────┐
    │   Node #1   │◄──────►│   Node #2   │
    │  (Active)   │        │  (Standby)  │
    └──────┬──────┘        └──────┬──────┘
           │                      │
           └──────────┬───────────┘
                      │
              ┌───────┴───────┐
              │    Redis      │
              │   (Shared)    │
              └───────────────┘

Configuration:
    ha:
      enabled: true
      mode: active_standby  # or active_active
      node_id: node-1
      redis_url: redis://localhost:6379
      heartbeat_interval: 5
      failover_timeout: 15
"""

from __future__ import annotations

import json
import logging
import os
import secrets
import socket
import threading
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any, Callable, Dict, List, Optional

logger = logging.getLogger(__name__)

# Optional Redis dependency
try:
    import redis

    HAS_REDIS = True
except ImportError:
    HAS_REDIS = False
    logger.debug("redis not installed — HA features unavailable.")


class HAMode(str, Enum):
    """High availability mode."""

    STANDALONE = "standalone"  # Single node, no HA
    ACTIVE_STANDBY = "active_standby"  # One active, one standby
    ACTIVE_ACTIVE = "active_active"  # Both active, load balanced


class NodeState(str, Enum):
    """Node state in HA cluster."""

    STARTING = "starting"
    ACTIVE = "active"
    STANDBY = "standby"
    FAILED = "failed"
    DRAINING = "draining"  # Graceful shutdown


@dataclass
class NodeInfo:
    """Information about a cluster node."""

    node_id: str
    hostname: str
    state: NodeState
    last_heartbeat: datetime
    started_at: datetime
    version: str = "1.8.0"

    # Metrics
    requests_handled: int = 0
    errors: int = 0
    avg_latency_ms: float = 0.0

    # Capabilities
    channels: List[str] = field(default_factory=list)

    def to_dict(self) -> Dict:
        d = asdict(self)
        d["state"] = self.state.value
        d["last_heartbeat"] = self.last_heartbeat.isoformat()
        d["started_at"] = self.started_at.isoformat()
        return d

    @classmethod
    def from_dict(cls, d: Dict) -> "NodeInfo":
        d = d.copy()
        d["state"] = NodeState(d["state"])
        d["last_heartbeat"] = datetime.fromisoformat(d["last_heartbeat"])
        d["started_at"] = datetime.fromisoformat(d["started_at"])
        return cls(**d)

    @property
    def is_alive(self) -> bool:
        """Check if node is alive (heartbeat within last 30 seconds)."""
        return (datetime.utcnow() - self.last_heartbeat).total_seconds() < 30


@dataclass
class HAConfig:
    """High availability configuration."""

    enabled: bool = False
    mode: HAMode = HAMode.STANDALONE

    # Node identity
    node_id: str = None  # Auto-generated if not set

    # Redis connection
    redis_url: str = "redis://localhost:6379"
    redis_db: int = 1  # Use separate DB for HA
    redis_prefix: str = "familiar:ha:"

    # Timing
    heartbeat_interval: int = 5  # Seconds between heartbeats
    failover_timeout: int = 15  # Seconds before failover
    election_timeout: int = 10  # Seconds for leader election

    # Callbacks
    on_become_active: Callable = None
    on_become_standby: Callable = None
    on_failover: Callable = None

    def __post_init__(self):
        if self.node_id is None:
            self.node_id = f"{socket.gethostname()}-{secrets.token_hex(4)}"

    @classmethod
    def from_dict(cls, d: Dict) -> "HAConfig":
        d = d.copy()
        if "mode" in d:
            d["mode"] = HAMode(d["mode"])
        return cls(**d)

    @classmethod
    def from_env(cls) -> "HAConfig":
        """Load from environment variables."""
        return cls(
            enabled=os.environ.get("HA_ENABLED", "").lower() == "true",
            mode=HAMode(os.environ.get("HA_MODE", "standalone")),
            node_id=os.environ.get("HA_NODE_ID"),
            redis_url=os.environ.get("REDIS_URL", "redis://localhost:6379"),
            heartbeat_interval=int(os.environ.get("HA_HEARTBEAT_INTERVAL", "5")),
            failover_timeout=int(os.environ.get("HA_FAILOVER_TIMEOUT", "15")),
        )


# ============================================================
# HEALTH CHECK
# ============================================================


@dataclass
class HealthStatus:
    """System health status."""

    healthy: bool
    status: str  # ok, degraded, unhealthy
    node_id: str
    node_state: NodeState
    uptime_seconds: float

    # Component health
    redis_connected: bool = False
    database_connected: bool = True
    llm_available: bool = True

    # Metrics
    requests_per_minute: float = 0.0
    error_rate: float = 0.0
    avg_latency_ms: float = 0.0

    # Cluster info (if HA enabled)
    cluster_size: int = 1
    active_nodes: int = 1

    checks: Dict[str, bool] = field(default_factory=dict)

    def to_dict(self) -> Dict:
        d = asdict(self)
        d["node_state"] = self.node_state.value
        return d


class HealthChecker:
    """
    Health check system for monitoring.

    Usage:
        checker = HealthChecker()

        # Add custom checks
        checker.add_check("database", lambda: db.ping())
        checker.add_check("llm", lambda: provider.is_available())

        # Get health status
        status = checker.check()
        print(f"Healthy: {status.healthy}")
    """

    def __init__(self, node_id: str = None):
        self.node_id = node_id or socket.gethostname()
        self.started_at = datetime.utcnow()
        self._checks: Dict[str, Callable[[], bool]] = {}
        self._metrics: Dict[str, float] = {}

        # Request tracking
        self._request_count = 0
        self._error_count = 0
        self._latency_sum = 0.0
        self._last_reset = time.time()

    def add_check(self, name: str, check_fn: Callable[[], bool]):
        """Add a health check function."""
        self._checks[name] = check_fn

    def record_request(self, latency_ms: float, error: bool = False):
        """Record a request for metrics."""
        self._request_count += 1
        self._latency_sum += latency_ms
        if error:
            self._error_count += 1

    def check(self, ha_manager: "HAManager" = None) -> HealthStatus:
        """
        Run all health checks and return status.

        Args:
            ha_manager: Optional HA manager for cluster info
        """
        checks = {}

        # Run all registered checks
        for name, check_fn in self._checks.items():
            try:
                checks[name] = check_fn()
            except Exception as e:
                logger.warning(f"Health check '{name}' failed: {e}")
                checks[name] = False

        # Calculate metrics
        elapsed = max(time.time() - self._last_reset, 1)
        rpm = (self._request_count / elapsed) * 60
        error_rate = self._error_count / max(self._request_count, 1)
        avg_latency = self._latency_sum / max(self._request_count, 1)

        # Determine overall status
        all_healthy = all(checks.values()) if checks else True

        if all_healthy and error_rate < 0.05:
            status = "ok"
            healthy = True
        elif all_healthy or error_rate < 0.20:
            status = "degraded"
            healthy = True
        else:
            status = "unhealthy"
            healthy = False

        # Get HA info
        node_state = NodeState.ACTIVE
        cluster_size = 1
        active_nodes = 1
        redis_connected = False

        if ha_manager:
            node_state = ha_manager.state
            cluster_size = len(ha_manager.get_nodes())
            active_nodes = len([n for n in ha_manager.get_nodes() if n.state == NodeState.ACTIVE])
            redis_connected = ha_manager._redis is not None

        return HealthStatus(
            healthy=healthy,
            status=status,
            node_id=self.node_id,
            node_state=node_state,
            uptime_seconds=(datetime.utcnow() - self.started_at).total_seconds(),
            redis_connected=redis_connected,
            database_connected=checks.get("database", True),
            llm_available=checks.get("llm", True),
            requests_per_minute=rpm,
            error_rate=error_rate,
            avg_latency_ms=avg_latency,
            cluster_size=cluster_size,
            active_nodes=active_nodes,
            checks=checks,
        )

    def reset_metrics(self):
        """Reset request metrics (call periodically)."""
        self._request_count = 0
        self._error_count = 0
        self._latency_sum = 0.0
        self._last_reset = time.time()


# ============================================================
# HA MANAGER
# ============================================================


class HAManager:
    """
    High availability manager for active/standby or active/active modes.

    Usage:
        config = HAConfig(
            enabled=True,
            mode=HAMode.ACTIVE_STANDBY,
            redis_url="redis://localhost:6379"
        )

        ha = HAManager(config)
        ha.start()

        # Check if this node should handle requests
        if ha.is_active:
            process_request()

        # Graceful shutdown
        ha.stop()
    """

    def __init__(self, config: HAConfig):
        self.config = config
        self.node_id = config.node_id
        self.state = NodeState.STARTING

        self._redis: Optional[redis.Redis] = None
        self._heartbeat_thread: Optional[threading.Thread] = None
        self._running = False
        self._leader_lock = threading.Lock()

        # Node info
        self._node_info = NodeInfo(
            node_id=self.node_id,
            hostname=socket.gethostname(),
            state=NodeState.STARTING,
            last_heartbeat=datetime.utcnow(),
            started_at=datetime.utcnow(),
        )

    @property
    def is_active(self) -> bool:
        """Check if this node is active (should handle requests)."""
        if not self.config.enabled:
            return True  # Standalone mode - always active

        return self.state == NodeState.ACTIVE

    @property
    def is_leader(self) -> bool:
        """Alias for is_active."""
        return self.is_active

    def start(self):
        """Start the HA manager."""
        if not self.config.enabled:
            self.state = NodeState.ACTIVE
            logger.info("HA disabled - running in standalone mode")
            return

        if not HAS_REDIS:
            raise RuntimeError("Redis required for HA mode.")

        # Connect to Redis
        self._connect_redis()

        # Register this node
        self._register_node()

        # Start heartbeat thread
        self._running = True
        self._heartbeat_thread = threading.Thread(target=self._heartbeat_loop, daemon=True)
        self._heartbeat_thread.start()

        # Initial leader election
        self._elect_leader()

        logger.info(f"HA manager started: node={self.node_id}, state={self.state.value}")

    def stop(self):
        """Stop the HA manager gracefully."""
        if not self.config.enabled:
            return

        logger.info(f"Stopping HA manager: node={self.node_id}")

        # Set draining state
        self.state = NodeState.DRAINING
        self._update_node_state()

        # Stop heartbeat
        self._running = False
        if self._heartbeat_thread:
            self._heartbeat_thread.join(timeout=5)

        # Unregister node
        self._unregister_node()

        # Force leader election for remaining nodes
        self._trigger_election()

        # Close Redis
        if self._redis:
            self._redis.close()

    def get_nodes(self) -> List[NodeInfo]:
        """Get list of all nodes in cluster."""
        if not self._redis:
            return [self._node_info]

        nodes = []
        try:
            pattern = f"{self.config.redis_prefix}node:*"
            for key in self._redis.scan_iter(pattern):
                data = self._redis.get(key)
                if data:
                    nodes.append(NodeInfo.from_dict(json.loads(data)))
        except Exception as e:
            logger.error(f"Failed to get nodes: {e}")

        return nodes

    def get_leader(self) -> Optional[str]:
        """Get current leader node ID."""
        if not self._redis:
            return self.node_id

        try:
            leader = self._redis.get(f"{self.config.redis_prefix}leader")
            return leader.decode() if leader else None
        except Exception as e:
            logger.error(f"Failed to get leader: {e}")
            return None

    def _connect_redis(self):
        """Connect to Redis."""
        try:
            self._redis = redis.from_url(
                self.config.redis_url, db=self.config.redis_db, decode_responses=False
            )
            self._redis.ping()
            logger.info(f"Connected to Redis: {self.config.redis_url}")
        except Exception as e:
            logger.error(f"Failed to connect to Redis: {e}")
            raise

    def _register_node(self):
        """Register this node in Redis."""
        self._node_info.last_heartbeat = datetime.utcnow()
        key = f"{self.config.redis_prefix}node:{self.node_id}"
        self._redis.setex(
            key, self.config.failover_timeout * 2, json.dumps(self._node_info.to_dict())
        )

    def _unregister_node(self):
        """Remove this node from Redis."""
        key = f"{self.config.redis_prefix}node:{self.node_id}"
        self._redis.delete(key)

    def _update_node_state(self):
        """Update node state in Redis."""
        self._node_info.state = self.state
        self._node_info.last_heartbeat = datetime.utcnow()
        self._register_node()

    def _heartbeat_loop(self):
        """Background thread for heartbeats."""
        while self._running:
            try:
                self._heartbeat()
            except Exception as e:
                logger.error(f"Heartbeat error: {e}")

            time.sleep(self.config.heartbeat_interval)

    def _heartbeat(self):
        """Send heartbeat and check cluster health."""
        # Update our node info
        self._node_info.last_heartbeat = datetime.utcnow()
        self._node_info.state = self.state
        self._register_node()

        # Check if leader is alive
        leader = self.get_leader()

        if leader and leader != self.node_id:
            # We're not the leader - check if leader is alive
            leader_key = f"{self.config.redis_prefix}node:{leader}"
            leader_data = self._redis.get(leader_key)

            if leader_data:
                leader_info = NodeInfo.from_dict(json.loads(leader_data))
                if not leader_info.is_alive:
                    logger.warning(f"Leader {leader} appears dead - triggering election")
                    self._elect_leader()
            else:
                logger.warning(f"Leader {leader} not found - triggering election")
                self._elect_leader()

        elif not leader:
            # No leader - elect one
            self._elect_leader()

    def _elect_leader(self):
        """Perform leader election."""
        with self._leader_lock:
            leader_key = f"{self.config.redis_prefix}leader"
            lock_key = f"{self.config.redis_prefix}election_lock"

            # Try to acquire election lock
            if self._redis.set(lock_key, self.node_id, nx=True, ex=self.config.election_timeout):
                try:
                    # We have the lock - check current leader
                    current_leader = self._redis.get(leader_key)

                    if current_leader:
                        current_leader = current_leader.decode()
                        # Check if current leader is alive
                        leader_data = self._redis.get(
                            f"{self.config.redis_prefix}node:{current_leader}"
                        )
                        if leader_data:
                            leader_info = NodeInfo.from_dict(json.loads(leader_data))
                            if leader_info.is_alive:
                                # Leader is fine
                                if current_leader == self.node_id:
                                    self._become_active()
                                else:
                                    self._become_standby()
                                return

                    # No valid leader - become leader
                    self._redis.setex(leader_key, self.config.failover_timeout * 2, self.node_id)
                    self._become_active()
                    logger.info(f"Node {self.node_id} elected as leader")

                finally:
                    self._redis.delete(lock_key)
            else:
                # Another node is running election
                time.sleep(1)

                # Check result
                leader = self.get_leader()
                if leader == self.node_id:
                    self._become_active()
                else:
                    self._become_standby()

    def _trigger_election(self):
        """Force a new election."""
        if self._redis:
            self._redis.delete(f"{self.config.redis_prefix}leader")

    def _become_active(self):
        """Transition to active state."""
        if self.state != NodeState.ACTIVE:
            old_state = self.state
            self.state = NodeState.ACTIVE
            self._update_node_state()

            logger.info(f"Node {self.node_id} became ACTIVE (was {old_state.value})")

            if self.config.on_become_active:
                try:
                    self.config.on_become_active()
                except Exception as e:
                    logger.error(f"on_become_active callback failed: {e}")

    def _become_standby(self):
        """Transition to standby state."""
        if self.state != NodeState.STANDBY:
            old_state = self.state
            self.state = NodeState.STANDBY
            self._update_node_state()

            logger.info(f"Node {self.node_id} became STANDBY (was {old_state.value})")

            if self.config.on_become_standby:
                try:
                    self.config.on_become_standby()
                except Exception as e:
                    logger.error(f"on_become_standby callback failed: {e}")


# ============================================================
# REDIS SESSION STORE
# ============================================================


class RedisSessionStore:
    """
    Redis-backed session store for shared state across nodes.

    Usage:
        store = RedisSessionStore("redis://localhost:6379")

        # Store session
        store.set("session:abc123", {"user_id": "user1"}, ttl=3600)

        # Get session
        data = store.get("session:abc123")
    """

    def __init__(self, redis_url: str, prefix: str = "familiar:session:"):
        if not HAS_REDIS:
            raise RuntimeError("Redis required for session storage.")

        self.prefix = prefix
        self._redis = redis.from_url(redis_url, decode_responses=True)

    def set(self, key: str, value: Any, ttl: int = 3600):
        """Set a value with optional TTL."""
        full_key = f"{self.prefix}{key}"
        self._redis.setex(full_key, ttl, json.dumps(value))

    def get(self, key: str) -> Optional[Any]:
        """Get a value."""
        full_key = f"{self.prefix}{key}"
        data = self._redis.get(full_key)
        return json.loads(data) if data else None

    def delete(self, key: str):
        """Delete a key."""
        full_key = f"{self.prefix}{key}"
        self._redis.delete(full_key)

    def exists(self, key: str) -> bool:
        """Check if key exists."""
        full_key = f"{self.prefix}{key}"
        return self._redis.exists(full_key) > 0

    def extend_ttl(self, key: str, ttl: int):
        """Extend TTL on existing key."""
        full_key = f"{self.prefix}{key}"
        self._redis.expire(full_key, ttl)

    def get_all_keys(self, pattern: str = "*") -> List[str]:
        """Get all keys matching pattern."""
        full_pattern = f"{self.prefix}{pattern}"
        keys = self._redis.keys(full_pattern)
        return [k.replace(self.prefix, "") for k in keys]


# ============================================================
# FLASK INTEGRATION
# ============================================================


def setup_ha_routes(app, ha_manager: HAManager, health_checker: HealthChecker):
    """
    Add HA routes to Flask app.

    Routes:
        GET /health - Basic health check (for load balancers)
        GET /health/detailed - Detailed health status
        GET /ha/status - HA cluster status
        POST /ha/drain - Drain node (graceful shutdown prep)
    """
    from flask import jsonify

    @app.route("/health")
    def health():
        """Simple health check for load balancers."""
        status = health_checker.check(ha_manager)

        if status.healthy:
            return jsonify({"status": "ok"}), 200
        else:
            return jsonify({"status": "unhealthy"}), 503

    def _require_ha_auth():
        """Check for HA shared secret on sensitive endpoints."""
        from flask import request as req

        ha_secret = os.environ.get("HA_SHARED_SECRET")
        if not ha_secret:
            return None  # No secret configured, allow (backward compat for dev)
        token = req.headers.get("X-HA-Token", "")
        import hmac as _hmac

        if not _hmac.compare_digest(token, ha_secret):
            return jsonify({"error": "Unauthorized"}), 401
        return None

    @app.route("/health/detailed")
    def health_detailed():
        """Detailed health status (requires HA auth)."""
        auth_err = _require_ha_auth()
        if auth_err:
            return auth_err
        status = health_checker.check(ha_manager)
        return jsonify(status.to_dict()), 200 if status.healthy else 503

    @app.route("/ha/status")
    def ha_status():
        """HA cluster status (requires HA auth)."""
        auth_err = _require_ha_auth()
        if auth_err:
            return auth_err
        if not ha_manager.config.enabled:
            return jsonify({"enabled": False, "mode": "standalone", "node_id": ha_manager.node_id})

        nodes = ha_manager.get_nodes()
        leader = ha_manager.get_leader()

        return jsonify(
            {
                "enabled": True,
                "mode": ha_manager.config.mode.value,
                "node_id": ha_manager.node_id,
                "state": ha_manager.state.value,
                "is_leader": ha_manager.is_leader,
                "leader": leader,
                "nodes": [n.to_dict() for n in nodes],
                "cluster_size": len(nodes),
                "active_nodes": len([n for n in nodes if n.state == NodeState.ACTIVE]),
            }
        )

    @app.route("/ha/drain", methods=["POST"])
    def ha_drain():
        """Drain this node (prepare for shutdown). Requires HA auth."""
        auth_err = _require_ha_auth()
        if auth_err:
            return auth_err
        ha_manager.state = NodeState.DRAINING
        ha_manager._update_node_state()
        ha_manager._trigger_election()

        return jsonify({"status": "draining", "node_id": ha_manager.node_id})

    return app


# ============================================================
# SYSTEMD WATCHDOG
# ============================================================


class SystemdWatchdog:
    """
    Systemd watchdog integration for process monitoring.

    Usage:
        watchdog = SystemdWatchdog()

        if watchdog.enabled:
            # Notify systemd we're ready
            watchdog.ready()

            # In main loop, send heartbeats
            while running:
                do_work()
                watchdog.ping()

    Systemd service file:
        [Service]
        WatchdogSec=30
        NotifyAccess=all
    """

    def __init__(self):
        self.enabled = "WATCHDOG_USEC" in os.environ
        self._socket_path = os.environ.get("NOTIFY_SOCKET")

    def ready(self):
        """Notify systemd that service is ready."""
        self._notify("READY=1")

    def ping(self):
        """Send watchdog ping."""
        self._notify("WATCHDOG=1")

    def stopping(self):
        """Notify systemd that service is stopping."""
        self._notify("STOPPING=1")

    def status(self, message: str):
        """Update status message in systemd."""
        self._notify(f"STATUS={message}")

    def _notify(self, message: str):
        """Send notification to systemd."""
        if not self._socket_path:
            return

        try:
            sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
            sock.connect(self._socket_path)
            sock.send(message.encode())
            sock.close()
        except Exception as e:
            logger.debug(f"Systemd notify failed: {e}")


# ============================================================
# SINGLETON
# ============================================================

_ha_manager: Optional[HAManager] = None
_health_checker: Optional[HealthChecker] = None


def get_ha_manager() -> Optional[HAManager]:
    """Get global HA manager."""
    return _ha_manager


def get_health_checker() -> HealthChecker:
    """Get or create global health checker."""
    global _health_checker
    if _health_checker is None:
        _health_checker = HealthChecker()
    return _health_checker


def init_ha(config: HAConfig = None) -> HAManager:
    """Initialize HA manager."""
    global _ha_manager

    if config is None:
        config = HAConfig.from_env()

    _ha_manager = HAManager(config)
    _ha_manager.start()

    return _ha_manager


def shutdown_ha():
    """Shutdown HA manager."""
    global _ha_manager
    if _ha_manager:
        _ha_manager.stop()
        _ha_manager = None
