path: root/bot.py

import discord
from discord import app_commands
from discord.ext import commands
import random
import asyncio
from proxmoxer import ProxmoxAPI
import config
from utils.database import Database

bot = commands.Bot(command_prefix="!", intents=discord.Intents.default())
tree = bot.tree

# List of cogs to load
COGS = [
    "cogs.commissions",
    "cogs.vm_management",
    "cogs.alerts",
]


def get_proxmox():
    return ProxmoxAPI(
        config.PROXMOX_HOST,
        user=config.PROXMOX_USER,
        token_name=config.PROXMOX_TOKEN_NAME,
        token_value=config.PROXMOX_TOKEN_VALUE,
        verify_ssl=config.PROXMOX_VERIFY_SSL,
    )


async def load_cogs():
    """Load all cogs."""
    for cog in COGS:
        try:
            await bot.load_extension(cog)
            print(f"Loaded cog: {cog}")
        except Exception as e:
            print(f"Failed to load cog {cog}: {e}")


@bot.event
async def on_ready():
    print(f"Logged in as {bot.user} (ID: {bot.user.id})")

    # Initialize database tables
    try:
        await Database.init_tables()
        print("Database tables initialized")
    except Exception as e:
        print(f"Failed to initialize database: {e}")

    try:
        synced = await tree.sync()
        print(f"Synced {len(synced)} command(s)")
    except Exception as e:
        print(f"Failed to sync commands: {e}")


# ============ FUN COMMANDS ============

@tree.command(name="pet", description="Pet the bot!")
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.allowed_installs(guilds=True, users=True)
async def pet(interaction: discord.Interaction):
    responses = [
        "*purrs contentedly* :3",
        "*happy bot noises*",
        "ehe~ thank you!",
        "*wags tail enthusiastically*",
        "nyaa~ more pets please!",
        "*leans into the pets*",
        "beep boop... I mean, *purr*",
    ]
    await interaction.response.send_message(random.choice(responses))


@tree.command(name="boop", description="Boop the bot's nose")
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.allowed_installs(guilds=True, users=True)
async def boop(interaction: discord.Interaction):
    responses = [
        "*surprised pikachu face*",
        "hey! >:(",
        "*boops you back*",
        "my nose! D:",
        "owo what's this?",
        "*sneezes*",
    ]
    await interaction.response.send_message(random.choice(responses))


@tree.command(name="hug", description="Give the bot a hug")
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.allowed_installs(guilds=True, users=True)
async def hug(interaction: discord.Interaction):
    responses = [
        f"*hugs {interaction.user.display_name} back tightly*",
        "warm... cozy... :)",
        "*melts into the hug*",
        f"aww, {interaction.user.display_name}! *squeeze*",
    ]
    await interaction.response.send_message(random.choice(responses))


@tree.command(name="cookie", description="Give the bot a cookie")
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.allowed_installs(guilds=True, users=True)
async def cookie(interaction: discord.Interaction):
    await interaction.response.send_message("*nom nom nom* thank you for the cookie! :cookie:")


@tree.command(name="headpat", description="Give the bot headpats")
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.allowed_installs(guilds=True, users=True)
async def headpat(interaction: discord.Interaction):
    responses = [
        "*happy rumbling*",
        "yes yes yes!!",
        "*closes eyes contentedly*",
        "more... MORE!",
        ":D",
    ]
    await interaction.response.send_message(random.choice(responses))


# ============ SERVER ANALYTICS ============

