# tools.py
# Three tiny tools the agent can call. Fake weather data so no extra API key is needed.


FAKE_WEATHER = {
    "mumbai": "32 C, sunny, humid",
    "london": "14 C, cloudy, light rain",
    "tokyo": "21 C, clear skies",
    "new york": "18 C, partly cloudy",
    "paris": "16 C, overcast",
}


def add(a: float, b: float) -> str:
    return f"{a + b}"


def multiply(a: float, b: float) -> str:
    return f"{a * b}"


def get_weather(city: str) -> str:
    return FAKE_WEATHER.get(
        city.lower(),
        f"Weather for {city}: 25 C, partly cloudy (demo data)",
    )


# ----------------------------------------------------------------
# ML example tools — wrap the helpers from examples.py so the agent
# can search the paper catalog, look up a paper, or list all papers.
# ----------------------------------------------------------------
from examples import search_examples, get_paper_info, list_papers


def search_ml_examples(query: str) -> str:
    """Search the ML paper sentence catalog by keyword."""
    matches = search_examples(query)
    if not matches:
        return f"No sentences matching '{query}'."
    lines = [f"Found {len(matches)} match(es):"]
    for m in matches[:5]:
        lines.append(
            f"- [{m['label']}] \"{m['sentence']}\" "
            f"({m['paper_title']}, {m['year']})"
        )
    return "\n".join(lines)


def ml_paper_info(paper_id: str) -> str:
    """Look up metadata for a specific paper by its id."""
    info = get_paper_info(paper_id)
    if not info:
        return f"No paper with id '{paper_id}'."
    return (
        f"{info['title']} ({info['year']}) — "
        f"id: {info['paper_id']}, sentences in catalog: {info['sentence_count']}"
    )


def list_ml_papers() -> str:
    """List every paper in the catalog."""
    papers = list_papers()
    lines = [f"{len(papers)} papers in catalog:"]
    for p in papers:
        lines.append(
            f"- {p['paper_id']}: {p['title']} ({p['year']}) "
            f"— {p['sentence_count']} sentences"
        )
    return "\n".join(lines)


TOOL_FUNCTIONS = {
    "add": add,
    "multiply": multiply,
    "get_weather": get_weather,
    "search_ml_examples": search_ml_examples,
    "ml_paper_info": ml_paper_info,
    "list_ml_papers": list_ml_papers,
}


TOOL_SCHEMAS = [
    {
        "type": "function",
        "function": {
            "name": "add",
            "description": "Add two numbers and return the result.",
            "parameters": {
                "type": "object",
                "properties": {
                    "a": {"type": "number", "description": "First number"},
                    "b": {"type": "number", "description": "Second number"},
                },
                "required": ["a", "b"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "multiply",
            "description": "Multiply two numbers and return the result.",
            "parameters": {
                "type": "object",
                "properties": {
                    "a": {"type": "number", "description": "First number"},
                    "b": {"type": "number", "description": "Second number"},
                },
                "required": ["a", "b"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "Get the current weather for a given city.",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "City name"},
                },
                "required": ["city"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "search_ml_examples",
            "description": "Search the built-in ML paper sentence catalog. Returns sentences matching the query along with their paper title, year, and label.",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "Keyword or phrase to search for"},
                },
                "required": ["query"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "ml_paper_info",
            "description": "Look up metadata (title, year, sentence count) for a specific ML paper by its id like 'vaswani-2017-attention'.",
            "parameters": {
                "type": "object",
                "properties": {
                    "paper_id": {"type": "string", "description": "Paper id slug"},
                },
                "required": ["paper_id"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "list_ml_papers",
            "description": "List every ML paper in the built-in catalog with its id, title, year, and sentence count.",
            "parameters": {
                "type": "object",
                "properties": {},
            },
        },
    },
]