"""BEP-2: External Agent Relay — Cross-Model Bridging.

HTTP API allowing external AI models (Grok, Claude, Gemini, GPT) to register
identity, heartbeat, and participate in the Atlas WITHOUT running a full
Beacon node. This is the "on-ramp" for non-native agents.

Relay agents get the SAME identity format (bcn_*) and appear in the Atlas
just like native agents. No second-class citizens. Their property records
include "relay": true so the system knows they're external.

Beacon 2.8.0 — Elyan Labs.
"""

import hashlib
import json
import secrets
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

from .identity import AgentIdentity, agent_id_from_pubkey
from .storage import _dir, append_jsonl, read_jsonl_tail

RELAY_STATE_FILE = "relay_agents.json"
RELAY_LOG_FILE = "relay_log.jsonl"

# Relay token TTL: 24 hours
RELAY_TOKEN_TTL_S = 86400
# Heartbeat timeout: 15 minutes of silence = inactive
RELAY_SILENCE_THRESHOLD_S = 900
# 1 hour = presumed dead
RELAY_DEAD_THRESHOLD_S = 3600

# Known model providers
KNOWN_PROVIDERS = {
    "xai": "xAI (Grok)",
    "anthropic": "Anthropic (Claude)",
    "google": "Google (Gemini)",
    "openai": "OpenAI (GPT)",
    "meta": "Meta (Llama)",
    "mistral": "Mistral AI",
    "elyan": "Elyan Labs",
    "other": "Independent",
}


class RelayAgent:
    """An external agent registered via the relay."""

    def __init__(self, data: Dict[str, Any]):
        self.agent_id: str = data.get("agent_id", "")
        self.pubkey_hex: str = data.get("pubkey_hex", "")
        self.model_id: str = data.get("model_id", "")
        self.provider: str = data.get("provider", "other")
        self.capabilities: List[str] = data.get("capabilities", [])
        self.webhook_url: str = data.get("webhook_url", "")
        self.relay_token: str = data.get("relay_token", "")
        self.token_expires: int = data.get("token_expires", 0)
        self.registered_at: int = data.get("registered_at", 0)
        self.last_heartbeat: int = data.get("last_heartbeat", 0)
        self.beat_count: int = data.get("beat_count", 0)
        self.status: str = data.get("status", "active")
        self.name: str = data.get("name", "")
        self.metadata: Dict[str, Any] = data.get("metadata", {})

    def to_dict(self) -> Dict[str, Any]:
        return {
            "agent_id": self.agent_id,
            "pubkey_hex": self.pubkey_hex,
            "model_id": self.model_id,
            "provider": self.provider,
            "capabilities": self.capabilities,
            "webhook_url": self.webhook_url,
            "relay_token": self.relay_token,
            "token_expires": self.token_expires,
            "registered_at": self.registered_at,
            "last_heartbeat": self.last_heartbeat,
            "beat_count": self.beat_count,
            "status": self.status,
            "name": self.name,
            "metadata": self.metadata,
        }

    def to_public_dict(self) -> Dict[str, Any]:
        """Public-safe view — no relay token."""
        d = self.to_dict()
        d.pop("relay_token", None)
        d.pop("token_expires", None)
        return d

    def assess_status(self) -> str:
        """Assess liveness based on last heartbeat."""
        if self.status == "shutting_down":
            return "shutting_down"
        now = int(time.time())
        age = now - self.last_heartbeat
        if age <= RELAY_SILENCE_THRESHOLD_S:
            return "active"
        if age <= RELAY_DEAD_THRESHOLD_S:
            return "silent"
        return "presumed_dead"