@tree.command(name="node", description="Check full Proxmox node health")
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.allowed_installs(guilds=True, users=True)
async def node(interaction: discord.Interaction):
    await interaction.response.defer()
    try:
        proxmox = get_proxmox()
        node_status = proxmox.nodes(config.PROXMOX_NODE).status.get()

        mem_used = node_status["memory"]["used"] / (1024**3)
        mem_total = node_status["memory"]["total"] / (1024**3)
        swap_used = node_status.get("swap", {}).get("used", 0) / (1024**3)
        swap_total = node_status.get("swap", {}).get("total", 1) / (1024**3)
        rootfs_used = node_status.get("rootfs", {}).get("used", 0) / (1024**3)
        rootfs_total = node_status.get("rootfs", {}).get("total", 1) / (1024**3)

        embed = discord.Embed(
            title=f"Node: {config.PROXMOX_NODE}",
            color=discord.Color.blue()
        )
        embed.add_field(
            name="CPU",
            value=f"{node_status['cpu'] * 100:.1f}% ({node_status.get('cpuinfo', {}).get('cpus', '?')} cores)",
            inline=True
        )
        embed.add_field(
            name="Memory",
            value=f"{mem_used:.1f} / {mem_total:.1f} GB\n({mem_used/mem_total*100:.0f}%)",
            inline=True
        )
        embed.add_field(
            name="Swap",
            value=f"{swap_used:.1f} / {swap_total:.1f} GB",
            inline=True
        )
        embed.add_field(
            name="Root FS",
            value=f"{rootfs_used:.1f} / {rootfs_total:.1f} GB\n({rootfs_used/rootfs_total*100:.0f}%)",
            inline=True
        )
        embed.add_field(
            name="Uptime",
            value=f"{node_status['uptime'] // 86400}d {(node_status['uptime'] % 86400) // 3600}h {(node_status['uptime'] % 3600) // 60}m",
            inline=True
        )
        embed.add_field(
            name="Load Avg",
            value=f"{', '.join(f'{float(l):.2f}' for l in node_status.get('loadavg', [0,0,0]))}",
            inline=True
        )

        vms = proxmox.nodes(config.PROXMOX_NODE).qemu.get()
        lxcs = proxmox.nodes(config.PROXMOX_NODE).lxc.get()
        running_vms = sum(1 for vm in vms if vm.get("status") == "running")
        running_lxc = sum(1 for lxc in lxcs if lxc.get("status") == "running")

        embed.add_field(name="VMs", value=f"{running_vms}/{len(vms)} running", inline=True)
        embed.add_field(name="Containers", value=f"{running_lxc}/{len(lxcs)} running", inline=True)
        embed.set_footer(text=node_status.get("pveversion", ""))

        await interaction.followup.send(embed=embed)
    except Exception as e:
        await interaction.followup.send(f"Failed to fetch node status: {e}")


@tree.command(name="vms", description="List all VMs and containers")
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.allowed_installs(guilds=True, users=True)
async def vms(interaction: discord.Interaction):
    await interaction.response.defer()
    try:
        proxmox = get_proxmox()
        qemu_vms = proxmox.nodes(config.PROXMOX_NODE).qemu.get()
        lxcs = proxmox.nodes(config.PROXMOX_NODE).lxc.get()

        embed = discord.Embed(
            title=f"VMs & Containers on {config.PROXMOX_NODE}",
            color=discord.Color.blue()
        )

        if qemu_vms:
            vm_list = []
            for vm in sorted(qemu_vms, key=lambda x: x.get("vmid", 0)):
                status = ":green_circle:" if vm.get("status") == "running" else ":red_circle:"
                mem = vm.get("mem", 0) / (1024**3) if vm.get("status") == "running" else 0
                maxmem = vm.get("maxmem", 1) / (1024**3)
                vm_list.append(f"{status} **{vm.get('vmid')}** {vm.get('name', 'unnamed')} ({mem:.1f}/{maxmem:.0f}GB)")
            embed.add_field(name="VMs", value="\n".join(vm_list) or "None", inline=False)

        if lxcs:
            lxc_list = []
            for lxc in sorted(lxcs, key=lambda x: x.get("vmid", 0)):
                status = ":green_circle:" if lxc.get("status") == "running" else ":red_circle:"
                mem = lxc.get("mem", 0) / (1024**3) if lxc.get("status") == "running" else 0
                maxmem = lxc.get("maxmem", 1) / (1024**3)
                lxc_list.append(f"{status} **{lxc.get('vmid')}** {lxc.get('name', 'unnamed')} ({mem:.1f}/{maxmem:.0f}GB)")
            embed.add_field(name="Containers", value="\n".join(lxc_list) or "None", inline=False)

        await interaction.followup.send(embed=embed)
    except Exception as e:
        await interaction.followup.send(f"Failed to fetch VMs: {e}")


# ============ NOTIFY COMMAND ============

@tree.command(name="notify", description="Send a DM notification (owner only)")
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.allowed_installs(guilds=True, users=True)
@app_commands.describe(user="User to notify", message="Message to send")
async def notify(interaction: discord.Interaction, user: discord.User, message: str):
    if interaction.user.id != config.OWNER_ID:
        await interaction.response.send_message("You don't have permission to use this command.", ephemeral=True)
        return

    try:
        await user.send(message)
        await interaction.response.send_message(f"Sent notification to {user.display_name}", ephemeral=True)
    except discord.Forbidden:
        await interaction.response.send_message(f"Couldn't DM {user.display_name} - they may have DMs disabled.", ephemeral=True)
    except Exception as e:
        await interaction.response.send_message(f"Failed to send notification: {e}", ephemeral=True)


async def main():
    async with bot:
        await load_cogs()
        await bot.start(config.DISCORD_TOKEN)


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