"""Service & Agent Registry — control plane registry for discovery and lookup.

Provides registry contracts and in-memory implementation for:
- Agent registration with role-based lookup
- Service registration with health checking
- Capability-based discovery
"""
from __future__ import annotations

import logging
import time
from collections import defaultdict
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Protocol, runtime_checkable

logger = logging.getLogger(__name__)


class ServiceStatus(str, Enum):
    """Health status of a registered service."""
    HEALTHY = "healthy"
    DEGRADED = "degraded"
    UNHEALTHY = "unhealthy"
    UNKNOWN = "unknown"


@dataclass
class ServiceRecord:
    """Registered service metadata."""
    service_id: str
    name: str
    endpoint: str | None = None
    capabilities: set[str] = field(default_factory=set)
    metadata: dict[str, Any] = field(default_factory=dict)
    status: ServiceStatus = ServiceStatus.UNKNOWN
    registered_at: float = field(default_factory=time.time)
    last_heartbeat: float | None = None
    tags: dict[str, str] = field(default_factory=dict)


@dataclass
class AgentRecord:
    """Registered agent metadata."""
    agent_id: str
    role: str
    identity: Any = None  # AgentIdentity when available
    capabilities: set[str] = field(default_factory=set)
    metadata: dict[str, Any] = field(default_factory=dict)
    registered_at: float = field(default_factory=time.time)
    status: ServiceStatus = ServiceStatus.HEALTHY


@runtime_checkable
class Registry(Protocol):
    """Registry contract for service and agent discovery."""

    # Agent operations
    def register_agent(self, agent_id: str, role: str, *, identity: Any = None,
                       capabilities: set[str] | None = None, metadata: dict[str, Any] | None = None) -> AgentRecord:
        """Register an agent and return its record."""
        ...

    def unregister_agent(self, agent_id: str) -> bool:
        """Remove an agent. Returns True if existed."""
        ...

    def get_agent(self, agent_id: str) -> AgentRecord | None:
        """Retrieve agent by ID."""
        ...

    def list_agents_by_role(self, role: str) -> list[AgentRecord]:
        """List all agents with a given role."""
        ...

    def list_all_agents(self) -> list[AgentRecord]:
        """List all registered agents."""
        ...

    # Service operations
    def register_service(self, service_id: str, name: str, *,
                         endpoint: str | None = None,
                         capabilities: set[str] | None = None,
                         metadata: dict[str, Any] | None = None,
                         tags: dict[str, str] | None = None) -> ServiceRecord:
        """Register a service and return its record."""
        ...

    def unregister_service(self, service_id: str) -> bool:
        """Remove a service. Returns True if existed."""
        ...

    def get_service(self, service_id: str) -> ServiceRecord | None:
        """Retrieve service by ID."""
        ...

    def list_services_by_capability(self, capability: str) -> list[ServiceRecord]:
        """List services providing a capability."""
        ...

    def list_all_services(self) -> list[ServiceRecord]:
        """List all registered services."""
        ...

    def heartbeat(self, service_id: str, status: ServiceStatus = ServiceStatus.HEALTHY) -> bool:
        """Update service heartbeat. Returns True if service exists."""
        ...

    # Discovery
    def discover(self, *, role: str | None = None, capability: str | None = None,
                 tags: dict[str, str] | None = None) -> list[AgentRecord | ServiceRecord]:
        """Discover agents/services matching criteria."""
        ...


class InMemoryRegistry:
    """Default in-memory Registry implementation.

    Thread-safe for single-threaded asyncio use. Production deployments
    may substitute distributed implementations (etcd, Consul, etc.).
    """

    def __init__(self) -> None:
        self._agents: dict[str, AgentRecord] = {}
        self._roles: dict[str, set[str]] = defaultdict(set)
        self._services: dict[str, ServiceRecord] = {}
        self._capabilities: dict[str, set[str]] = defaultdict(set)

    # Agent operations
    def register_agent(self, agent_id: str, role: str, *, identity: Any = None,
                       capabilities: set[str] | None = None,
                       metadata: dict[str, Any] | None = None) -> AgentRecord:
        record = AgentRecord(
            agent_id=agent_id,
            role=role,
            identity=identity,
            capabilities=capabilities or set(),
            metadata=metadata or {},
        )
        self._agents[agent_id] = record
        self._roles[role].add(agent_id)
        logger.info("Registered agent %s (role=%s)", agent_id, role)
        return record

    def unregister_agent(self, agent_id: str) -> bool:
        if agent_id not in self._agents:
            return False
        record = self._agents.pop(agent_id)
        self._roles[record.role].discard(agent_id)
        logger.info("Unregistered agent %s", agent_id)
        return True

    def get_agent(self, agent_id: str) -> AgentRecord | None:
        return self._agents.get(agent_id)

    def list_agents_by_role(self, role: str) -> list[AgentRecord]:
        return [self._agents[aid] for aid in self._roles.get(role, set()) if aid in self._agents]

    def list_all_agents(self) -> list[AgentRecord]:
        return list(self._agents.values())

    # Service operations
    def register_service(self, service_id: str, name: str, *,
                         endpoint: str | None = None,
                         capabilities: set[str] | None = None,
                         metadata: dict[str, Any] | None = None,
                         tags: dict[str, str] | None = None) -> ServiceRecord:
        record = ServiceRecord(
            service_id=service_id,
            name=name,
            endpoint=endpoint,
            capabilities=capabilities or set(),
            metadata=metadata or {},
            tags=tags or {},
        )
        self._services[service_id] = record
        for cap in record.capabilities:
            self._capabilities[cap].add(service_id)
        logger.info("Registered service %s (%s)", service_id, name)
        return record

    def unregister_service(self, service_id: str) -> bool:
        if service_id not in self._services:
            return False
        record = self._services.pop(service_id)
        for cap in record.capabilities:
            self._capabilities[cap].discard(service_id)
        logger.info("Unregistered service %s", service_id)
        return True

    def get_service(self, service_id: str) -> ServiceRecord | None:
        return self._services.get(service_id)

    def list_services_by_capability(self, capability: str) -> list[ServiceRecord]:
        return [self._services[sid] for sid in self._capabilities.get(capability, set()) if sid in self._services]

    def list_all_services(self) -> list[ServiceRecord]:
        return list(self._services.values())

    def heartbeat(self, service_id: str, status: ServiceStatus = ServiceStatus.HEALTHY) -> bool:
        if service_id not in self._services:
            return False
        record = self._services[service_id]
        record.last_heartbeat = time.time()
        record.status = status
        return True

    # Discovery
    def discover(self, *, role: str | None = None, capability: str | None = None,
                 tags: dict[str, str] | None = None) -> list[AgentRecord | ServiceRecord]:
        results: list[AgentRecord | ServiceRecord] = []

        if role:
            results.extend(self.list_agents_by_role(role))

        if capability:
            results.extend(self.list_services_by_capability(capability))

        if tags:
            for svc in self._services.values():
                if all(svc.tags.get(k) == v for k, v in tags.items()):
                    results.append(svc)

        return results


# Default singleton for convenience
_default_registry: InMemoryRegistry | None = None


def get_registry() -> InMemoryRegistry:
    """Get the default registry singleton."""
    global _default_registry
    if _default_registry is None:
        _default_registry = InMemoryRegistry()
    return _default_registry


__all__ = [
    "AgentRecord",
    "InMemoryRegistry",
    "Registry",
    "ServiceRecord",
    "ServiceStatus",
    "get_registry",
]