### MCP Server to sourcing agent cards.
import json
import logging
import os
from pathlib import Path

import numpy as np
from langchain_ollama import OllamaEmbeddings
from mcp.server import FastMCP
import pandas as pd


# BASE_DIR = Path(__file__).resolve().parent.parent  # goes from automa_ai/mcp_servers/ -> automa_ai/
# AGENT_CARDS_DIR = BASE_DIR / "agent_cards"
MODEL = "ollama_chat/llama3.1:8b"

logger = logging.getLogger(__name__)

MCP_NAME = "agent_card_mcp"


def generate_embeddings(text):
    """
    Generates embeddings for the given text using ollama
    :param text:
    :return:
    """
    embed_model = OllamaEmbeddings(model="mxbai-embed-large")
    return embed_model.embed_query(text)


def load_agent_cards(agent_card_dir: str):
    """Loads agent card data from JSON files within a specified directory
    agent_card_dir: directory to agent cards
    Returns:
        A list containing JSON data from an agent card file found in the specified directory.
        Returns an empty list if the directory is empty, contains no '.json' files,
        or if all '.json' files encounter errors during processing.
    """
    card_uris = []
    agent_cards = []
    dir_path = Path(agent_card_dir)
    if not dir_path.is_dir():
        logger.error(
            f"Agent cards directory not found or is not a directory: {agent_card_dir}"
        )
        return agent_cards

    logger.info(f"Loading agent cards from card repo: {agent_card_dir}")

    for filename in os.listdir(agent_card_dir):
        if filename.lower().endswith(".json"):
            file_path = dir_path / filename

            if file_path.is_file():
                logger.info(f"Reading file: {filename}")
                try:
                    with file_path.open("r", encoding="utf-8") as f:
                        data = json.load(f)
                        card_uris.append(
                            f"resource://agent_cards/{Path(filename).stem}"
                        )
                        agent_cards.append(data)
                except json.JSONDecodeError as jde:
                    logger.error(f"JSON Decoder Error {jde}")
                except OSError as e:
                    logger.error(f"Error reading file {filename}: {e}.")
                except Exception as e:
                    logger.error(
                        f"An unexpected error occurred processing {filename}: {e}",
                        exc_info=True,
                    )

    logger.info(f"Finished loading agent cards. Found {len(agent_cards)} cards.")
    return card_uris, agent_cards


def build_agent_card_embeddings(agent_card_dir: str) -> pd.DataFrame:
    card_uris, agent_cards = load_agent_cards(agent_card_dir)

    if not agent_cards:
        return pd.DataFrame()

    texts = [json.dumps(card) for card in agent_cards]
    embedding_model = OllamaEmbeddings(model="mxbai-embed-large")
    embeddings = embedding_model.embed_documents(texts)  # shape: (N, D)

    df = pd.DataFrame(
        {"card_uri": card_uris, "agent_card": agent_cards, "embedding": embeddings}
    )
    return df


def find_best_match(df: pd.DataFrame, query: str) -> dict:
    if df.empty:
        raise ValueError("No agent cards loaded.")
    embedding_model = OllamaEmbeddings(model="mxbai-embed-large")
    query_vec = embedding_model.embed_query(query)
    scores = df["embedding"].apply(lambda emb: np.dot(emb, query_vec))

    best_idx = scores.idxmax()
    best_card = df.loc[best_idx, "agent_card"]
    best_uri = df.loc[best_idx, "card_uri"]
    best_score = scores[best_idx]

    logger.info(f"Best match: {best_uri} with score {best_score:.4f}")
    return best_card


def get_card_by_uri(df: pd.DataFrame, uri: str) -> dict | None:
    result = df[df["card_uri"] == uri]
    if result.empty:
        return None
    return result.iloc[0]["agent_card"]


def serve(host, port, transport, agent_cards_dir: str):
    """Initialize and runs the agent cards mcp_servers server.
    Args:
        host: The hostname or IP address to bind the server to.
        port: The port number to bind the server to.
        transport: The transport mechanism for the MCP server (e.g., 'stdio', 'sse')
        agent_cards_dir: directory to agent_cards

    Raises:
        ValueError
    """
    logger.info("Starting Agent Cards MCP Server")
    mcp = FastMCP("agent-cards", host=host, port=port)

    df = build_agent_card_embeddings(agent_cards_dir)

    @mcp.tool(
        name="find_agent",
        description="Finds the most relevant agent card based on a natural language query string.",
    )
    def find_agent(query: str) -> dict:
        """
        Finds the most relevant agent card based on a query string.

        This function takes a user query, typically a natural language question or a task generated by an agent,
        generates its embedding, and compares it against the
        pre-computed embedding of the loaded agent cards. It uses the dot product to measure similarity and identifies the agent card with the highest similarity score.

        Args:
            query: The natual language query string used to search for a relevant agent.

        Returns:
            The json representing the agent card deemed most relevant to the input query based on embedding similarity.
        """
        return find_best_match(df, query)

    @mcp.resource("resource://agent_cards/{card_name}", mime_type="application/json")
    def get_agent_card(card_name: str) -> dict:
        """Retrieves an agent card as a json / dictionary for the MCP resource endpoint.

        This function serves as the handler for the MCP resource identified by
            the URI 'resource://agent_cards/{card_name}'.

        Returns:
            A json / dictionary
        """
        uri = f"resource://agent_cards/{card_name}"
        card = get_card_by_uri(df, uri)
        if card:
            return {"agent_card": card}
        return {}

    logger.info(f"Agent cards MCP Server at {host}:{port} and transport {transport}")
    mcp.run(transport=transport)
