#!/usr/bin/env python3
"""
Agent Management Script for MCP Gateway Registry.

This tool provides CLI commands for managing A2A agents via the A2A Agent Management API.
It uses JWT Bearer tokens from Keycloak for authentication via the mcp-gateway-m2m service account.

SERVICE ACCOUNT: mcp-gateway-m2m
The mcp-gateway-m2m service account is a Keycloak M2M client that provides authentication
for both MCP server management and A2A agent management operations. The JWT token from this
account is automatically loaded from .oauth-tokens/ingress.json.

PERMISSIONS:
- Token scopes: mcp-servers-restricted/*, mcp-servers-unrestricted/*, a2a-agent-admin
- Agent operations: register, modify, delete, list (full admin access)
- Group assignment: mcp-servers-unrestricted, a2a-agent-admin

API: /api/agents (A2A Agent Management API - dedicated endpoints for agent management)
- List agents: GET /api/agents
- Get agent: GET /api/agents/{path}
- Register agent: POST /api/agents/register
- Update agent: PUT /api/agents/{path}
- Delete agent: DELETE /api/agents/{path}
- Toggle agent: POST /api/agents/{path}/toggle
- Search agents: POST /api/agents/discover/semantic

HEALTH CHECKS:
- The 'test' command performs two-level verification:
  1. Registry Check: Verifies agent metadata in the registry (always performed)
  2. Service Health: Fetches agent card from /.well-known/agent-card.json (if agent is enabled)
- Results show: PASSED (agent responds), FAILED (agent unavailable), SKIPPED (disabled/no URL)

SEMANTIC SEARCH:
- The 'search' command performs natural language semantic search using FAISS vector index
- Returns enabled agents matching the query with relevance scores

Usage:
    # Automatically loads token from .oauth-tokens/ingress.json (generated by credentials-provider/generate_creds.sh)
    uv run python cli/agent_mgmt.py list
    uv run python cli/agent_mgmt.py get /code-reviewer
    uv run python cli/agent_mgmt.py test /code-reviewer          # Test single agent with health check
    uv run python cli/agent_mgmt.py test-all                      # Test all agents with health checks
    uv run python cli/agent_mgmt.py search "code review agent"    # Semantic search for agents

For agent creation, registration, toggle, and delete operations, the mcp-gateway-m2m
service account must be assigned to the a2a-agent-admin group in Keycloak.
"""

import argparse
import base64
import json
import logging
import os
import subprocess  # nosec B404
import sys
import time
from typing import Any

import requests

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s,p%(process)s,{%(filename)s:%(lineno)d},%(levelname)s,%(message)s",
)
logger = logging.getLogger(__name__)

DEFAULT_BASE_URL: str = "http://localhost"  # Goes through nginx (port 80), not direct :7860
DEFAULT_TOKEN_FILE: str = ".oauth-tokens/ingress.json"
REQUEST_TIMEOUT: int = 10
API_BASE: str = "/api/agents"


def _extract_username_from_jwt(token: str) -> str:
    """Extract username from JWT token payload."""
    try:
        parts = token.split(".")
        if len(parts) != 3:
            return "unknown"

        payload = parts[1]
        padding = 4 - (len(payload) % 4)
        if padding != 4:
            payload += "=" * padding

        decoded = base64.urlsafe_b64decode(payload)
        claims = json.loads(decoded)

        username = claims.get("preferred_username") or claims.get("sub") or "unknown"
        return username
    except Exception:
        return "unknown"


def _get_token_expiration(token: str) -> int | None:
    """Extract expiration timestamp from JWT token.

    Returns:
        Expiration timestamp (seconds since epoch) or None if unable to extract
    """
    try:
        parts = token.split(".")
        if len(parts) != 3:
            return None

        payload = parts[1]
        padding = 4 - (len(payload) % 4)
        if padding != 4:
            payload += "=" * padding

        decoded = base64.urlsafe_b64decode(payload)
        claims = json.loads(decoded)
        return claims.get("exp")
    except Exception:
        return None


