import discord
from discord import app_commands
from discord.ext import commands
from typing import Literal
import asyncio
from proxmoxer import ProxmoxAPI
import config


class VMManagement(commands.Cog):
    """VM and LXC container management commands."""

    def __init__(self, bot: commands.Bot):
        self.bot = bot

    def get_proxmox(self) -> ProxmoxAPI:
        """Get a Proxmox API connection."""
        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 _check_owner(self, interaction: discord.Interaction) -> bool:
        """Check if user is owner, send error if not."""
        if interaction.user.id != config.OWNER_ID:
            await interaction.response.send_message(
                "You don't have permission to use this command.", ephemeral=True
            )
            return False
        return True

    def _get_vm_type(self, proxmox: ProxmoxAPI, vmid: int) -> str | None:
        """Determine if VMID is a QEMU VM or LXC container."""
        # Check QEMU VMs
        try:
            vms = proxmox.nodes(config.PROXMOX_NODE).qemu.get()
            for vm in vms:
                if vm.get("vmid") == vmid:
                    return "qemu"
        except Exception:
            pass

        # Check LXC containers
        try:
            lxcs = proxmox.nodes(config.PROXMOX_NODE).lxc.get()
            for lxc in lxcs:
                if lxc.get("vmid") == vmid:
                    return "lxc"
        except Exception:
            pass

        return None

    def _get_vm_name(self, proxmox: ProxmoxAPI, vmid: int, vm_type: str) -> str:
        """Get the name of a VM or container."""
        try:
            if vm_type == "qemu":
                config_data = proxmox.nodes(config.PROXMOX_NODE).qemu(vmid).config.get()
            else:
                config_data = proxmox.nodes(config.PROXMOX_NODE).lxc(vmid).config.get()
            return config_data.get("hostname", config_data.get("name", f"VM {vmid}"))
        except Exception:
            return f"VM {vmid}"

    vm_group = app_commands.Group(
        name="vm",
        description="Manage VMs and containers (owner only)",
        allowed_contexts=app_commands.AppCommandContext(guild=True, dm_channel=True, private_channel=True),
        allowed_installs=app_commands.AppInstallationType(guild=True, user=True),
    )

    @vm_group.command(name="start", description="Start a VM or container")
    @app_commands.describe(vmid="The VMID to start")
    async def start_vm(self, interaction: discord.Interaction, vmid: int):
        if not await self._check_owner(interaction):
            return

        await interaction.response.defer(ephemeral=True)

        try:
            proxmox = self.get_proxmox()
            vm_type = self._get_vm_type(proxmox, vmid)

            if not vm_type:
                await interaction.followup.send(f"VM/Container {vmid} not found.", ephemeral=True)
                return

            vm_name = self._get_vm_name(proxmox, vmid, vm_type)

            # Start the VM/container
            if vm_type == "qemu":
                await asyncio.to_thread(
                    lambda: proxmox.nodes(config.PROXMOX_NODE).qemu(vmid).status.start.post()
                )
            else:
                await asyncio.to_thread(
                    lambda: proxmox.nodes(config.PROXMOX_NODE).lxc(vmid).status.start.post()
                )

            type_name = "VM" if vm_type == "qemu" else "Container"
            await interaction.followup.send(
                f":green_circle: {type_name} **{vmid}** ({vm_name}) is starting...", ephemeral=True
            )

        except Exception as e:
            await interaction.followup.send(f"Failed to start VM {vmid}: {e}", ephemeral=True)

    @vm_group.command(name="stop", description="Stop a VM or container")
    @app_commands.describe(vmid="The VMID to stop")
    async def stop_vm(self, interaction: discord.Interaction, vmid: int):
        if not await self._check_owner(interaction):
            return

        await interaction.response.defer(ephemeral=True)

        try:
            proxmox = self.get_proxmox()
            vm_type = self._get_vm_type(proxmox, vmid)

            if not vm_type:
                await interaction.followup.send(f"VM/Container {vmid} not found.", ephemeral=True)
                return

            vm_name = self._get_vm_name(proxmox, vmid, vm_type)

            # Stop the VM/container
            if vm_type == "qemu":
                await asyncio.to_thread(
                    lambda: proxmox.nodes(config.PROXMOX_NODE).qemu(vmid).status.stop.post()
                )
            else:
                await asyncio.to_thread(
                    lambda: proxmox.nodes(config.PROXMOX_NODE).lxc(vmid).status.stop.post()
                )

            type_name = "VM" if vm_type == "qemu" else "Container"
            await interaction.followup.send(
                f":red_circle: {type_name} **{vmid}** ({vm_name}) is stopping...", ephemeral=True
            )

        except Exception as e:
            await interaction.followup.send(f"Failed to stop VM {vmid}: {e}", ephemeral=True)

    @vm_group.command(name="restart", description="Restart a VM or container")
    @app_commands.describe(vmid="The VMID to restart")
    async def restart_vm(self, interaction: discord.Interaction, vmid: int):
        if not await self._check_owner(interaction):
            return

        await interaction.response.defer(ephemeral=True)

        try:
            proxmox = self.get_proxmox()
            vm_type = self._get_vm_type(proxmox, vmid)

            if not vm_type:
                await interaction.followup.send(f"VM/Container {vmid} not found.", ephemeral=True)
                return

            vm_name = self._get_vm_name(proxmox, vmid, vm_type)

            # Reboot the VM/container
            if vm_type == "qemu":
                await asyncio.to_thread(
                    lambda: proxmox.nodes(config.PROXMOX_NODE).qemu(vmid).status.reboot.post()
                )
            else:
                await asyncio.to_thread(
                    lambda: proxmox.nodes(config.PROXMOX_NODE).lxc(vmid).status.reboot.post()
                )

            type_name = "VM" if vm_type == "qemu" else "Container"
            await interaction.followup.send(
                f":arrows_counterclockwise: {type_name} **{vmid}** ({vm_name}) is restarting...", ephemeral=True
            )

        except Exception as e:
            await interaction.followup.send(f"Failed to restart VM {vmid}: {e}", ephemeral=True)


async def setup(bot: commands.Bot):
    await bot.add_cog(VMManagement(bot))