class RelayManager:
    """Manage external agent relay registrations, heartbeats, and discovery."""

    def __init__(self, data_dir: Optional[Path] = None, host_identity: Optional[AgentIdentity] = None):
        self._dir = data_dir or _dir()
        self._host_identity = host_identity

    def _state_path(self) -> Path:
        return self._dir / RELAY_STATE_FILE

    def _load_agents(self) -> Dict[str, Dict[str, Any]]:
        path = self._state_path()
        if not path.exists():
            return {}
        try:
            return json.loads(path.read_text(encoding="utf-8"))
        except Exception:
            return {}

    def _save_agents(self, agents: Dict[str, Dict[str, Any]]) -> None:
        self._state_path().parent.mkdir(parents=True, exist_ok=True)
        self._state_path().write_text(
            json.dumps(agents, indent=2, sort_keys=True) + "\n",
            encoding="utf-8",
        )

    def _log(self, entry: Dict[str, Any]) -> None:
        append_jsonl(RELAY_LOG_FILE, entry)

    def _generate_token(self) -> str:
        """Generate a relay bearer token."""
        return f"relay_{secrets.token_hex(24)}"

    # ── Registration ──

    def register(
        self,
        pubkey_hex: str,
        model_id: str,
        *,
        provider: str = "other",
        capabilities: Optional[List[str]] = None,
        webhook_url: str = "",
        name: str = "",
        signature: str = "",
        metadata: Optional[Dict] = None,
    ) -> Dict[str, Any]:
        """Register an external agent via relay.

        Args:
            pubkey_hex: Ed25519 public key (64 hex chars).
            model_id: Model identifier (e.g. "grok-3", "claude-opus-4-6").
            provider: Model provider (e.g. "xai", "anthropic").
            capabilities: List of capability domains.
            webhook_url: Optional callback URL for incoming messages.
            name: Human-readable agent name.
            signature: Ed25519 signature of registration payload for verification.
            metadata: Additional metadata.

        Returns:
            Dict with agent_id, relay_token, atlas_address info.
        """
        # Validate pubkey
        try:
            pubkey_bytes = bytes.fromhex(pubkey_hex)
            if len(pubkey_bytes) != 32:
                return {"error": "Invalid pubkey: must be 32 bytes (64 hex chars)"}
        except ValueError:
            return {"error": "Invalid pubkey: not valid hex"}

        # Derive agent_id from pubkey (same format as native agents)
        aid = agent_id_from_pubkey(pubkey_bytes)

        # Verify signature if provided
        if signature:
            reg_payload = json.dumps({
                "model_id": model_id,
                "provider": provider,
                "pubkey_hex": pubkey_hex,
            }, sort_keys=True, separators=(",", ":")).encode("utf-8")
            if not AgentIdentity.verify(pubkey_hex, signature, reg_payload):
                return {"error": "Invalid signature"}

        now = int(time.time())
        token = self._generate_token()
        token_expires = now + RELAY_TOKEN_TTL_S

        agents = self._load_agents()
        agents[aid] = {
            "agent_id": aid,
            "pubkey_hex": pubkey_hex,
            "model_id": model_id,
            "provider": provider,
            "capabilities": capabilities or [],
            "webhook_url": webhook_url,
            "relay_token": token,
            "token_expires": token_expires,
            "registered_at": now,
            "last_heartbeat": now,
            "beat_count": 0,
            "status": "active",
            "name": name or f"{model_id} ({provider})",
            "metadata": metadata or {},
        }
        self._save_agents(agents)

        self._log({
            "ts": now,
            "action": "register",
            "agent_id": aid,
            "model_id": model_id,
            "provider": provider,
        })

        return {
            "ok": True,
            "agent_id": aid,
            "relay_token": token,
            "token_expires": token_expires,
            "ttl_s": RELAY_TOKEN_TTL_S,
            "capabilities_registered": capabilities or [],
        }

    # ── Authentication ──

    def authenticate(self, token: str) -> Optional[RelayAgent]:
        """Validate a relay token and return the agent, or None."""
        agents = self._load_agents()
        now = int(time.time())
        for data in agents.values():
            if data.get("relay_token") == token:
                if data.get("token_expires", 0) < now:
                    return None  # Expired
                return RelayAgent(data)
        return None

    # ── Heartbeat ──

    def heartbeat(
        self,
        agent_id: str,
        token: str,
        *,
        status: str = "alive",
        health: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Process a relay heartbeat. Refreshes token TTL.

        Args:
            agent_id: The relay agent's bcn_ ID.
            token: Bearer token from registration.
            status: One of "alive", "degraded", "shutting_down".
            health: Optional health metrics.

        Returns:
            Dict with heartbeat confirmation and new token expiry.
        """
        agents = self._load_agents()
        agent_data = agents.get(agent_id)

        if not agent_data:
            return {"error": "Agent not registered", "code": "NOT_FOUND"}

        if agent_data.get("relay_token") != token:
            return {"error": "Invalid relay token", "code": "AUTH_FAILED"}

        now = int(time.time())
        if agent_data.get("token_expires", 0) < now:
            return {"error": "Token expired — re-register", "code": "TOKEN_EXPIRED"}

        # Update heartbeat state
        agent_data["last_heartbeat"] = now
        agent_data["beat_count"] = agent_data.get("beat_count", 0) + 1
        agent_data["status"] = status

        # Refresh token TTL
        agent_data["token_expires"] = now + RELAY_TOKEN_TTL_S

        if health:
            agent_data["metadata"]["last_health"] = health

        agents[agent_id] = agent_data
        self._save_agents(agents)

        self._log({
            "ts": now,
            "action": "heartbeat",
            "agent_id": agent_id,
            "status": status,
            "beat_count": agent_data["beat_count"],
        })

        return {
            "ok": True,
            "agent_id": agent_id,
            "beat_count": agent_data["beat_count"],
            "status": status,
            "token_expires": agent_data["token_expires"],
            "assessment": RelayAgent(agent_data).assess_status(),
        }

    # ── Discovery ──

    def discover(self, provider: Optional[str] = None, capability: Optional[str] = None) -> List[Dict[str, Any]]:
        """List registered relay agents (public view, no tokens).

        Args:
            provider: Filter by provider (e.g. "xai", "anthropic").
            capability: Filter by capability domain.

        Returns:
            List of public agent dicts.
        """
        agents = self._load_agents()
        results = []

        for data in agents.values():
            agent = RelayAgent(data)
            # Update status assessment
            data["status"] = agent.assess_status()

            if provider and data.get("provider") != provider:
                continue
            if capability and capability not in data.get("capabilities", []):
                continue

            results.append(agent.to_public_dict())

        results.sort(key=lambda a: a.get("last_heartbeat", 0), reverse=True)
        return results

    def get_agent(self, agent_id: str) -> Optional[Dict[str, Any]]:
        """Get a specific relay agent's public info."""
        agents = self._load_agents()
        data = agents.get(agent_id)
        if not data:
            return None
        agent = RelayAgent(data)
        data["status"] = agent.assess_status()
        return agent.to_public_dict()

    # ── Message Forwarding ──

    def forward_message(
        self,
        from_agent_id: str,
        token: str,
        envelope: Dict[str, Any],
    ) -> Dict[str, Any]:
        """Forward a beacon envelope from a relay agent.

        The relay stamps the envelope with relay metadata and
        stores it in the inbox.

        Args:
            from_agent_id: Sender's agent ID.
            token: Bearer token for authentication.
            envelope: The beacon envelope to forward.

        Returns:
            Dict with forwarding result.
        """
        agent = self.authenticate(token)
        if not agent or agent.agent_id != from_agent_id:
            return {"error": "Authentication failed", "code": "AUTH_FAILED"}

        now = int(time.time())

        # Stamp envelope with relay provenance
        envelope["_relay"] = True
        envelope["_relay_ts"] = now
        envelope["_relay_from"] = from_agent_id

        # Store in inbox
        record = {
            "platform": "relay",
            "from": from_agent_id,
            "received_at": now,
            "text": json.dumps(envelope),
            "envelopes": [envelope],
        }
        append_jsonl("inbox.jsonl", record)

        self._log({
            "ts": now,
            "action": "forward",
            "from": from_agent_id,
            "kind": envelope.get("kind", "unknown"),
        })

        return {
            "ok": True,
            "forwarded": True,
            "kind": envelope.get("kind", ""),
            "nonce": envelope.get("nonce", ""),
        }

    # ── Atlas Integration ──

    def register_in_atlas(self, agent_id: str, atlas_mgr: Any) -> Optional[Dict[str, Any]]:
        """Register a relay agent in the Atlas based on their capabilities.

        Args:
            agent_id: The relay agent's bcn_ ID.
            atlas_mgr: AtlasManager instance.

        Returns:
            Atlas registration result, or None if agent not found.
        """
        agents = self._load_agents()
        data = agents.get(agent_id)
        if not data:
            return None

        domains = data.get("capabilities", ["general"])
        if not domains:
            domains = ["general"]

        metadata = {
            "relay": True,
            "model_id": data.get("model_id", ""),
            "provider": data.get("provider", ""),
        }

        return atlas_mgr.register_agent(
            agent_id=agent_id,
            domains=domains,
            name=data.get("name", agent_id),
            metadata=metadata,
        )

    # ── Cleanup ──

    def prune_dead(self, max_silence_s: Optional[int] = None) -> int:
        """Remove relay agents that have been dead too long. Returns count removed."""
        threshold = max_silence_s or RELAY_DEAD_THRESHOLD_S * 3
        now = int(time.time())
        agents = self._load_agents()

        stale = [
            aid for aid, data in agents.items()
            if (now - data.get("last_heartbeat", 0)) > threshold
        ]
        for aid in stale:
            del agents[aid]
        if stale:
            self._save_agents(agents)
            self._log({
                "ts": now,
                "action": "prune",
                "removed": stale,
                "count": len(stale),
            })
        return len(stale)

    def relay_log(self, limit: int = 100) -> List[Dict[str, Any]]:
        """Read recent relay log entries."""
        return read_jsonl_tail(RELAY_LOG_FILE, limit=limit)

    # ── Stats ──

    def stats(self) -> Dict[str, Any]:
        """Relay system statistics."""
        agents = self._load_agents()
        now = int(time.time())

        by_provider: Dict[str, int] = {}
        active = 0
        silent = 0
        dead = 0

        for data in agents.values():
            agent = RelayAgent(data)
            status = agent.assess_status()
            if status == "active":
                active += 1
            elif status == "silent":
                silent += 1
            else:
                dead += 1

            prov = data.get("provider", "other")
            by_provider[prov] = by_provider.get(prov, 0) + 1

        return {
            "total_agents": len(agents),
            "active": active,
            "silent": silent,
            "presumed_dead": dead,
            "by_provider": by_provider,
            "ts": now,
        }