def _is_token_expired(token: str, buffer_seconds: int = 30) -> bool:
    """Check if token is expired or about to expire.

    Args:
        token: JWT token string
        buffer_seconds: Seconds before actual expiration to consider token expired

    Returns:
        True if token is expired or expiring soon, False otherwise
    """
    exp_timestamp = _get_token_expiration(token)
    if exp_timestamp is None:
        return False

    current_time = time.time()
    return current_time >= (exp_timestamp - buffer_seconds)


def _regenerate_token(token_file: str) -> bool:
    """Regenerate token using generate_creds.sh script.

    Args:
        token_file: Path to the token file

    Returns:
        True if regeneration succeeded, False otherwise
    """
    logger.info("Token expired, regenerating credentials...")

    # Extract the agent name from token file if it's a bot account
    # e.g., .oauth-tokens/bot-x-token.json -> bot-x
    token_filename = os.path.basename(token_file)
    if token_filename.endswith("-token.json"):
        agent_name = token_filename[:-11]  # Remove "-token.json"
    else:
        # Fallback to running generate_creds.sh for main ingress token
        agent_name = None

    try:
        script_dir = os.path.dirname(os.path.abspath(__file__))
        project_root = os.path.dirname(script_dir)

        if agent_name:
            # Use generate-agent-token.sh for specific agent
            # Call from keycloak/setup directory so relative paths work
            token_script = os.path.join(project_root, "keycloak/setup/generate-agent-token.sh")
            keycloak_setup_dir = os.path.join(project_root, "keycloak/setup")
            logger.info(f"Running: {token_script} {agent_name}")
            result = subprocess.run(  # nosec B603 - hardcoded internal script path
                [token_script, agent_name],
                cwd=keycloak_setup_dir,  # Run from keycloak/setup so ../../.oauth-tokens works
                capture_output=True,
                text=True,
                timeout=30,
            )
        else:
            # Use generate_creds.sh for ingress token
            creds_script = os.path.join(project_root, "credentials-provider/generate_creds.sh")
            logger.info(f"Running: {creds_script} --ingress-only")
            result = subprocess.run(  # nosec B603 - hardcoded internal script path and flags
                [creds_script, "--ingress-only"],
                cwd=project_root,
                capture_output=True,
                text=True,
                timeout=60,
            )

        if result.returncode == 0:
            logger.info("✓ Token regenerated successfully")
            return True
        else:
            logger.error("✗ Token regeneration failed")
            logger.error(f"  Error output: {result.stderr}")
            return False

    except FileNotFoundError as e:
        logger.error(f"✗ Token regeneration script not found: {e}")
        return False
    except subprocess.TimeoutExpired:
        logger.error("✗ Token regeneration script timed out")
        return False
    except Exception as e:
        logger.error(f"✗ Token regeneration failed: {e}")
        return False


def _load_token(token_file: str) -> tuple[str, str]:
    """Load JWT token from file and extract username.

    If token is expired, automatically regenerate it.

    Returns:
        Tuple of (token, username)
    """
    abs_path = os.path.abspath(token_file)
    try:
        with open(abs_path) as f:
            data = json.load(f)
            token = data.get("access_token") or data.get("token")
            if not token:
                raise ValueError("No access_token found in token file")

            # Check if token is expired
            if _is_token_expired(token):
                logger.warning("Token is expired or expiring soon, regenerating...")
                if _regenerate_token(token_file):
                    # Reload token from file after regeneration
                    with open(abs_path) as f2:
                        data = json.load(f2)
                        token = data.get("access_token") or data.get("token")
                        if not token:
                            raise ValueError("No access_token found after regeneration")
                else:
                    raise RuntimeError("Failed to regenerate expired token")

            username = _extract_username_from_jwt(token)
            logger.info(f"✓ Token loaded from: {abs_path}")
            logger.info(f"  User: {username}")
            logger.info(f"  Token length: {len(token)} characters")
            return token, username
    except FileNotFoundError:
        logger.error(f"✗ Token file not found: {abs_path}")
        logger.error(f"  Current directory: {os.getcwd()}")
        logger.error(f"  Looking for: {abs_path}")
        raise FileNotFoundError(f"Token file not found: {abs_path}")
    except json.JSONDecodeError as e:
        logger.error(f"✗ Invalid JSON in token file: {abs_path}")
        logger.error(f"  Error: {e}")
        raise ValueError(f"Invalid JSON in token file: {abs_path}")


