# Copyright 2024-2026 The DNS-AID Authors
# SPDX-License-Identifier: Apache-2.0

"""
A2A Agent Card — Google's Agent-to-Agent protocol agent description format.

The Agent Card is a JSON document served at `/.well-known/agent.json` that
describes an agent's capabilities, authentication requirements, and metadata.

Reference: https://google.github.io/A2A/
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any
from urllib.parse import urljoin

import httpx
import structlog

logger = structlog.get_logger(__name__)

# Well-known path for A2A Agent Cards
A2A_AGENT_CARD_PATH = "/.well-known/agent.json"


@dataclass
class A2AProvider:
    """Agent provider/organization information."""

    organization: str
    url: str | None = None


@dataclass
class A2ASkill:
    """A single skill/capability the agent can perform."""

    id: str
    name: str
    description: str | None = None
    input_modes: list[str] = field(default_factory=lambda: ["text"])
    output_modes: list[str] = field(default_factory=lambda: ["text"])
    tags: list[str] = field(default_factory=list)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> A2ASkill:
        """Parse a skill from JSON dict."""
        return cls(
            id=data.get("id", ""),
            name=data.get("name", ""),
            description=data.get("description"),
            input_modes=data.get("inputModes", ["text"]),
            output_modes=data.get("outputModes", ["text"]),
            tags=data.get("tags", []),
        )


@dataclass
class A2AAuthentication:
    """Authentication requirements for the agent."""

    schemes: list[str] = field(default_factory=list)
    credentials: str | None = None

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> A2AAuthentication:
        """Parse authentication from JSON dict."""
        return cls(
            schemes=data.get("schemes", []),
            credentials=data.get("credentials"),
        )


@dataclass
class A2AAgentCard:
    """
    Google A2A Agent Card — the canonical agent description format.

    Fetched from: https://{domain}/.well-known/agent.json

    Attributes:
        name: Human-readable agent name.
        url: Agent's endpoint URL.
        version: Agent version string.
        description: Human-readable description.
        provider: Organization/provider info.
        skills: List of capabilities the agent offers.
        authentication: Auth requirements.
        default_input_modes: Default input formats (text, data, audio, video).
        default_output_modes: Default output formats.
        metadata: Additional fields not in the core schema.
    """

    name: str
    url: str
    version: str = "1.0.0"
    description: str | None = None
    provider: A2AProvider | None = None
    skills: list[A2ASkill] = field(default_factory=list)
    authentication: A2AAuthentication | None = None
    default_input_modes: list[str] = field(default_factory=lambda: ["text"])
    default_output_modes: list[str] = field(default_factory=lambda: ["text"])
    metadata: dict[str, Any] = field(default_factory=dict)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> A2AAgentCard:
        """Parse an Agent Card from JSON dict."""
        # Parse provider
        provider = None
        if "provider" in data and isinstance(data["provider"], dict):
            provider = A2AProvider(
                organization=data["provider"].get("organization", ""),
                url=data["provider"].get("url"),
            )

        # Parse skills
        skills = []
        if "skills" in data and isinstance(data["skills"], list):
            skills = [A2ASkill.from_dict(s) for s in data["skills"] if isinstance(s, dict)]

        # Parse authentication
        auth = None
        if "authentication" in data and isinstance(data["authentication"], dict):
            auth = A2AAuthentication.from_dict(data["authentication"])

        # Collect unknown fields as metadata
        known_keys = {
            "name",
            "url",
            "version",
            "description",
            "provider",
            "skills",
            "authentication",
            "defaultInputModes",
            "defaultOutputModes",
        }
        metadata = {k: v for k, v in data.items() if k not in known_keys}

        return cls(
            name=data.get("name", ""),
            url=data.get("url", ""),
            version=data.get("version", "1.0.0"),
            description=data.get("description"),
            provider=provider,
            skills=skills,
            authentication=auth,
            default_input_modes=data.get("defaultInputModes", ["text"]),
            default_output_modes=data.get("defaultOutputModes", ["text"]),
            metadata=metadata,
        )

    @property
    def skill_ids(self) -> list[str]:
        """Get list of skill IDs (convenience for capability matching)."""
        return [s.id for s in self.skills]

    @property
    def skill_names(self) -> list[str]:
        """Get list of skill names."""
        return [s.name for s in self.skills]

    def to_capabilities(self) -> list[str]:
        """Convert skills to DNS-AID capability format (skill IDs)."""
        return self.skill_ids


async def fetch_agent_card(
    endpoint: str,
    timeout: float = 10.0,
) -> A2AAgentCard | None:
    """
    Fetch an A2A Agent Card from the well-known location.

    Given an agent endpoint (e.g., "https://agent.example.com"), fetches
    the Agent Card from "https://agent.example.com/.well-known/agent.json".

    Args:
        endpoint: Agent's base URL (scheme + host).
        timeout: HTTP request timeout in seconds.

    Returns:
        A2AAgentCard if successfully fetched and parsed, None otherwise.

    Example:
        >>> card = await fetch_agent_card("https://payment.example.com")
        >>> print(card.skills[0].name)
        "Process Payment"
    """
    # Construct the well-known URL
    if not endpoint.startswith("https://"):
        endpoint = f"https://{endpoint}"

    card_url = urljoin(endpoint.rstrip("/") + "/", A2A_AGENT_CARD_PATH.lstrip("/"))

    # SSRF protection: validate URL before fetching
    try:
        from dns_aid.utils.url_safety import UnsafeURLError, validate_fetch_url

        validate_fetch_url(card_url)
    except UnsafeURLError as e:
        logger.warning("Agent Card URL blocked by SSRF protection", url=card_url, error=str(e))
        return None

    logger.debug("Fetching A2A Agent Card", url=card_url)

    try:
        async with httpx.AsyncClient(timeout=timeout, follow_redirects=False) as client:
            response = await client.get(card_url)

            if response.status_code != 200:
                logger.debug(
                    "Agent Card fetch failed",
                    url=card_url,
                    status_code=response.status_code,
                )
                return None

            data = response.json()

            if not isinstance(data, dict):
                logger.debug("Agent Card is not a JSON object", url=card_url)
                return None

            card = A2AAgentCard.from_dict(data)

            logger.debug(
                "Agent Card fetched successfully",
                url=card_url,
                name=card.name,
                skills_count=len(card.skills),
            )
            return card

    except httpx.TimeoutException:
        logger.debug("Agent Card fetch timed out", url=card_url)
        return None
    except httpx.ConnectError:
        logger.debug("Agent Card connection failed", url=card_url)
        return None
    except Exception as e:
        logger.debug("Agent Card fetch error", url=card_url, error=str(e))
        return None


async def fetch_agent_card_from_domain(
    domain: str,
    timeout: float = 10.0,
) -> A2AAgentCard | None:
    """
    Fetch an A2A Agent Card from a domain's well-known location.

    Convenience function that constructs the full URL from just a domain.

    Args:
        domain: Domain name (e.g., "example.com").
        timeout: HTTP request timeout in seconds.

    Returns:
        A2AAgentCard if successfully fetched and parsed, None otherwise.

    Example:
        >>> card = await fetch_agent_card_from_domain("example.com")
    """
    return await fetch_agent_card(f"https://{domain}", timeout=timeout)
