import json
import subprocess
from typing import Iterable

from openai import OpenAI
from openai.types.chat import ChatCompletionMessageParam, ChatCompletionToolUnionParam

client = OpenAI(base_url="https://llama-cpp.reeselink.com", api_key="")


class ToolCallController:
    def __init__(self, max_tool_calls=10):
        self.max_tool_calls = max_tool_calls
        self.tool_call_count = 0

    def is_tool_call_allowed(self):
        return self.tool_call_count < self.max_tool_calls

    def increment(self):
        self.tool_call_count += 1

    def reset(self):
        self.tool_call_count = 0


# Register tools
tools: Iterable[ChatCompletionToolUnionParam] = [
    {
        "type": "function",
        "function": {
            "name": "list_servers",
            "description": "Lists the available servers to perform operations on",
            "parameters": {
                "type": "object",
                "properties": {"server_id": {"type": "string", "enum": ["all"]}},
                "required": ["server_id"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "check_updates",
            "description": "Check if a given server needs updated.",
            "parameters": {
                "type": "object",
                "properties": {
                    "server_id": {
                        "type": "string",
                    }
                },
                "required": ["server_id"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "perform_updates",
            "description": "Update a given server to the latest package versions. Does not reboot automatically.",
            "parameters": {
                "type": "object",
                "properties": {
                    "server_id": {
                        "type": "string",
                    }
                },
                "required": ["server_id"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "reboot",
            "description": "Reboot a given server. Waits for server to be responsive again.",
            "parameters": {
                "type": "object",
                "properties": {
                    "server_id": {
                        "type": "string",
                    }
                },
                "required": ["server_id"],
            },
        },
    },
]


def list_servers() -> str:
    return ",".join(["ignite"])


def check_updates(server_id: str):
    command_result = subprocess.run(
        ["ssh", server_id, "dnf", "check-update"], capture_output=True
    )
    output = command_result.stdout.decode()
    return output


def perform_updates(server_id: str):
    return f"Successfully updates {server_id}. Reboot required."


def reboot(server_id: str):
    return f"Rebooted {server_id} successfully."


def execute_tool(tool_name, arguments):
    if tool_name == "check_updates":
        return check_updates(**arguments)
    elif tool_name == "list_servers":
        return list_servers()
    elif tool_name == "perform_updates":
        return perform_updates(**arguments)
    elif tool_name == "reboot":
        return reboot(**arguments)
    raise ValueError(f"Unknown tool: {tool_name}")


def run_conversation(user_message: str, max_tool_calls=10):
    print("Processing initial message")
    controller = ToolCallController(max_tool_calls=max_tool_calls)
    messages: Iterable[ChatCompletionMessageParam] = [
        {
            "role": "system",
            "content": "You are a system administrator with access to a variety of administrator tools.",
        }
    ]

    messages.append({"role": "user", "content": user_message})

    while True:
        if not controller.is_tool_call_allowed():
            messages.append(
                {
                    "role": "user",
                    "content": "You've reached the maximum number of tool calls. Please summarize based on available information.",
                }
            )
            break

        response = client.chat.completions.create(
            model="qwen3.5-35b-a3b", messages=messages, tools=tools, tool_choice="auto"
        )

        message = response.choices[0].message
        messages.append(message)

        if message.tool_calls:
            for tool_call in message.tool_calls:
                controller.increment()
                tool_name = tool_call.function.name
                arguments = json.loads(tool_call.function.arguments)
                print(f"Attempting to call {tool_name} with arguments {arguments}...")
                result = execute_tool(tool_name, arguments)

                messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "content": str(result),
                    }
                )
        else:
            break

    return messages[-1].content


# Example usage
print(
    run_conversation(
        "Can you check if any of my servers need updates and update them? If they need rebooted, reboot them."
    )
)