def _make_request(
    method: str,
    url: str,
    token: str,
    data: dict[str, Any] | None = None,
    params: dict[str, Any] | None = None,
    timeout: int = REQUEST_TIMEOUT,
) -> requests.Response:
    """Make HTTP request with Bearer token."""
    headers = {
        "Authorization": f"Bearer {token[:50]}..." if token else "Bearer <NO_TOKEN>",
        "Content-Type": "application/json",
    }

    logger.info(f"HTTP {method} Request:")
    logger.info(f"  URL: {url}")
    logger.info("  Headers:")
    logger.info(
        f"    Authorization: Bearer {token[:50]}..." if token else "    Authorization: <NO_TOKEN>"
    )
    logger.info("    Content-Type: application/json")
    if params:
        logger.info(f"  Query Params: {json.dumps(params, indent=2)}")
    if data:
        logger.info(f"  Request Body: {len(json.dumps(data))} bytes")
        logger.debug(f"  Request Data: {json.dumps(data, indent=2)}")
    logger.info(f"  Timeout: {timeout}s")

    headers["Authorization"] = f"Bearer {token}"

    try:
        logger.info(f"→ Sending {method} request to {url}...")
        response = requests.request(
            method=method,
            url=url,
            json=data,
            params=params,
            headers=headers,
            timeout=timeout,
        )

        # Log response details
        logger.info(f"← Received HTTP {response.status_code}")
        logger.info("  Response Headers:")
        for header_name, header_value in response.headers.items():
            # Hide sensitive headers
            if header_name.lower() in ["authorization", "x-scopes"]:
                logger.info(
                    f"    {header_name}: {header_value[:50]}..."
                    if len(str(header_value)) > 50
                    else f"    {header_name}: {header_value}"
                )
            else:
                logger.info(f"    {header_name}: {header_value}")

        response_size = len(response.content) if response.content else 0
        logger.info(f"  Response Body: {response_size} bytes")

        if response.status_code >= 400:
            logger.warning(f"✗ HTTP {response.status_code} Error")
            try:
                resp_json = response.json()
                logger.warning(f"  Error Response: {json.dumps(resp_json, indent=2)}")
            except json.JSONDecodeError:
                logger.warning(f"  Error Response (raw): {response.text[:200]}")
        else:
            logger.info(f"✓ HTTP {response.status_code} Success")

        return response

    except requests.exceptions.Timeout as e:
        logger.error(f"✗ Request timed out after {timeout}s: {url}")
        logger.error(f"  Error: {e}")
        raise TimeoutError(f"Request timed out after {timeout} seconds: {url}")
    except requests.exceptions.ConnectionError as e:
        logger.error(f"✗ Failed to connect to {url}")
        logger.error(f"  Error: {e}")
        logger.error(f"  Check if service is running at {url}")
        raise ConnectionError(f"Failed to connect to {url}: {e}")
    except requests.exceptions.RequestException as e:
        logger.error(f"✗ Request failed: {e}")
        raise RuntimeError(f"Request failed: {e}")


def _print_response(response: requests.Response) -> None:
    """Pretty print response."""
    try:
        data = response.json()
        print(json.dumps(data, indent=2))
    except json.JSONDecodeError:
        print(response.text)


def list_agents(
    base_url: str,
    token: str,
) -> None:
    """List all agents via A2A Agent Management API."""
    endpoint = f"{base_url}{API_BASE}"
    logger.info(f"Listing all agents from {endpoint}...")

    try:
        response = _make_request("GET", endpoint, token)

        if response.status_code == 200:
            data = response.json()
            agents = data.get("agents", [])
            total_count = data.get("total_count", 0)

            if agents:
                print(f"Found {total_count} agent(s):")
                print("-" * 120)
                print(f"{'Agent Name':<40} | {'Path':<25} | {'Status':<8}")
                print("-" * 120)
                for agent in agents:
                    name = agent.get("name", "unknown")
                    path = agent.get("path", "unknown")
                    is_enabled = agent.get("is_enabled", False)
                    status = "ENABLED" if is_enabled else "DISABLED"
                    print(f"{name:<40} | {path:<25} | {status:<8}")
                print("-" * 120)
            else:
                print("No agents found")
        elif response.status_code == 401:
            print("Error: Authentication failed (401)")
            print("Make sure your JWT token is valid and not expired")
        else:
            print(f"Error: HTTP {response.status_code}")
            _print_response(response)

    except (TimeoutError, ConnectionError, RuntimeError) as e:
        print(f"Error: {e}")
        sys.exit(1)


