#!/usr/bin/env python3
"""
Sixpert K2 - Function Calling Agent Example
============================================
Demonstrates how to use Sixpert K2's function calling capabilities
in an agentic loop. K2's MoE architecture provides specialized experts
for tool selection and argument formatting.

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


TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "search_research_papers",
            "description": "Search for academic research papers on a topic",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "Research topic"},
                    "year_from": {"type": "integer", "description": "Start year"},
                    "max_results": {"type": "integer", "default": 10},
                },
                "required": ["query"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "run_python_code",
            "description": "Execute Python code in a sandboxed environment",
            "parameters": {
                "type": "object",
                "properties": {
                    "code": {"type": "string", "description": "Python code to execute"},
                    "timeout": {"type": "integer", "description": "Timeout in seconds", "default": 30},
                },
                "required": ["code"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "analyze_data",
            "description": "Analyze a dataset and provide statistical insights",
            "parameters": {
                "type": "object",
                "properties": {
                    "data_description": {"type": "string", "description": "Description of the data"},
                    "analysis_type": {"type": "string", "enum": ["descriptive", "inferential", "predictive"], "description": "Type of analysis"},
                },
                "required": ["data_description", "analysis_type"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "calculate_math",
            "description": "Perform mathematical calculations",
            "parameters": {
                "type": "object",
                "properties": {
                    "expression": {"type": "string", "description": "Mathematical expression"},
                    "precision": {"type": "integer", "description": "Decimal places", "default": 6},
                },
                "required": ["expression"],
            },
        },
    },
]


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"  Tool call: {name}({json.dumps(args, indent=2)})")

    if name == "search_research_papers":
        return json.dumps({
            "papers": [
                {"title": f"Advances in {args['query']}", "year": 2025, "citations": 42},
                {"title": f"Survey of {args['query']} Methods", "year": 2024, "citations": 128},
            ],
            "total_found": 1247,
        })
    elif name == "run_python_code":
        return json.dumps({"stdout": "42", "stderr": "", "exit_code": 0})
    elif name == "analyze_data":
        return json.dumps({
            "summary": "Dataset shows normal distribution with mean=0.0, std=1.0",
            "insights": ["No significant outliers detected", "Data is well-scaled"],
        })
    elif name == "calculate_math":
        return json.dumps({"result": 3.141593, "expression": args["expression"]})

    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=16384,
        n_gpu_layers=-1,
        verbose=False,
    )

    messages = [
        {
            "role": "system",
            "content": (
                "You are Sixpert K2, a deep reasoning engine with advanced "
                "agentic capabilities. When the user asks a question that "
                "requires external tools, use the available functions. "
                "Think step-by-step and plan your approach before calling tools. "
                "You can call multiple tools in sequence to solve complex problems."
            ),
        },
        {"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.6,
            stream=False,
        )

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

        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[:150]}...")

                messages.append({
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [tool_call],
                })
                messages.append({
                    "role": "tool",
                    "tool_call_id": tool_call["id"],
                    "content": tool_result,
                })
        else:
            print(f"\nSixpert K2: {message['content']}")
            break
    else:
        print("\nReached maximum turns.")


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Sixpert K2 Function Calling")
    parser.add_argument("--model", type=str, default="SixpertK2.gguf", help="Path to GGUF model")
    parser.add_argument("--query", type=str, default="Calculate pi to 6 decimal places and search for research on numerical methods", help="User query")

    args = parser.parse_args()

    print("=" * 60)
    print("  Sixpert K2 - Function Calling Agent")
    print("  MoE Architecture | Deep Reasoning")
    print("=" * 60)

    run_agent(args.model, args.query)


if __name__ == "__main__":
    main()