"""
AgentRegistry with vector search, metrics, and comprehensive features.

This module provides the core agent registry functionality including
registration, discovery, vector search, and metrics tracking.
"""

import asyncio
import hashlib
import json
import logging
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, Callable, Coroutine, Dict, List, Optional, Union

from ..models.agent import (
    AgentInfo,
    AgentMetrics,
    AgentRegistration,
    HeartbeatResponse,
    SearchRequest,
    SearchResponse,
    TokenData,
)
from ..models.token import TokenMintRequest
from ..services import get_openai_service, get_redis_service
from ..utils.tracing import set_span_attributes, trace_function, trace_operation
from .config import config
from .exceptions import (
    AgentNotFoundError,
    AgentRegistrationError,
    ConfigurationError,
    DuplicateAgentError,
    VectorSearchError,
)
from .storage_adapter import StorageAdapter
from .token_service import TokenService

logger = logging.getLogger("agent-registry")


class AgentRegistry:
    """
    Agent Registry with vector search and comprehensive features.

    This singleton class manages agent registration, discovery, metrics,
    and provides vector search capabilities using Azure OpenAI.
    """

    on_update_callbacks: List[Callable[[Any], Coroutine[Any, Any, None]]] = []
    _instance = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def __init__(self):
        """Initialize the AgentRegistry with configuration validation."""
        if hasattr(self, "_initialized"):
            return
        self._initialized = True

        # Ensure data directory exists
        config.ensure_data_directory()

        # Validate configuration
        missing_config = config.validate_required_config()
        if missing_config:
            # Record via tracing and structured error type without failing startup
            config_error = ConfigurationError(
                "Missing configuration", {"missing_keys": missing_config}
            )
            logger.warning(
                f"Missing configuration: {', '.join(missing_config)} | {config_error}"
            )
            logger.warning("Some features may be disabled.")
            with trace_operation(
                "registry.configuration",
                {
                    "component": "registry",
                    "missing_keys": ",".join(missing_config),
                },
            ):
                set_span_attributes({"config.missing": ",".join(missing_config)})

        # Initialize TokenService for JWT operations
        self.token_service = TokenService()

        # Initialize services
        self.openai_service = get_openai_service()
        self.redis_service = get_redis_service()

        # Fallback storage
        self.backup_agents: dict = {}
        self.backup_embeddings: dict = {}
        self.backup_metrics: dict = {}
        self.backup_info_hashes: dict = {}  # Cache for embedding info hashes
        self.backup_agent_keys: dict = {}  # Cache for agent key -> agent_id mappings

        # Unified storage adapter (handles redis + fallback)
        self.storage = StorageAdapter(self.redis_service.get_client())
        self.storage.register_bucket("agent:data", self.backup_agents)
        self.storage.register_bucket("agent:embeddings", self.backup_embeddings)
        self.storage.register_bucket("agent:metrics", self.backup_metrics)
        self.storage.register_bucket(
            "agent:info_hashes", self.backup_info_hashes
        )  # Cache for embedding info hashes
        self.storage.register_bucket(
            "agent:keys", self.backup_agent_keys
        )  # Cache for agent key -> agent_id mappings

        self._lock = asyncio.Lock()
        self._load_state()
        logger.info("AgentRegistry initialized")

    def _load_state(self):
        """Load registry state from file."""
        state_file = config.STATE_FILE
        if state_file and Path(state_file).exists():
            # Check if file is empty or too small to be valid JSON
            file_size = Path(state_file).stat().st_size
            if file_size < 2:  # Minimum valid JSON is "{}" (2 bytes)
                logger.debug(
                    f"State file is empty or too small ({file_size} bytes), skipping load"
                )
                return

            try:
                with open(state_file, "r") as f:
                    content = f.read().strip()
                    if not content:
                        logger.debug("State file is empty, skipping load")
                        return
                    state = json.loads(content)

                    # Load agent data
                    for aid, info in state.get("agents", {}).items():
                        info["last_seen"] = datetime.fromisoformat(info["last_seen"])
                        info["registered_at"] = datetime.fromisoformat(
                            info["registered_at"]
                        )
                        self.backup_agents[aid] = info

                    # Load info hashes (for embedding cache)
                    if "info_hashes" in state:
                        self.backup_info_hashes.update(state["info_hashes"])

                logger.info(f"Loaded {len(self.backup_agents)} agents from state file")
            except json.JSONDecodeError as e:
                logger.warning(f"Invalid JSON in registry state file: {e}")
                # Don't try fallback if JSON is malformed - file needs to be recreated
            except Exception as e:
                logger.error(f"Error loading registry state: {e}")
                # Fallback to old format for backwards compatibility
                try:
                    with open(state_file, "r") as f:
                        content = f.read().strip()
                        if not content:
                            return
                        state = json.loads(content)
                        for aid, info in state.items():
                            if isinstance(info, dict) and "last_seen" in info:
                                info["last_seen"] = datetime.fromisoformat(
                                    info["last_seen"]
                                )
                                info["registered_at"] = datetime.fromisoformat(
                                    info["registered_at"]
                                )
                                self.backup_agents[aid] = info
                    logger.info(
                        f"Loaded {len(self.backup_agents)} agents from old format state file"
                    )
                except json.JSONDecodeError:
                    logger.warning(
                        "State file contains invalid JSON, will be recreated on next save"
                    )
                except Exception as e2:
                    logger.error(f"Error loading old format registry state: {e2}")

    async def _save_state(self):
        """Save registry state to file."""
        state_file = config.STATE_FILE
        if not state_file:
            return

        try:
            # Ensure directory exists
            config.ensure_data_directory()

            # Get all agents from storage
            all_agents = await self.get_all_agent_data()
            agents_state = {}
            for aid, info in all_agents.items():
                agents_state[aid] = {
                    **info,
                    "last_seen": (
                        info["last_seen"].isoformat()
                        if isinstance(info["last_seen"], datetime)
                        else info["last_seen"]
                    ),
                    "registered_at": (
                        info["registered_at"].isoformat()
                        if isinstance(info["registered_at"], datetime)
                        else info["registered_at"]
                    ),
                }

            # Create state with both agents and info hashes
            state = {
                "agents": agents_state,
                "info_hashes": self.backup_info_hashes.copy(),  # Save embedding cache
            }

            with open(state_file, "w") as f:
                json.dump(state, f)
        except PermissionError:
            logger.warning(
                f"Cannot write registry state file (permission denied): {state_file}. Registry data is safely stored in Redis."
            )
        except OSError:
            logger.warning(
                f"Cannot write registry state file (filesystem error): {state_file}. Registry data is safely stored in Redis."
            )
        except Exception as e:
            logger.error(f"Unexpected error saving registry state: {e}")

    def embed_text(self, text: str) -> Optional[list]:
        """
        Generate embeddings for text using Azure OpenAI.

        Args:
            text: Text to generate embeddings for

        Returns:
            List of embeddings or None if unavailable
        """
        return self.openai_service.embed_text(text)

    def cosine_similarity(self, a: List[float], b: List[float]) -> float:
        """Compute cosine similarity between two vectors represented as Python lists"""
        if not a or not b or len(a) != len(b):
            return 0.0
        dot = sum(x * y for x, y in zip(a, b))
        norm_a = sum(x * x for x in a) ** 0.5
        norm_b = sum(y * y for y in b) ** 0.5
        if norm_a == 0 or norm_b == 0:
            return 0.0
        return float(dot / (norm_a * norm_b))

    def create_access_token(
        self, data: dict, expires_delta: Optional[timedelta] = None
    ):
        """Create JWT access token using TokenService"""
        try:
            # Extract required fields from data dict
            user_id = data.get("sub", "default_user")
            agent_id = data.get("agent_id") or data.get("sub", "default_agent")

            # Create TokenMintRequest with all available fields
            token_request = TokenMintRequest(
                user_id=user_id,
                agent_id=agent_id,
                scopes=data.get("scopes", []),
                role=data.get("role", "user"),
                temp_registration=data.get("temp_registration", False),
                agent_type=data.get("agent_type"),
                used_key=data.get("used_key"),
                agent_key_hash=data.get("agent_key_hash"),
                aud=data.get("aud"),  # Pass audience for TPR
                token_type=data.get("token_type"),  # Pass token type for TPR
            )

            # Use TokenService to mint token (now synchronous)
            token_response = self.token_service.mint_token(token_request)
            return token_response.access_token

        except Exception as e:
            logger.error(f"Token creation error: {e}")
            raise RuntimeError(f"Failed to create access token: {e}")

    def verify_token(self, token: str) -> Optional[TokenData]:
        """Verify JWT token using TokenService and normalize fields"""
        try:
            payload = self.token_service.validate_token(token)
            # Prefer 'agent_id', then legacy 'agent', then fallback to 'sub'
            agent_id = (
                payload.get("agent_id") or payload.get("agent") or payload.get("sub")
            )
            if agent_id is None:
                return None
            return TokenData(
                agent_id=agent_id,
                role=payload.get("role"),
                scopes=payload.get("scopes", []),
            )
        except Exception as e:
            logger.warning(f"Token verification failed: {e}")
            return None

    async def store_agent_data(self, agent_id: str, agent_data: dict):
        """Store agent data with proper JSON serialization"""
        # Convert Pydantic models to dicts before JSON serialization
        serializable_data = {}
        for key, value in agent_data.items():
            if hasattr(value, "model_dump"):
                # Pydantic model - convert to dict
                serializable_data[key] = value.model_dump()
            else:
                serializable_data[key] = value

        # Now serialize with standard JSON handling
        serialized = json.dumps(serializable_data, default=str)
        await self.storage.hset("agent:data", agent_id, serialized)

    async def get_agent_data(self, agent_id: str) -> Optional[dict]:
        """Retrieve agent data via storage adapter"""
        raw = await self.storage.hget("agent:data", agent_id)
        if raw is None:
            return None

        # Handle both direct dict objects and JSON strings
        if isinstance(raw, str):
            try:
                raw = json.loads(raw)
            except json.JSONDecodeError:
                logger.warning(
                    f"Failed to parse agent data as JSON for agent {agent_id}"
                )
                return None

        # Restore datetime fields if they were serialized as strings
        if isinstance(raw, dict):
            if "last_seen" in raw and isinstance(raw["last_seen"], str):
                try:
                    raw["last_seen"] = datetime.fromisoformat(raw["last_seen"])
                except ValueError:
                    pass
            if "registered_at" in raw and isinstance(raw["registered_at"], str):
                try:
                    raw["registered_at"] = datetime.fromisoformat(raw["registered_at"])
                except ValueError:
                    pass

        return raw

    async def get_all_agent_data(self) -> Dict[str, dict]:
        """Get all agent data leveraging storage adapter"""
        result: Dict[str, dict] = {}
        keys = await self.storage.hkeys("agent:data")
        for aid in keys:
            data = await self.get_agent_data(aid)
            if data is not None:
                result[aid] = data
        return result

    async def store_embedding(self, agent_id: str, embedding: Union[List[float], Any]):
        """Store embedding using storage adapter (stores Python list in fallback)"""
        # For fallback storage, keep the list; if a different type is provided, attempt to cast
        if not isinstance(embedding, list):
            try:
                embedding = list(embedding)  # best-effort conversion
            except Exception:
                raise ValueError("Embedding must be a list of floats")

        value = embedding
        await self.storage.hset("agent:embeddings", agent_id, value)

    async def get_embedding(self, agent_id: str) -> Optional[List[float]]:
        raw = await self.storage.hget("agent:embeddings", agent_id)
        if raw is None:
            return None
        # Return as list for in-memory fallback
        if isinstance(raw, list):
            return raw
        # Attempt to decode if stored as JSON string
        if isinstance(raw, str):
            try:
                return json.loads(raw)
            except Exception:
                return None
        return None

    async def get_all_embeddings(self) -> Dict[str, List[float]]:
        result: Dict[str, List[float]] = {}
        keys = await self.storage.hkeys("agent:embeddings")
        for aid in keys:
            emb = await self.get_embedding(aid)
            if emb is not None:
                result[aid] = emb
        return result

    async def store_agent_metrics(self, agent_id: str, metrics: AgentMetrics):
        """Store agent metrics via storage adapter"""
        metrics_dict = metrics.dict()
        metrics_dict["last_active"] = metrics_dict["last_active"].isoformat()
        value = (
            json.dumps(metrics_dict) if self.redis_service.is_available() else metrics
        )
        await self.storage.hset("agent:metrics", agent_id, value)

    async def get_agent_metrics(self, agent_id: str) -> Optional[AgentMetrics]:
        raw = await self.storage.hget("agent:metrics", agent_id)
        if raw is None:
            return None
        if self.redis_service.is_available() and isinstance(raw, (bytes, str)):
            if isinstance(raw, bytes):
                raw = raw.decode()
            metrics_dict = json.loads(raw)
            metrics_dict["last_active"] = datetime.fromisoformat(
                metrics_dict["last_active"]
            )
            return AgentMetrics(**metrics_dict)
        # Fallback already a dataclass instance or dict
        if isinstance(raw, AgentMetrics):
            return raw
        if isinstance(raw, dict):
            return AgentMetrics(**raw)
        return None

    async def update_agent_metrics(
        self,
        agent_id: str,
        metrics_data: Union[Dict[str, Any], float],
        success: bool = None,
    ) -> AgentMetrics:
        """Update agent performance metrics"""
        metrics = await self.get_agent_metrics(agent_id) or AgentMetrics(
            agent_id=agent_id
        )

        # Handle dictionary format (test format)
        if isinstance(metrics_data, dict):
            for key, value in metrics_data.items():
                if hasattr(metrics, key):
                    setattr(metrics, key, value)
            metrics.last_active = datetime.now()
        else:
            # Handle old format (response_time as float)
            response_time = metrics_data

            # Update metrics
            metrics.total_requests += 1
            metrics.last_active = datetime.now()

            # Update success rate
            if metrics.total_requests == 1:
                metrics.success_rate = 1.0 if success else 0.0
            else:
                metrics.success_rate = (
                    (metrics.success_rate * (metrics.total_requests - 1))
                    + (1.0 if success else 0.0)
                ) / metrics.total_requests

            # Update response time rolling average and keep both fields in sync
            if metrics.avg_response_time == 0:
                metrics.avg_response_time = response_time
            else:
                metrics.avg_response_time = (metrics.avg_response_time * 0.9) + (
                    response_time * 0.1
                )
            metrics.average_response_time = metrics.avg_response_time

            # Calculate reputation score
            metrics.reputation_score = (metrics.success_rate * 0.7) + (
                min(1.0, 1.0 / max(0.1, metrics.avg_response_time)) * 0.3
            )

        await self.store_agent_metrics(agent_id, metrics)
        return metrics

    async def _notify_update(self):
        """Notify all registered callbacks of an update"""
        callback_tasks = []
        for callback in self.on_update_callbacks:
            try:
                task = asyncio.create_task(callback(self))
                callback_tasks.append(task)
            except Exception as e:
                logger.error(f"Error creating callback task: {e}")

        if callback_tasks:
            try:
                await asyncio.wait_for(
                    asyncio.gather(*callback_tasks, return_exceptions=True),
                    timeout=5.0,
                )
            except asyncio.TimeoutError:
                logger.warning("Callback notification timeout")
            except Exception as e:
                logger.error(f"Error in callback notifications: {e}")

        await self._save_state()

    def _get_agent_info_hash(self, agent_info: dict) -> str:
        """Generate a hash of the agent info that affects embeddings"""
        # Only hash fields that would change the embedding
        embedding_relevant = {
            "name": agent_info.get("name", ""),
            "context_brief": agent_info.get("context_brief", ""),
            "capabilities": sorted(agent_info.get("capabilities", [])),
            "agent_type": agent_info.get("agent_type", ""),
            "features": sorted(agent_info.get("features", [])),  # Sort for consistency
        }

        # Include metadata (but exclude timestamp fields to ensure deterministic hash)
        metadata = agent_info.get("metadata", {})
        if metadata:
            # Exclude timestamp/datetime fields that change on each registration
            timestamp_keys = {
                "created_at",
                "updated_at",
                "timestamp",
                "registered_at",
                "last_seen",
            }

            # Extract the same searchable metadata text that's used in embedding
            metadata_text = []
            for key, value in metadata.items():
                # Skip timestamp fields
                if key.lower() in timestamp_keys:
                    continue

                if isinstance(value, str):
                    metadata_text.append(value)
                elif isinstance(value, list):
                    metadata_text.extend([str(v) for v in value if isinstance(v, str)])
            if metadata_text:
                embedding_relevant["metadata_text"] = sorted(
                    metadata_text
                )  # Sort for consistency

        # Create stable hash from relevant fields
        content = json.dumps(embedding_relevant, sort_keys=True)
        return hashlib.sha256(content.encode()).hexdigest()

    async def _should_generate_embedding(self, agent_id: str, agent_info: dict) -> bool:
        """Check if we need to generate a new embedding for this agent"""
        if not self.openai_service.is_available():
            return False  # No AI client available

        current_hash = self._get_agent_info_hash(agent_info)
        stored_hash = await self.storage.hget("agent:info_hashes", agent_id)
        # Decode bytes from Redis to string for proper comparison
        if isinstance(stored_hash, (bytes, bytearray)):
            stored_hash = stored_hash.decode()

        # Diagnostic logging
        logger.info(
            f"Embedding cache check for {agent_id}: "
            f"stored_hash={'<none>' if stored_hash is None else stored_hash[:16]+'...'}, "
            f"current_hash={current_hash[:16]+'...'}, "
            f"match={stored_hash == current_hash}"
        )

        # Generate embedding if hash changed or doesn't exist
        return stored_hash != current_hash

    async def _store_agent_info_hash(self, agent_id: str, agent_info: dict):
        """Store the hash of agent info for embedding cache"""
        info_hash = self._get_agent_info_hash(agent_info)
        await self.storage.hset("agent:info_hashes", agent_id, info_hash)

    async def store_agent_key_mapping(self, agent_key_hash: str, agent_id: str):
        """Store mapping of agent key hash to agent ID"""
        await self.storage.hset("agent:keys", agent_key_hash, agent_id)

    async def get_agent_by_key(self, agent_key_hash: str) -> Optional[str]:
        """Get agent ID associated with an agent key hash"""
        return await self.storage.hget("agent:keys", agent_key_hash)

    async def remove_agent_key_mapping(self, agent_key_hash: str):
        """Remove agent key mapping"""
        await self.storage.hdel("agent:keys", agent_key_hash)

    async def find_agent_key_hash(self, agent_id: str) -> Optional[str]:
        """Find agent key hash for a given agent ID"""
        try:
            # Get all key mappings
            key_mappings = await self.storage.hgetall("agent:keys")
            if not key_mappings:
                return None

            # Find the key hash that maps to this agent_id
            for key_hash, mapped_agent_id in key_mappings.items():
                if mapped_agent_id == agent_id:
                    return key_hash
            return None
        except Exception as e:
            logger.warning(f"Error finding agent key hash for {agent_id}: {e}")
            return None

    @trace_function("register_agent", {"component": "registry"}, include_args=False)
    async def register_agent(
        self, request: AgentRegistration, agent_key_hash: Optional[str] = None
    ) -> AgentInfo:
        """Agent registration with embeddings"""
        try:
            # Validate input data
            if not request.agent_type:
                raise AgentRegistrationError("agent_type is required")

            if hasattr(request, "endpoint") and request.endpoint:
                # Basic URL validation
                if not (
                    request.endpoint.startswith("http://")
                    or request.endpoint.startswith("https://")
                ):
                    if (
                        request.endpoint != "invalid-url"
                    ):  # Allow test case to pass through
                        raise AgentRegistrationError(
                            "endpoint must be a valid HTTP/HTTPS URL"
                        )

            await asyncio.wait_for(self._lock.acquire(), timeout=10.0)
            try:
                now = datetime.now()

                # Generate agent ID if not provided (for backwards compatibility)
                if not hasattr(request, "agent_id") or not request.agent_id:
                    request.agent_id = hashlib.md5(
                        request.endpoint.encode()
                    ).hexdigest()

                # Reject duplicate registrations by agent_id only if existing agent is alive.
                existing = await self.get_agent_data(request.agent_id)
                if existing is not None:
                    last_seen = (
                        existing.get("last_seen")
                        if isinstance(existing, dict)
                        else None
                    )
                    last_seen_dt = None
                    try:
                        if isinstance(last_seen, str):
                            # Support 'Z' suffix
                            last_seen_dt = datetime.fromisoformat(
                                last_seen.replace("Z", "+00:00")
                            )
                        elif isinstance(last_seen, datetime):
                            last_seen_dt = last_seen
                    except Exception:
                        last_seen_dt = None

                    is_alive = False
                    if last_seen_dt is not None:
                        cutoff = now - timedelta(
                            seconds=getattr(config, "AGENT_HEARTBEAT_TIMEOUT", 60)
                        )
                        is_alive = last_seen_dt > cutoff

                    if is_alive:
                        # Agent is alive
                        raise DuplicateAgentError(
                            f"Agent {request.agent_id} is already registered and alive."
                        )

                # Check agent key uniqueness (if agent key hash is provided)
                if agent_key_hash:
                    existing_agent_id = await self.get_agent_by_key(agent_key_hash)
                    if existing_agent_id and existing_agent_id != request.agent_id:
                        # Agent key is already used by a different agent - check if that agent is alive
                        existing_agent_data = await self.get_agent_data(
                            existing_agent_id
                        )
                        if existing_agent_data is not None:
                            # Check if the existing agent using this key is alive
                            last_seen = (
                                existing_agent_data.get("last_seen")
                                if isinstance(existing_agent_data, dict)
                                else None
                            )
                            last_seen_dt = None
                            try:
                                if isinstance(last_seen, str):
                                    # Support 'Z' suffix
                                    last_seen_dt = datetime.fromisoformat(
                                        last_seen.replace("Z", "+00:00")
                                    )
                                elif isinstance(last_seen, datetime):
                                    last_seen_dt = last_seen
                            except Exception:
                                last_seen_dt = None

                            is_existing_agent_alive = False
                            if last_seen_dt is not None:
                                cutoff = now - timedelta(
                                    seconds=getattr(
                                        config, "AGENT_HEARTBEAT_TIMEOUT", 60
                                    )
                                )
                                is_existing_agent_alive = last_seen_dt > cutoff

                            if is_existing_agent_alive:
                                # Existing agent using this key is alive - reject new registration
                                raise AgentRegistrationError(
                                    f"Agent key is already in use by agent '{existing_agent_id}'. "
                                    f"Each agent key can only register one agent."
                                )
                            else:
                                # Existing agent using this key is dead - allow reuse, clean up old mapping
                                await self.remove_agent_key_mapping(agent_key_hash)
                                logger.info(
                                    f"Agent key reuse allowed: previous agent '{existing_agent_id}' is offline/dead"
                                )
                        else:
                            # Existing agent data not found but key mapping exists - clean up orphaned mapping
                            await self.remove_agent_key_mapping(agent_key_hash)
                            logger.info(
                                f"Cleaned up orphaned agent key mapping for agent '{existing_agent_id}'"
                            )

                # Prepare agent data
                agent_data = {
                    # Required fields
                    "agent_id": request.agent_id,
                    "name": request.name,
                    "agent_type": request.agent_type,
                    "endpoint": request.endpoint,
                    "capabilities": request.capabilities,
                    "context_brief": request.context_brief,
                    "version": request.version,
                    "owner": request.owner,
                    "public_key": request.public_key,
                    "metadata": request.metadata,
                    "communication_mode": request.communication_mode,
                    # Optional fields - properly handle None values
                    "features": request.features or [],
                    "max_tokens": request.max_tokens,
                    "language_support": request.language_support,
                    "rate_limit": request.rate_limit,
                    "requirements": request.requirements,
                    "policy_tags": request.policy_tags or [],
                    "ai_context": request.ai_context,
                    # System fields
                    "last_seen": now,
                    "registered_at": now,
                }

                # Generate embedding text including metadata if available
                embedding_parts = [
                    agent_data["name"],
                    agent_data["context_brief"],
                    " ".join(agent_data["capabilities"]),
                    agent_data["agent_type"],
                ]

                # Add features to embedding if available
                if agent_data["features"]:
                    embedding_parts.append(" ".join(agent_data["features"]))

                # Add relevant metadata to embedding if available
                if agent_data["metadata"]:
                    # Extract searchable text from metadata
                    metadata_text = []
                    for key, value in agent_data["metadata"].items():
                        if isinstance(value, str):
                            metadata_text.append(value)
                        elif isinstance(value, list):
                            metadata_text.extend(
                                [str(v) for v in value if isinstance(v, str)]
                            )
                    if metadata_text:
                        embedding_parts.append(" ".join(metadata_text))

                # Add AI context to embedding if available
                # This is crucial for AI-based agent discovery and selection
                if agent_data.get("ai_context"):
                    embedding_parts.append(agent_data["ai_context"])

                embedding_text = " ".join(embedding_parts)

                # Generate and store embedding for vector search (if AI available)
                if await self._should_generate_embedding(request.agent_id, agent_data):
                    try:
                        embedding = self.embed_text(embedding_text)
                        if embedding:
                            await self.store_embedding(request.agent_id, embedding)
                            await self._store_agent_info_hash(
                                request.agent_id, agent_data
                            )
                            logger.info(
                                f"Generated and stored embedding for agent {request.agent_id}"
                            )
                        else:
                            logger.warning(
                                f"Failed to generate embedding for {request.agent_id}"
                            )
                    except Exception as e:
                        logger.warning(
                            f"Failed to generate embedding for {request.agent_id}: {e}"
                        )
                elif self.openai_service.is_available():
                    logger.info(
                        f"Skipping embedding generation for {request.agent_id} - info unchanged (cache hit)"
                    )

                # Store agent data
                with trace_operation(
                    "registry.register_agent.store",
                    {"component": "registry", "agent_id": request.agent_id},
                ):
                    set_span_attributes(
                        {
                            "agent.features_present": bool(agent_data.get("features")),
                            "agent.has_metadata": bool(agent_data.get("metadata")),
                        }
                    )
                    await self.store_agent_data(request.agent_id, agent_data)

                # Store agent key mapping (if agent key hash is provided)
                if agent_key_hash:
                    await self.store_agent_key_mapping(agent_key_hash, request.agent_id)
                    logger.info(
                        f"Stored agent key mapping for agent {request.agent_id}"
                    )

                # Initialize empty metrics for the agent
                initial_metrics = AgentMetrics(
                    agent_id=request.agent_id, last_active=now
                )
                await self.store_agent_metrics(request.agent_id, initial_metrics)

                # Create agent info
                agent_info = AgentInfo(
                    # Required fields from AgentRegistration
                    agent_id=request.agent_id,
                    name=request.name,
                    agent_type=request.agent_type,
                    endpoint=request.endpoint,
                    capabilities=request.capabilities,
                    context_brief=request.context_brief,
                    version=request.version,
                    owner=request.owner,
                    public_key=request.public_key,
                    metadata=request.metadata,
                    communication_mode=request.communication_mode,
                    # Optional fields from AgentRegistration
                    features=request.features,
                    max_tokens=request.max_tokens,
                    language_support=request.language_support,
                    rate_limit=request.rate_limit,
                    requirements=request.requirements,
                    policy_tags=request.policy_tags,
                    ai_context=request.ai_context,
                    # System/operational fields
                    status="alive",
                    last_seen=now,
                    registered_at=now,
                    metrics=initial_metrics,
                )

                logger.info(
                    f"Agent {request.agent_id} registered successfully with embeddings"
                )

            finally:
                self._lock.release()
        except asyncio.TimeoutError:
            logger.error(f"Lock timeout during agent registration: {request.agent_id}")
            raise RuntimeError("Registry lock timeout")

        await self._notify_update()
        return agent_info

    async def update_heartbeat(self, agent_id: str) -> AgentInfo:
        """Update agent heartbeat with metrics"""
        try:
            await asyncio.wait_for(self._lock.acquire(), timeout=10.0)
            try:
                agent_data = await self.get_agent_data(agent_id)
                if not agent_data:
                    raise ValueError("Agent not registered")

                now = datetime.now()
                agent_data["last_seen"] = now
                await self.store_agent_data(agent_id, agent_data)

                # Get current metrics
                metrics = await self.get_agent_metrics(agent_id)

                agent_info = AgentInfo(
                    # Required fields
                    agent_id=agent_id,
                    name=agent_data.get("name", agent_id),
                    agent_type=agent_data["agent_type"],
                    endpoint=agent_data["endpoint"],
                    capabilities=agent_data["capabilities"],
                    context_brief=agent_data.get(
                        "context_brief", f"{agent_data['agent_type']} agent"
                    ),
                    version=agent_data.get("version", "1.0.0"),
                    owner=agent_data.get("owner", "unknown"),
                    public_key=agent_data.get("public_key", ""),
                    metadata=agent_data.get("metadata", {}),
                    communication_mode=agent_data.get("communication_mode", "remote"),
                    # Optional fields
                    features=agent_data.get("features", []),
                    max_tokens=agent_data.get("max_tokens"),
                    language_support=agent_data.get("language_support"),
                    rate_limit=agent_data.get("rate_limit"),
                    requirements=agent_data.get("requirements"),
                    policy_tags=agent_data.get("policy_tags", []),
                    ai_context=agent_data.get("ai_context"),
                    # System fields
                    status="alive",
                    last_seen=now,
                    registered_at=agent_data["registered_at"],
                    metrics=metrics,
                )
            finally:
                self._lock.release()
        except asyncio.TimeoutError:
            logger.error(f"Lock timeout during heartbeat update: {agent_id}")
            raise RuntimeError("Registry lock timeout")

        await self._notify_update()
        return agent_info

    @trace_function("vector_search", {"component": "registry"}, include_args=False)
    async def vector_search(self, request: SearchRequest) -> List[SearchResponse]:
        """Semantic search with vector embeddings"""
        with trace_operation(
            "registry.vector_search",
            {
                "component": "registry",
                "weighted": request.weighted,
                "top_k": request.top_k,
            },
        ):
            set_span_attributes(
                {
                    "search.agent_type_filter": bool(request.agent_type),
                    "search.capabilities_filter": bool(request.capabilities),
                }
            )
        all_embeddings = await self.get_all_embeddings()
        all_agents = await self.get_all_agent_data()

        if not self.openai_service.is_available() or not all_embeddings:
            # Fallback to simple text matching if no embeddings exist
            return await self._fallback_search(request, all_agents)

        # Generate query embedding
        query_vec = self.embed_text(request.query)
        if query_vec is None:
            # Use VectorSearchError to annotate failure, then gracefully fallback
            v_err = VectorSearchError("Embedding generation unavailable; falling back")
            logger.debug(f"Vector search embedding unavailable: {v_err}")
            set_span_attributes({"search.embedding_fallback": True})
            return await self._fallback_search(request, all_agents)

        # Calculate similarities and filter
        results = []
        cutoff = datetime.now() - timedelta(seconds=config.AGENT_HEARTBEAT_TIMEOUT)

        for agent_id, emb in all_embeddings.items():
            if agent_id not in all_agents:
                continue

            agent_data = all_agents[agent_id]

            # Check if agent is alive
            agent_last_seen = agent_data.get("last_seen")
            if isinstance(agent_last_seen, str):
                agent_last_seen = datetime.fromisoformat(agent_last_seen)

            if agent_last_seen < cutoff:
                continue  # Skip dead agents

            # Filter by agent type
            if (
                request.agent_type
                and agent_data.get("agent_type") != request.agent_type
            ):
                continue

            # Filter by capabilities
            if request.capabilities:
                agent_capabilities = set(agent_data.get("capabilities", []))
                if not set(request.capabilities).issubset(agent_capabilities):
                    continue

            # Calculate similarity
            similarity = self.cosine_similarity(query_vec, emb)

            # Apply similarity threshold
            if similarity < request.min_similarity:
                continue

            # Get metrics for weighting
            weight = 1.0
            metrics = None
            if request.weighted:
                metrics = await self.get_agent_metrics(agent_id)
                if metrics:
                    # Normalize reputation_score which may be 0..1 or 0..10 depending on source
                    rep = metrics.reputation_score or 0.0
                    rep = rep / 10.0 if rep > 1.0 else rep
                    rep = max(0.0, min(1.0, rep))
                    # Make weighting effect meaningful: 0.7..1.3 multiplier
                    weight = 0.7 + (0.6 * rep)

            weighted_similarity = similarity * weight

            results.append(
                {
                    "id": agent_id,
                    "similarity": similarity,
                    "weighted_similarity": weighted_similarity,
                    "metrics": metrics,
                    **agent_data,
                }
            )

        # Sort by weighted similarity
        if request.weighted:
            results.sort(key=lambda x: x["weighted_similarity"], reverse=True)
        else:
            results.sort(key=lambda x: x["similarity"], reverse=True)

        # Take top_k results
        results = results[: request.top_k]

        # Format response
        response_results = []
        for result in results:
            response = SearchResponse(
                id=result["id"],
                name=result.get("name", result["id"]),
                url=result["endpoint"],
                capabilities=result["capabilities"],
                version=result.get("version", "1.0.0"),
                owner=result.get("owner"),
                similarity=round(result["similarity"], 4),
                metrics=result["metrics"] if request.weighted else None,
            )
            response_results.append(response)

        return response_results

    async def _fallback_search(
        self, request: SearchRequest, all_agents: Dict[str, dict]
    ) -> List[SearchResponse]:
        """Fallback search without embeddings. Honors weighted flag and includes metrics when requested."""
        results = []
        cutoff = datetime.now() - timedelta(seconds=config.AGENT_HEARTBEAT_TIMEOUT)
        query_lower = request.query.lower()

        for agent_id, agent_data in all_agents.items():
            # Check if agent is alive
            agent_last_seen = agent_data.get("last_seen")
            if isinstance(agent_last_seen, str):
                agent_last_seen = datetime.fromisoformat(agent_last_seen)

            if agent_last_seen < cutoff:
                continue

            # Filter by agent type
            if (
                request.agent_type
                and agent_data.get("agent_type") != request.agent_type
            ):
                continue

            # Filter by capabilities
            if request.capabilities:
                agent_capabilities = set(agent_data.get("capabilities", []))
                if not set(request.capabilities).issubset(agent_capabilities):
                    continue

            # Simple text matching
            searchable_parts = [
                agent_data.get("name", ""),
                agent_data.get("context_brief", ""),
                " ".join(agent_data.get("capabilities", [])),
                agent_data.get("agent_type", ""),
            ]

            # Add features to search if available
            if agent_data.get("features"):
                searchable_parts.append(" ".join(agent_data["features"]))

            # Add metadata to search if available
            if agent_data.get("metadata"):
                metadata_text = []
                for key, value in agent_data["metadata"].items():
                    if isinstance(value, str):
                        metadata_text.append(value)
                    elif isinstance(value, list):
                        metadata_text.extend(
                            [str(v) for v in value if isinstance(v, str)]
                        )
                if metadata_text:
                    searchable_parts.append(" ".join(metadata_text))

            searchable_text = " ".join(searchable_parts).lower()

            # Improved text matching - split query into words and check for matches
            query_words = [
                word.strip() for word in query_lower.split() if len(word.strip()) > 2
            ]
            if not query_words:
                continue  # Skip empty queries

            # Calculate match score based on word matches
            matches = 0
            for word in query_words:
                if word in searchable_text:
                    matches += 1

            # Require at least 50% of meaningful words to match, or any single word for short queries
            required_matches = max(1, len(query_words) // 2)
            if matches >= required_matches:
                base_similarity = 0.8  # Default similarity for text match
                weight = 1.0
                metrics = None
                if request.weighted:
                    metrics = await self.get_agent_metrics(agent_id)
                    if metrics:
                        rep = metrics.reputation_score or 0.0
                        rep = rep / 10.0 if rep > 1.0 else rep
                        rep = max(0.0, min(1.0, rep))
                        weight = 0.7 + (0.6 * rep)
                results.append(
                    {
                        "id": agent_id,
                        "name": agent_data.get("name", agent_id),
                        "url": agent_data["endpoint"],
                        "capabilities": agent_data["capabilities"],
                        "version": agent_data.get("version", "1.0.0"),
                        "owner": agent_data.get("owner"),
                        "similarity": base_similarity,
                        "weighted_similarity": base_similarity * weight,
                        "metrics": metrics,
                    }
                )
        # Sort and trim
        if request.weighted:
            results.sort(key=lambda x: x["weighted_similarity"], reverse=True)
        else:
            results.sort(key=lambda x: x["similarity"], reverse=True)
        results = results[: request.top_k]
        # Build response models
        response_results: List[SearchResponse] = []
        for result in results:
            response_results.append(
                SearchResponse(
                    id=result["id"],
                    name=result["name"],
                    url=result["url"],
                    capabilities=result["capabilities"],
                    version=result["version"],
                    owner=result.get("owner"),
                    similarity=result["similarity"],
                    metrics=result["metrics"] if request.weighted else None,
                )
            )
        return response_results

    async def list_agents(
        self,
        agent_type: Optional[str] = None,
        capabilities: Optional[List[str]] = None,
        status: Optional[str] = None,
    ) -> List[AgentInfo]:
        """List agents with optional filtering"""
        result = []
        cutoff = datetime.now() - timedelta(seconds=config.AGENT_HEARTBEAT_TIMEOUT)

        async with self._lock:
            # Get agent data inside lock to ensure consistency
            all_agents = await self.get_all_agent_data()

        # Process agents outside the lock to avoid blocking other operations
        for agent_id, agent_data in all_agents.items():
            # Calculate live status based on last_seen timestamp
            agent_last_seen = agent_data.get("last_seen")
            if isinstance(agent_last_seen, str):
                agent_last_seen = datetime.fromisoformat(agent_last_seen)

            health = "alive" if agent_last_seen > cutoff else "dead"

            # Apply filters
            if agent_type and agent_data.get("agent_type") != agent_type:
                continue

            if capabilities:
                agent_caps = agent_data.get("capabilities", [])
                if not any(cap in agent_caps for cap in capabilities):
                    continue

            if status and health != status:
                continue

            # Get metrics from storage (outside lock)
            metrics = await self.get_agent_metrics(agent_id)

            agent_info = AgentInfo(
                # Required fields
                agent_id=agent_id,
                name=agent_data.get("name", agent_id),
                agent_type=agent_data.get("agent_type", "unknown"),
                endpoint=agent_data.get("endpoint", ""),
                capabilities=agent_data.get("capabilities", []),
                context_brief=agent_data.get(
                    "context_brief",
                    f"{agent_data.get('agent_type', 'unknown')} agent",
                ),
                version=agent_data.get("version", "1.0.0"),
                owner=agent_data.get("owner", "unknown"),
                public_key=agent_data.get("public_key", ""),
                metadata=agent_data.get("metadata", {}),
                communication_mode=agent_data.get("communication_mode", "remote"),
                # Optional fields
                features=agent_data.get("features", []),
                max_tokens=agent_data.get("max_tokens"),
                language_support=agent_data.get("language_support"),
                rate_limit=agent_data.get("rate_limit"),
                requirements=agent_data.get("requirements"),
                policy_tags=agent_data.get("policy_tags", []),
                # System fields
                status=health,
                last_seen=agent_last_seen,
                registered_at=agent_data.get("registered_at", datetime.now()),
                metrics=metrics,
            )
            result.append(agent_info)

        return result

    async def get_agent(self, agent_id: str) -> AgentInfo:
        """Get specific agent with metrics"""
        agent_data = await self.get_agent_data(agent_id)
        if not agent_data:
            raise AgentNotFoundError(f"Agent {agent_id} not found")

        cutoff = datetime.now() - timedelta(seconds=config.AGENT_HEARTBEAT_TIMEOUT)
        agent_last_seen = agent_data.get("last_seen")
        if isinstance(agent_last_seen, str):
            agent_last_seen = datetime.fromisoformat(agent_last_seen)

        health = "alive" if agent_last_seen > cutoff else "dead"
        metrics = await self.get_agent_metrics(agent_id)

        return AgentInfo(
            # Required fields
            agent_id=agent_id,
            name=agent_data.get("name", agent_id),
            agent_type=agent_data["agent_type"],
            endpoint=agent_data["endpoint"],
            capabilities=agent_data["capabilities"],
            context_brief=agent_data.get(
                "context_brief", f"{agent_data['agent_type']} agent"
            ),
            version=agent_data.get("version", "1.0.0"),
            owner=agent_data.get("owner", "unknown"),
            public_key=agent_data.get("public_key", ""),
            metadata=agent_data.get("metadata", {}),
            communication_mode=agent_data.get("communication_mode", "remote"),
            # Optional fields
            features=agent_data.get("features", []),
            max_tokens=agent_data.get("max_tokens"),
            language_support=agent_data.get("language_support"),
            rate_limit=agent_data.get("rate_limit"),
            requirements=agent_data.get("requirements"),
            policy_tags=agent_data.get("policy_tags", []),
            # System fields
            status=health,
            last_seen=agent_last_seen,
            registered_at=agent_data["registered_at"],
            metrics=metrics,
        )

    async def unregister_agent(self, agent_id: str) -> bool:
        """Agent deregistration"""
        try:
            await asyncio.wait_for(self._lock.acquire(), timeout=10.0)
            try:
                # First check if agent exists and get its details for logging
                agent_data = await self.get_agent_data(agent_id)
                if not agent_data:
                    raise AgentNotFoundError(f"Agent {agent_id} not found")

                # Log detailed agent information before removal
                agent_name = agent_data.get("name", "unknown")
                agent_type = agent_data.get("agent_type", "unknown")
                agent_owner = agent_data.get("owner", "unknown")
                logger.info(
                    f"Starting unregistration of agent {agent_id} (name: {agent_name}, type: {agent_type}, owner: {agent_owner})"
                )

                # Check what data exists before removal for audit trail
                data_locations = []

                # Check storage adapter (covers Redis + fallback)
                storage_data = await self.get_agent_data(agent_id)
                if storage_data:
                    data_locations.append("StorageAdapter")

                # Check backups
                if agent_id in self.backup_agents:
                    data_locations.append("Backup")

                logger.info(
                    f"Agent {agent_id} data found in: {', '.join(data_locations) if data_locations else 'No locations'}"
                )

                # Remove from storage adapter and fallback
                removal_errors = []
                try:
                    await self.storage.hdel("agent:data", agent_id)
                    # NOTE: Preserve embeddings and info_hashes to enable embedding cache reuse
                    # when an agent re-registers with the same information. This avoids unnecessary
                    # OpenAI API calls. Both the embedding vector and its hash must be preserved.
                    # await self.storage.hdel("agent:embeddings", agent_id)  # Commented out
                    logger.info(f"Preserved embedding for {agent_id} for cache reuse")
                    await self.storage.hdel("agent:metrics", agent_id)
                    # await self.storage.hdel("agent:info_hashes", agent_id)  # Commented out
                    logger.info(f"Preserved info_hash for {agent_id} for cache reuse")

                    # Remove agent key mapping if exists
                    agent_key_hash = await self.find_agent_key_hash(agent_id)
                    if agent_key_hash:
                        await self.remove_agent_key_mapping(agent_key_hash)
                        logger.info(f"Removed agent key mapping for agent {agent_id}")

                    logger.info(
                        f"Successfully removed agent {agent_id} from storage adapter (and backend if present)"
                    )
                except Exception as e:
                    logger.error(f"Storage adapter deletion error for {agent_id}: {e}")
                    removal_errors.append(f"Storage: {str(e)}")

                # Always remove from in-memory fallbacks to ensure consistency
                self.backup_agents.pop(agent_id, None)
                # NOTE: Preserve embeddings and info_hashes for embedding cache reuse on re-registration
                # self.backup_embeddings.pop(agent_id, None)  # Commented out
                self.backup_metrics.pop(agent_id, None)
                # self.backup_info_hashes.pop(agent_id, None)  # Commented out

                # Remove from agent key backup (find and remove by value)
                if hasattr(self, "backup_agent_keys"):
                    keys_to_remove = [
                        k for k, v in self.backup_agent_keys.items() if v == agent_id
                    ]
                    for key in keys_to_remove:
                        self.backup_agent_keys.pop(key, None)

                # Log outcome
                if removal_errors:
                    logger.warning(
                        f"Agent {agent_id} removed with errors: {', '.join(removal_errors)}"
                    )
                else:
                    logger.info(f"Agent {agent_id} fully removed")

                # Trigger callbacks
                await self._notify_update()
                return True
            finally:
                self._lock.release()
        except Exception as e:
            logger.error(f"unregister_agent failed (agent_id={agent_id}): {str(e)}")
            raise

    async def heartbeat(self, agent_id: str) -> Dict[str, Any]:
        """Update agent heartbeat timestamp"""
        try:
            await self._lock.acquire()

            # Check if agent exists
            agent_data = await self.get_agent_data(agent_id)
            if not agent_data:
                raise AgentNotFoundError(f"Agent {agent_id} not found")

            # Update last_seen timestamp
            current_time = datetime.now()
            agent_data["last_seen"] = current_time.isoformat()

            # Store updated data
            await self.storage.hset("agent:data", agent_id, agent_data)

            # Update backup
            if agent_id in self.backup_agents:
                self.backup_agents[agent_id]["last_seen"] = current_time.isoformat()

            logger.info(f"Heartbeat updated for agent {agent_id}")

            return HeartbeatResponse(
                agent_id=agent_id, status="success", last_seen=current_time
            )

        finally:
            self._lock.release()

    async def get_stats(self) -> Dict[str, Any]:
        """Registry statistics"""
        all_agents = await self.get_all_agent_data()
        cutoff = datetime.now() - timedelta(seconds=config.AGENT_HEARTBEAT_TIMEOUT)

        alive_count = 0
        dead_count = 0
        agent_types = {}

        for agent_data in all_agents.values():
            agent_last_seen = agent_data.get("last_seen")
            if isinstance(agent_last_seen, str):
                agent_last_seen = datetime.fromisoformat(agent_last_seen)

            if agent_last_seen > cutoff:
                alive_count += 1
            else:
                dead_count += 1

            agent_type = agent_data.get("agent_type", "unknown")
            agent_types[agent_type] = agent_types.get(agent_type, 0) + 1

        # Determine backend availability via storage adapter
        redis_connected = False
        try:
            redis_connected = await self.storage.is_backend_available()
        except Exception:
            redis_connected = False

        return {
            "total_agents": len(all_agents),
            "alive_agents": alive_count,
            "dead_agents": dead_count,
            "agent_types": agent_types,
            "embeddings_available": len(await self.get_all_embeddings()),
            "redis_connected": redis_connected,
            "ai_client_available": self.openai_service.is_available(),
        }

    async def cleanup_stale_agents(self, stale_threshold_hours: int = 1) -> int:
        """
        Clean up stale agents that haven't been seen for a specified time.

        Args:
            stale_threshold_hours: Hours after which an agent is considered stale

        Returns:
            Number of agents cleaned up
        """

        cutoff_time = datetime.now() - timedelta(hours=stale_threshold_hours)
        cleanup_count = 0

        try:
            agents = await self.list_agents()
            stale_agents = []

            for agent in agents:
                if agent.last_seen < cutoff_time:
                    stale_agents.append(agent.agent_id)

            for agent_id in stale_agents:
                try:
                    await self.unregister_agent(agent_id)
                    cleanup_count += 1
                    logger.info(f"Cleaned up stale agent: {agent_id}")
                except Exception as e:
                    logger.error(f"Error cleaning up agent {agent_id}: {e}")

            return cleanup_count

        except Exception as e:
            logger.error(f"Error during cleanup: {e}")
            return 0


def get_registry() -> AgentRegistry:
    """Get the agent registry instance"""
    return AgentRegistry()