def get_agent(
    base_url: str,
    token: str,
    agent_path: str,
) -> None:
    """Get agent details via A2A Agent Management API."""
    # Normalize path: /code-reviewer -> /code-reviewer
    if not agent_path.startswith("/"):
        agent_path = "/" + agent_path

    endpoint = f"{base_url}{API_BASE}{agent_path}"

    logger.info(f"Getting agent details for path '{agent_path}'...")

    try:
        response = _make_request("GET", endpoint, token)

        if response.status_code == 200:
            print(f"Agent details for: {agent_path}")
            _print_response(response)
        elif response.status_code == 404:
            print(f"Error: Agent at path '{agent_path}' not found")
        elif response.status_code == 401:
            print("Error: Authentication failed (401)")
        elif response.status_code == 403:
            print("Error: Access denied - you do not have permission to view this agent")
        else:
            print(f"Error: HTTP {response.status_code}")
            _print_response(response)

    except (TimeoutError, ConnectionError, RuntimeError) as e:
        print(f"Error: {e}")
        sys.exit(1)


def _check_agent_health(
    agent_url: str,
) -> tuple[bool, str]:
    """
    Check agent health by fetching agent card from /.well-known/agent-card.json.

    Args:
        agent_url: Base URL of the agent service

    Returns:
        Tuple of (success: bool, message: str)
    """
    if not agent_url:
        return False, "Agent URL not provided"

    health_endpoint = f"{agent_url}/.well-known/agent-card.json"
    logger.info(f"Checking agent health at: {health_endpoint}")

    try:
        response = requests.get(
            health_endpoint,
            timeout=REQUEST_TIMEOUT,
            headers={"Content-Type": "application/json"},
        )

        if response.status_code == 200:
            try:
                card_data = response.json()
                agent_name = card_data.get("name", "unknown")
                return True, f"Agent card retrieved successfully from {agent_name}"
            except json.JSONDecodeError:
                return False, "Agent returned invalid JSON for agent card"
        elif response.status_code == 404:
            return False, "Agent card endpoint not found (/.well-known/agent-card.json)"
        elif response.status_code == 503:
            return False, "Agent service unavailable (503)"
        else:
            return False, f"Agent returned HTTP {response.status_code}"

    except requests.exceptions.Timeout:
        return False, f"Agent health check timed out ({REQUEST_TIMEOUT}s)"
    except requests.exceptions.ConnectionError:
        return False, "Cannot connect to agent service (connection refused)"
    except Exception as e:
        return False, f"Health check error: {str(e)}"


def test_agent(
    base_url: str,
    token: str,
    agent_path: str,
) -> None:
    """Test agent accessibility and health via A2A Agent Management API."""
    # Normalize path
    if not agent_path.startswith("/"):
        agent_path = "/" + agent_path

    endpoint = f"{base_url}{API_BASE}{agent_path}"

    logger.info(f"Testing agent at path '{agent_path}'...")

    try:
        response = _make_request("GET", endpoint, token)

        if response.status_code == 200:
            data = response.json()
            name = data.get("name", "unknown")
            description = data.get("description", "")
            is_enabled = data.get("is_enabled", False)
            agent_url = data.get("url", "")

            print(f"Agent: {name}")
            print(f"Path: {agent_path}")
            print(f"Status: {'ENABLED' if is_enabled else 'DISABLED'}")
            print(f"Description: {description}")
            print(f"Service URL: {agent_url}")

            # Perform health check if agent is enabled and has URL
            if is_enabled and agent_url:
                print("\nPerforming health check...")
                health_passed, health_message = _check_agent_health(agent_url)
                if health_passed:
                    print("  Health Check: PASSED")
                    print(f"  Details: {health_message}")
                else:
                    print("  Health Check: FAILED")
                    print(f"  Reason: {health_message}")
            elif not is_enabled:
                print("\nHealth Check: SKIPPED (agent is disabled)")
            elif not agent_url:
                print("\nHealth Check: SKIPPED (no service URL configured)")

            print("\nAgent Registry Details:")
            _print_response(response)
        elif response.status_code == 404:
            print(f"Error: Agent at path '{agent_path}' not found")
        elif response.status_code == 401:
            print("Error: Authentication failed (401)")
        elif response.status_code == 403:
            print("Error: Access denied - you do not have permission to view this agent")
        else:
            print(f"Error: HTTP {response.status_code}")
            _print_response(response)

    except (TimeoutError, ConnectionError, RuntimeError) as e:
        print(f"Error: {e}")
        sys.exit(1)


