#!/usr/bin/env python3
"""
Sixpert K1 - Function Calling Example
======================================
Demonstrates how to use Sixpert K1's function calling capabilities.

Usage:
    python function_calling.py
"""

import json
import sys

try:
    from llama_cpp import Llama
except ImportError:
    print("Installing llama-cpp-python...")
    import subprocess
    subprocess.check_call([sys.executable, "-m", "pip", "install", "llama-cpp-python"])
    from llama_cpp import Llama


# Define available functions
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "Get the current weather for a given location",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "The city name, e.g. 'San Francisco'",
                    },
                    "unit": {
                        "type": "string",
                        "enum": ["celsius", "fahrenheit"],
                        "description": "Temperature unit",
                    },
                },
                "required": ["location"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "execute_code",
            "description": "Execute Python code and return the result",
            "parameters": {
                "type": "object",
                "properties": {
                    "code": {
                        "type": "string",
                        "description": "The Python code to execute",
                    },
                },
                "required": ["code"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "search_web",
            "description": "Search the web for information on a topic",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "The search query",
                    },
                    "max_results": {
                        "type": "integer",
                        "description": "Maximum number of results to return",
                        "default": 5,
                    },
                },
                "required": ["query"],
            },
        },
    },
]


def mock_execute_tool(tool_call: dict) -> str:
    """Mock execution of a tool call. Replace with real implementations."""
    name = tool_call["function"]["name"]
    args = json.loads(tool_call["function"]["arguments"])

    print(f"  Executing: {name}({args})")

    if name == "get_weather":
        return json.dumps({
            "location": args["location"],
            "temperature": 22,
            "condition": "Partly cloudy",
            "unit": args.get("unit", "celsius"),
        })
    elif name == "execute_code":
        return json.dumps({"result": "42", "success": True})
    elif name == "search_web":
        return json.dumps({
            "results": [
                {"title": f"Result 1 for {args['query']}", "url": "https://example.com"},
                {"title": f"Result 2 for {args['query']}", "url": "https://example.org"},
            ]
        })
    return json.dumps({"error": f"Unknown tool: {name}"})


def run_agent(model_path: str, user_query: str, max_turns: int = 5):
    """Run an agentic loop with function calling."""
    print(f"\nUser Query: {user_query}")
    print("-" * 50)

    llm = Llama(
        model_path=model_path,
        n_ctx=8192,
        n_gpu_layers=-1,
        verbose=False,
    )

    messages = [
        {
            "role": "system",
            "content": (
                "You are Sixpert K1, a precision logic engine. "
                "When the user asks a question that requires external tools, "
                "use the available functions to gather information. "
                "Think step-by-step before calling any tools."
            ),
        },
        {"role": "user", "content": user_query},
    ]

    for turn in range(max_turns):
        print(f"\n--- Turn {turn + 1} ---")

        response = llm.create_chat_completion(
            messages=messages,
            tools=TOOLS,
            tool_choice="auto",
            temperature=0.7,
            stream=False,
        )

        choice = response["choices"][0]
        message = choice["message"]

        # Check if model wants to call a tool
        if message.get("tool_calls"):
            for tool_call in message["tool_calls"]:
                print(f"  Tool call: {tool_call['function']['name']}")
                tool_result = mock_execute_tool(tool_call)
                print(f"  Result: {tool_result[:100]}...")

                # Add assistant message with tool call
                messages.append({
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [tool_call],
                })
                # Add tool result
                messages.append({
                    "role": "tool",
                    "tool_call_id": tool_call["id"],
                    "content": tool_result,
                })
        else:
            # Model responded directly
            print(f"\nSixpert K1: {message['content']}")
            break
    else:
        print("\nReached maximum turns.")


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Sixpert K1 Function Calling")
    parser.add_argument("--model", type=str, default="SixpertK1.gguf", help="Path to GGUF model")
    parser.add_argument("--query", type=str, default="What's the weather in Tokyo?", help="User query")

    args = parser.parse_args()

    print("=" * 60)
    print("  Sixpert K1 - Function Calling Agent")
    print("=" * 60)

    run_agent(args.model, args.query)


if __name__ == "__main__":
    main()