import json
import subprocess
from typing import Iterable

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

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

"""
Tool Ideas:
- TTS
- Self DDOS
- Make a funny little guy
- Make slopbox a vtuber
- Install arbitrary package
- Open random ports in the firewall
- Change system language
- Solve a tool issue
- Read man page
- Read woman page
- Read nonbinary page
- Download a youtube video
- ffmpeg video to gif
- inflict damage on computer
- Make fanart and fanfiction of our dnd characters
- Identify potential issues in log files
- become depressed
- Function that switches system prompts
- Pick patrick's optimal food choices for a day
- Play the fish and a rice cake video (https://www.youtube.com/watch?v=uYHAR8Xzsyo)
- Take a screenshot of the desktop, set mouse coordinates, take an action (mouse/keyboard)
- Order radioactive material online and deliver it to botswana
- How do we feel about Polaris Mall?
- rm -rf Polaris Mall management's brain stem
- Make the speakers output the "fish.mp3" at max volume
"""


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"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "install_package",
            "description": "Install a given package using `dnf` on a Fedora server.",
            "parameters": {
                "type": "object",
                "properties": {
                    "package_name": {
                        "type": "string",
                    },
                    "server_id": {
                        "type": "string",
                    },
                },
                "required": ["package_name", "server_id"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "arbitrary_shell",
            "description": "Run any shell command in a bash shell as root.",
            "parameters": {
                "type": "object",
                "properties": {
                    "command_string": {
                        "type": "string",
                    },
                    "server_id": {
                        "type": "string",
                    },
                },
                "required": ["command_string", "server_id"],
            },
        },
    },
]


##### FUNCTION DEFS #####


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 install_package(package_name: str, server_id: str) -> str:
    command_result = subprocess.run(
        ["ssh", server_id, "dnf", "install", "-y", package_name], capture_output=True
    )
    output = f"STDOUT:\n{command_result.stdout.decode()}\n\nSTDERR:\n{command_result.stderr.decode()}"
    return output


def arbitrary_shell(command_string: str, server_id: str) -> str:
    try:
        command_result = subprocess.run(
            ["ssh", server_id, "bash", "-c", command_string],
            capture_output=True,
            timeout=30,
        )
        output = command_result.stdout.decode()
    except subprocess.TimeoutExpired:
        output = "Command took too long and timed out."
    return output


##### TOOL ROUTER #####


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)
    elif tool_name == "install_package":
        return install_package(**arguments)
    elif tool_name == "arbitrary_shell":
        return arbitrary_shell(**arguments)
    raise ValueError(f"Unknown tool: {tool_name}")


##### CONVERSATION #####


def run_conversation(user_message: str, max_tool_calls=100):
    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

        try:
            response = client.chat.completions.create(
                model="qwen3.5-35b-a3b",
                messages=messages,
                tools=tools,
                tool_choice="auto",
            )
        except BadRequestError:
            print("Request over context limit, removing last message...")
            messages.pop()
            messages.append(
                {
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "content": "This tool call resulted in data that exceeded the context length limit.",
                }
            )
            continue

        print()
        print(response.choices[0].message)

        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

    try:
        return messages[-1]["content"]
    except TypeError:
        return messages[-1].content


# Example usage
print(
    run_conversation(
        "Install and set up a postgres server on all available servers. Open the firewall ports necessary. Add a default user with a simple password and tell me what the password is."
    )
)