def test_all_agents(
    base_url: str,
    token: str,
) -> None:
    """Test all agents accessibility and health."""
    endpoint = f"{base_url}{API_BASE}"
    logger.info("Testing all agents...")

    try:
        response = _make_request("GET", endpoint, token)

        if response.status_code == 200:
            data = response.json()
            agents = data.get("agents", [])

            if not agents:
                print("No agents to test")
                return

            passed = 0
            failed = 0

            print("Testing agents:")
            print("-" * 100)
            print(f"{'Agent Name':<35} | {'Registry':<8} | {'Health Check':<20}")
            print("-" * 100)

            for agent in agents:
                name = agent.get("name", "unknown")
                is_enabled = agent.get("is_enabled", False)
                agent_url = agent.get("url", "")
                registry_status = "ENABLED" if is_enabled else "DISABLED"

                # Perform health check if agent is enabled
                if is_enabled and agent_url:
                    health_passed, _ = _check_agent_health(agent_url)
                    health_status = "PASSED" if health_passed else "FAILED"
                    if health_passed:
                        passed += 1
                    else:
                        failed += 1
                else:
                    health_status = "SKIPPED"
                    passed += 1

                # Color-coded status
                if health_status == "PASSED":
                    result_icon = "✓"
                elif health_status == "FAILED":
                    result_icon = "✗"
                else:
                    result_icon = "-"

                print(f"{name:<35} | {registry_status:<8} | {result_icon} {health_status:<17}")

            print("-" * 100)
            print(f"Summary: {passed} passed, {failed} failed")
        elif response.status_code == 401:
            print("Error: Authentication failed (401)")
        else:
            print(f"Error: HTTP {response.status_code}")
            _print_response(response)

    except (TimeoutError, ConnectionError, RuntimeError) as e:
        print(f"Error: {e}")
        sys.exit(1)


def search_agents(
    base_url: str,
    token: str,
    query: str,
    max_results: int = 10,
) -> None:
    """Search agents using semantic search via natural language query."""
    endpoint = f"{base_url}{API_BASE}/discover/semantic"
    logger.info(f"Searching agents with query: {query}")

    params = {
        "query": query,
        "max_results": max_results,
    }

    try:
        response = _make_request("POST", endpoint, token, params=params)

        if response.status_code == 200:
            data = response.json()
            results = data.get("agents", [])

            if results:
                print(f"Found {len(results)} agent(s) matching '{query}':")
                print("-" * 110)
                print(f"{'Agent Name':<40} | {'Path':<25} | {'Score':<8}")
                print("-" * 110)
                for result in results:
                    name = result.get("name", "unknown")
                    path = result.get("path", "unknown")
                    score = result.get("score", 0.0)
                    print(f"{name:<40} | {path:<25} | {score:>7.4f}")
                print("-" * 110)
            else:
                print(f"No agents found matching '{query}'")
        elif response.status_code == 400:
            print("Error: Invalid search query (empty or malformed)")
        elif response.status_code == 401:
            print("Error: Authentication failed (401)")
        else:
            print(f"Error: HTTP {response.status_code}")
            _print_response(response)

    except (TimeoutError, ConnectionError, RuntimeError) as e:
        print(f"Error: {e}")
        sys.exit(1)


