# tool_registry_client.py
import json
from typing import Any, Dict, List, Optional

import backoff
import httpx
from guardianhub import get_logger

from guardianhub.config.settings import settings

logger = get_logger(__name__)


class ToolRegistryClient:
    """HTTP client for the tool registry service with retry logic."""

    def __init__(self):
        self._base_url = settings.endpoints.TOOL_REGISTRY_URL
        self._timeout = settings.endpoints.get("TOOL_REGISTRY_TIMEOUT", 10)
        self._service_port = 8002  # Default port for tool registry

    async def _request(self, method: str, path: str, token: str = None, **kwargs) -> Any:
        """Base request method with retry logic and proper authorization."""
        url = f"{self._base_url.rstrip('/')}{path}"

        # Set up headers - read baggage within the current context
        from opentelemetry import baggage, propagate
        from opentelemetry import context as otel_context
        
        # Get current context and read baggage from it
        current_ctx = otel_context.get_current()
        tenant_id = baggage.get_baggage("tenant_id", current_ctx)
        session_id = baggage.get_baggage("session_id", current_ctx)
        user_id = baggage.get_baggage("user_id", current_ctx)
        access_token = baggage.get_baggage("access_token", current_ctx)
        
        # Use provided token or fallback to baggage token
        token = token or access_token
        
        logger.info(f"Tool Registry Client: tenant_id={tenant_id}, session_id={session_id}, user_id={user_id}, token={token}")
        
        headers = kwargs.get('headers', {})
        if token:
            headers['Authorization'] = f'{token}'
        
        if tenant_id:
            logger.info("Adding X-Tenant-ID header")
            headers['X-Tenant-ID'] = tenant_id

        # Inject trace context and baggage into headers for propagation
        propagate.inject(headers)
        
        kwargs['headers'] = headers

        # Ensure the request runs in the current context
        async with httpx.AsyncClient(timeout=self._timeout) as client:
            response = await client.request(method, url, **kwargs)
            response.raise_for_status()
            return response.json()

    async def load_agentic_capabilities(self) -> Dict[str, Any]:
        """Get all agent specifications from registry."""
        # First get all agent names (This is actually returning a List[Dict])
        agents_data = await self._request("GET", "/agent_registry/agents/")

        # Then fetch specs for each agent
        agent_specs = {}
        for agent in agents_data:
            # 🎯 Extract the name string from the dictionary
            agent_name = agent.get('name')

            if not agent_name:
                logger.warning(f"⚠️ Found agent entry without a name: {agent}")
                continue

            spec = await self._request("GET", f"/agent_registry/agents/{agent_name}/spec")
            agent_specs[agent_name] = spec

        return agent_specs


    async def get_agent_spec(self, agent_name: str) -> Dict[str, Any]:
        """Get agent specification by name."""
        return await self._request("GET", f"/agent_registry/agents/{agent_name}/spec")

    async def search_tools(
            self,
            query: str,
            agent_id: str,  # 🚀 MANDATORY for security wall
            intent: str,  # 🚀 MANDATORY for context filtering
            n_results: int = 5,
            collection: str = "sutram_agentic_tools"
    ) -> List[Dict[str, Any]]:
        """
        Performs Secure Semantic Discovery.
        """
        payload = {
            "query": query,
            "agent_id": agent_id,
            "intent": intent,
            "n_results": n_results,
            "collection": collection
        }

        # Returns the normalized List[Dict] we perfected in the VectorServiceClient
        response = await self._request("POST", "/data_registry/search/vector", json=payload)
        return response.get("results", [])

    async def list_tools(self) -> List[Dict[str, Any]]:
        """List all available tools."""
        return await self._request("GET", "/tool_registry/api/tools")
