"""
Nexus Demo — Tool Calling (Function Calling)

Demonstrates agents using tools through the ReAct loop:
  1. Simple tools   — calculator, clock
  2. Async tools    — HTTP fetch
  3. Multi-tool     — agent composes multiple tools to complete a task

The agent's ReAct loop (Think -> Act -> Observe -> Repeat) automatically
invokes tools and feeds results back until it produces a final answer.

Usage:
    python examples/demo_tools.py
"""

import asyncio
import json
import math
import sys
from datetime import datetime, timezone

from _config import get_model_config

from nexus import AgentCapability, Team
from nexus.core.agent import AgentTool
from nexus.orchestration import TopologyMode


# ================================================================
# Tool definitions
# ================================================================


def calculate(expression: str) -> str:
    """Evaluate a mathematical expression safely.

    Args:
        expression: A Python math expression, e.g. "2**10", "math.sqrt(144)"
    """
    allowed = {"__builtins__": {}, "math": math, "abs": abs, "round": round}
    try:
        result = eval(expression, allowed)  # noqa: S307
        return str(result)
    except Exception as exc:
        return f"Error: {exc}"


def get_current_time() -> str:
    """Return the current date and time in ISO-8601 format."""
    return datetime.now(tz=timezone.utc).isoformat()


def search_knowledge_base(query: str, max_results: int = 3) -> str:
    """Search an internal knowledge base for relevant information.

    Args:
        query: The search query
        max_results: Maximum number of results to return
    """
    kb = {
        "python": "Python is a high-level, interpreted language created by Guido van Rossum.",
        "rust": "Rust is a systems programming language focused on safety and performance.",
        "javascript": "JavaScript is a dynamic language primarily used for web development.",
        "machine learning": "ML is a subset of AI that learns patterns from data.",
        "nexus framework": "Nexus is a protocol-native multi-agent framework for building AI agent systems.",
        "quicksort": "Quicksort is a divide-and-conquer algorithm with O(n log n) average time complexity.",
        "bubble sort": "Bubble sort is a simple sorting algorithm with O(n^2) time complexity.",
    }
    results = []
    for key, value in kb.items():
        if any(word in key for word in query.lower().split()):
            results.append({"title": key, "content": value})
    if not results:
        return json.dumps({"message": "No results found", "query": query})
    return json.dumps(results[:max_results], ensure_ascii=False)


async def fetch_url(url: str) -> str:
    """Fetch the content of a URL and return the first 500 characters.

    Args:
        url: The URL to fetch
    """
    import httpx

    try:
        async with httpx.AsyncClient(follow_redirects=True, timeout=10) as client:
            resp = await client.get(url)
            resp.raise_for_status()
            text = resp.text[:500]
            return f"HTTP {resp.status_code} | {text}"
    except Exception as exc:
        return f"Fetch error: {exc}"


def write_file(filename: str, content: str) -> str:
    """Write content to a file (simulated).

    Args:
        filename: Name of the file to write
        content: Content to write
    """
    # In a real app this would write to disk
    return f"Wrote {len(content)} chars to {filename}"


# ================================================================
# Demo 1: Simple tool calling
# ================================================================


async def demo_simple_tools():
    """Agent uses calculator and clock tools."""
    print("=" * 60)
    print("Demo 1: Simple Tool Calling (calculator + clock)")
    print("=" * 60)

    team = Team(**get_model_config(), mode=TopologyMode.ROUTER, verbose=False)

    team.add_agent(
        "assistant",
        instructions=(
            "You are a helpful assistant with access to a calculator and a clock. "
            "Use your tools to answer questions accurately. Show your work."
        ),
        tools=[calculate, get_current_time],
    )

    prompts = [
        "What is 2^32 - 1?",
        "What time is it right now in UTC?",
        "Calculate the compound interest: principal=10000, rate=5%, years=10. "
        "Use the formula: P * (1 + r)^t",
    ]
    for prompt in prompts:
        print(f"\nUser: {prompt}")
        result = await team.run(prompt)
        print(f"Assistant: {result[:300]}")

    await team.close()


# ================================================================
# Demo 2: Async tool (HTTP fetch)
# ================================================================


async def demo_async_tool():
    """Agent uses an async HTTP fetch tool."""
    print("\n" + "=" * 60)
    print("Demo 2: Async Tool (HTTP fetch)")
    print("=" * 60)

    team = Team(**get_model_config(), mode=TopologyMode.ROUTER, verbose=False)

    team.add_agent(
        "researcher",
        instructions=(
            "You are a web researcher. Use the fetch_url tool to retrieve "
            "information from URLs. Summarize what you find."
        ),
        tools=[fetch_url],
    )

    print("\nUser: Fetch https://httpbin.org/json and summarize its content")
    result = await team.run("Fetch https://httpbin.org/json and summarize its content")
    print(f"Assistant: {result[:400]}")

    await team.close()


# ================================================================
# Demo 3: Multi-tool composition (research + calculate + write)
# ================================================================


async def demo_multi_tool():
    """Agent composes multiple tools to complete a complex task."""
    print("\n" + "=" * 60)
    print("Demo 3: Multi-tool Composition (search + calculate + write)")
    print("=" * 60)

    team = Team(**get_model_config(), mode=TopologyMode.ROUTER, verbose=False)

    team.add_agent(
        "analyst",
        instructions=(
            "You are a research analyst. You can search a knowledge base, "
            "perform calculations, and write files. "
            "When asked to produce a report, use ALL your tools: "
            "search for info, compute any numbers, then compile results."
        ),
        tools=[search_knowledge_base, calculate, write_file],
    )

    print(
        "\nUser: Search for info about quicksort and bubble sort, compare their "
        "complexity, calculate 2^20 (a typical large-n scenario), and write "
        "a summary report to report.md"
    )
    result = await team.run(
        "Search for info about quicksort and bubble sort, compare their "
        "complexity, calculate 2^20 (a typical large-n scenario), and write "
        "a summary report to report.md"
    )
    print(f"Assistant: {result[:500]}")

    await team.close()


# ================================================================
# Main
# ================================================================


async def main():
    demos = {
        "simple": demo_simple_tools,
        "async": demo_async_tool,
        "multi": demo_multi_tool,
    }
    choice = sys.argv[1] if len(sys.argv) > 1 else "all"

    if choice == "all":
        for fn in demos.values():
            await fn()
    elif choice in demos:
        await demos[choice]()
    else:
        print(f"Usage: python demo_tools.py [{'|'.join(demos.keys())}|all]")


if __name__ == "__main__":
    asyncio.run(main())