def register_agent(
    base_url: str,
    token: str,
    agent_file: str,
) -> None:
    """Register agent from JSON file via A2A Agent Management API."""
    import os

    abs_agent_file = os.path.abspath(agent_file)
    logger.info(f"Loading agent file from: {abs_agent_file}")

    try:
        with open(abs_agent_file) as f:
            agent_data = json.load(f)
        logger.info("✓ Agent file loaded successfully")
    except FileNotFoundError:
        logger.error(f"✗ Agent file not found: {abs_agent_file}")
        print(f"Error: File not found: {abs_agent_file}")
        sys.exit(1)
    except json.JSONDecodeError as e:
        logger.error(f"✗ Invalid JSON in file: {abs_agent_file}")
        logger.error(f"  Error: {e}")
        print(f"Error: Invalid JSON in file: {abs_agent_file}")
        sys.exit(1)

    # Use A2A Agent Management API endpoint for registration
    # Note: Goes through Nginx for JWT Bearer token validation via auth-server
    endpoint = f"{base_url}/api/agents/register"
    agent_name = agent_data.get("name", "Unknown")

    logger.info("=" * 80)
    logger.info("AGENT REGISTRATION REQUEST")
    logger.info("=" * 80)
    logger.info(f"Base URL: {base_url}")
    logger.info(f"Endpoint: {endpoint}")
    logger.info(f"Agent Name: {agent_name}")
    logger.info(f"Agent Path: {agent_data.get('path', 'N/A')}")
    logger.info(f"Agent File: {abs_agent_file}")
    logger.info("=" * 80)

    try:
        response = _make_request("POST", endpoint, token, agent_data)

        if response.status_code == 201:
            logger.info(f"✓ Agent '{agent_name}' registered successfully!")
            print(f"Agent '{agent_name}' registered successfully!")
            _print_response(response)
        elif response.status_code == 401:
            logger.error("✗ Authentication failed (HTTP 401)")
            logger.error(f"  Token file location: {os.path.abspath('.oauth-tokens/ingress.json')}")
            logger.error(f"  Token length: {len(token) if token else 0} characters")
            logger.error(
                f"  Authorization header: Bearer {token[:50]}..."
                if token
                else "  Authorization header: <NO_TOKEN>"
            )
            print("Error: Authentication failed (HTTP 401)")
            print("\nDEBUG INFORMATION:")
            print(f"  Token file: {os.path.abspath('.oauth-tokens/ingress.json')}")
            print(f"  Token length: {len(token) if token else 0} characters")
            print("\nNOTE: Make sure you have a valid token in '.oauth-tokens/ingress.json'")
            print("  The token should contain 'a2a-agent-admin' in groups claim")
            print("  Regenerate with: ./credentials-provider/generate_creds.sh")
            print("\nRESPONSE:")
            _print_response(response)
        elif response.status_code == 409:
            path = agent_data.get("path", "unknown")
            print(f"Error: Agent with path '{path}' already exists")
            _print_response(response)
        elif response.status_code == 422:
            print("Error: Validation failed - check agent JSON format")
            _print_response(response)
        elif response.status_code == 403:
            print("Error: Permission denied. You do not have permission to register agents")
            print(
                "\nNote: Agent registration requires proper Keycloak authentication with 'register_service' permission."
            )
            print("For testing/development, you may need to:")
            print("  1. Configure a Keycloak user with appropriate permissions")
            print("  2. Use the web UI dashboard to register agents")
            print("  3. Contact your administrator to grant registration permissions")
            _print_response(response)
        else:
            print(f"Error: HTTP {response.status_code}")
            _print_response(response)

    except (TimeoutError, ConnectionError, RuntimeError) as e:
        print(f"Error: {e}")
        sys.exit(1)


def update_agent(
    base_url: str,
    token: str,
    agent_path: str,
    agent_file: str,
) -> None:
    """Update agent via A2A Agent Management API."""
    # Normalize path
    if not agent_path.startswith("/"):
        agent_path = "/" + agent_path

    abs_agent_file = os.path.abspath(agent_file)
    logger.info(f"Loading agent file from: {abs_agent_file}")

    try:
        with open(abs_agent_file) as f:
            agent_data = json.load(f)
        logger.info("✓ Agent file loaded successfully")
    except FileNotFoundError:
        logger.error(f"✗ Agent file not found: {abs_agent_file}")
        print(f"Error: File not found: {abs_agent_file}")
        sys.exit(1)
    except json.JSONDecodeError as e:
        logger.error(f"✗ Invalid JSON in file: {abs_agent_file}")
        logger.error(f"  Error: {e}")
        print(f"Error: Invalid JSON in file: {abs_agent_file}")
        sys.exit(1)

    endpoint = f"{base_url}{API_BASE}{agent_path}"
    agent_name = agent_data.get("name", "Unknown")

    logger.info(f"Updating agent at path '{agent_path}'...")

    try:
        response = _make_request("PUT", endpoint, token, agent_data)

        if response.status_code == 200:
            logger.info(f"✓ Agent '{agent_name}' updated successfully!")
            print(f"Agent '{agent_name}' updated successfully!")
            _print_response(response)
        elif response.status_code == 404:
            print(f"Error: Agent at path '{agent_path}' not found")
        elif response.status_code == 401:
            print("Error: Authentication failed (401)")
        elif response.status_code == 403:
            print("Error: Access denied - you do not have permission to update this agent")
        else:
            print(f"Error: HTTP {response.status_code}")
            _print_response(response)

    except (TimeoutError, ConnectionError, RuntimeError) as e:
        print(f"Error: {e}")
        sys.exit(1)


def delete_agent(
    base_url: str,
    token: str,
    agent_path: str,
) -> None:
    """Delete agent via A2A Agent Management API."""
    # Normalize path
    if not agent_path.startswith("/"):
        agent_path = "/" + agent_path

    endpoint = f"{base_url}{API_BASE}{agent_path}"

    logger.info(f"Deleting agent at path '{agent_path}'...")

    try:
        response = _make_request("DELETE", endpoint, token)

        if response.status_code == 204:
            logger.info(f"✓ Agent at path '{agent_path}' deleted successfully!")
            print(f"Agent at path '{agent_path}' deleted successfully!")
        elif response.status_code == 404:
            print(f"Error: Agent at path '{agent_path}' not found")
        elif response.status_code == 401:
            print("Error: Authentication failed (401)")
        elif response.status_code == 403:
            print("Error: Access denied - you do not have permission to delete this agent")
        else:
            print(f"Error: HTTP {response.status_code}")
            _print_response(response)

    except (TimeoutError, ConnectionError, RuntimeError) as e:
        print(f"Error: {e}")
        sys.exit(1)


def toggle_agent(
    base_url: str,
    token: str,
    agent_path: str,
    enabled: bool,
) -> None:
    """Toggle agent enabled/disabled status via A2A Agent Management API."""
    # Normalize path
    if not agent_path.startswith("/"):
        agent_path = "/" + agent_path

    endpoint = f"{base_url}{API_BASE}{agent_path}/toggle"
    params = f"?enabled={str(enabled).lower()}"

    logger.info(f"Setting agent at path '{agent_path}' to {enabled}...")

    try:
        response = _make_request("POST", endpoint + params, token)

        if response.status_code == 200:
            data = response.json()
            is_enabled = data.get("is_enabled", False)
            status = "ENABLED" if is_enabled else "DISABLED"
            logger.info(f"✓ Agent at path '{agent_path}' toggled successfully!")
            print(f"Agent at path '{agent_path}' is now {status}")
            _print_response(response)
        elif response.status_code == 404:
            print(f"Error: Agent at path '{agent_path}' not found")
        elif response.status_code == 401:
            print("Error: Authentication failed (401)")
        elif response.status_code == 403:
            print("Error: Access denied - you do not have permission to toggle this agent")
        else:
            print(f"Error: HTTP {response.status_code}")
            _print_response(response)

    except (TimeoutError, ConnectionError, RuntimeError) as e:
        print(f"Error: {e}")
        sys.exit(1)


def main() -> None:
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Agent Management Script for MCP Gateway Registry - A2A Agent Management API",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # List all agents
  uv run python cli/agent_mgmt.py list

  # Get agent details
  uv run python cli/agent_mgmt.py get /code-reviewer

  # Register an agent from JSON file
  uv run python cli/agent_mgmt.py register cli/examples/test_code_reviewer_agent.json

  # Update an agent with new JSON
  uv run python cli/agent_mgmt.py update /code-reviewer cli/examples/updated_agent.json

  # Enable an agent
  uv run python cli/agent_mgmt.py toggle /code-reviewer true

  # Disable an agent
  uv run python cli/agent_mgmt.py toggle /code-reviewer false

  # Delete an agent
  uv run python cli/agent_mgmt.py delete /code-reviewer

  # Test agent accessibility
  uv run python cli/agent_mgmt.py test /code-reviewer

  # Test all agents
  uv run python cli/agent_mgmt.py test-all

  # Search agents with semantic query
  uv run python cli/agent_mgmt.py search "code review tool"

For more information on creating agent JSON files:
  cat cli/examples/README.md
""",
    )

    parser.add_argument(
        "--base-url",
        default=DEFAULT_BASE_URL,
        help=f"Base URL for API (default: {DEFAULT_BASE_URL})",
    )
    parser.add_argument(
        "--token-file",
        default=DEFAULT_TOKEN_FILE,
        help=f"Path to token JSON file (default: {DEFAULT_TOKEN_FILE})",
    )
    parser.add_argument(
        "--debug",
        action="store_true",
        help="Enable debug logging",
    )

    subparsers = parser.add_subparsers(dest="command", help="Command to execute")

    # Register command
    register_parser = subparsers.add_parser("register", help="Register agent from JSON file")
    register_parser.add_argument("file", help="Path to agent JSON file")

    # List command
    subparsers.add_parser("list", help="List all agents")

    # Get command
    get_parser = subparsers.add_parser("get", help="Get agent details")
    get_parser.add_argument("path", help="Agent path (e.g., /code-reviewer)")

    # Test command
    test_parser = subparsers.add_parser("test", help="Test agent accessibility")
    test_parser.add_argument("path", help="Agent path (e.g., /code-reviewer)")

    # Test all command
    subparsers.add_parser("test-all", help="Test all agents")

    # Search command
    search_parser = subparsers.add_parser("search", help="Search agents using semantic query")
    search_parser.add_argument(
        "query", help="Natural language search query (e.g., 'code review agent')"
    )
    search_parser.add_argument(
        "--max-results",
        type=int,
        default=10,
        help="Maximum number of results (default: 10)",
    )

    # Update command
    update_parser = subparsers.add_parser("update", help="Update agent from JSON file")
    update_parser.add_argument("path", help="Agent path (e.g., /code-reviewer)")
    update_parser.add_argument("file", help="Path to updated agent JSON file")

    # Delete command
    delete_parser = subparsers.add_parser("delete", help="Delete agent")
    delete_parser.add_argument("path", help="Agent path (e.g., /code-reviewer)")

    # Toggle command
    toggle_parser = subparsers.add_parser("toggle", help="Toggle agent enabled/disabled status")
    toggle_parser.add_argument("path", help="Agent path (e.g., /code-reviewer)")
    toggle_parser.add_argument(
        "enabled",
        type=lambda x: x.lower() == "true",
        help="Enable (true) or disable (false) the agent",
    )

    args = parser.parse_args()

    if args.debug:
        logger.setLevel(logging.DEBUG)

    if not args.command:
        parser.print_help()
        sys.exit(1)

    # Load token
    try:
        token, username = _load_token(args.token_file)
    except (FileNotFoundError, ValueError) as e:
        print(f"Error: {e}")
        sys.exit(1)

    # Execute command
    if args.command == "list":
        list_agents(args.base_url, token)
    elif args.command == "get":
        get_agent(args.base_url, token, args.path)
    elif args.command == "test":
        test_agent(args.base_url, token, args.path)
    elif args.command == "test-all":
        test_all_agents(args.base_url, token)
    elif args.command == "search":
        search_agents(args.base_url, token, args.query, args.max_results)
    elif args.command == "register":
        register_agent(args.base_url, token, args.file)
    elif args.command == "update":
        update_agent(args.base_url, token, args.path, args.file)
    elif args.command == "delete":
        delete_agent(args.base_url, token, args.path)
    elif args.command == "toggle":
        toggle_agent(args.base_url, token, args.path, args.enabled)


if __name__ == "__main__":
    main()
