import os
import sys
from typing import Dict, Any
import socket
import struct

# Ensure the plugin can import the StreamDeck library
sys.path.append(os.path.dirname(__file__))

try:
    from StreamDeck.DeviceManager import DeviceManager
    from StreamDeck.Device import Device
    from StreamDeck.Devices.DeviceManager import DeviceManager
except ImportError:
    print("StreamDeck library not found. Please install it using: pip install streamdeck")
    sys.exit(1)

class WakeOnLANPlugin:
    def __init__(self, streamdeck: Device):
        self.streamdeck = streamdeck
        # Store machine configurations
        self.machines: Dict[int, Dict[str, Any]] = {}

    def send_wol_packet(self, mac_address: str, broadcast_ip: str = '255.255.255.255', port: int = 9):
        """
        Send a Wake-on-LAN magic packet to wake up a machine
        
        :param mac_address: MAC address of the target machine (format: XX:XX:XX:XX:XX:XX)
        :param broadcast_ip: Broadcast IP to send the packet to (default is broadcast)
        :param port: UDP port for WoL packet (default is 9)
        """
        # Remove any colons or hyphens from the MAC address
        mac_address = mac_address.replace(':', '').replace('-', '')
        
        # Convert MAC address to bytes
        try:
            mac_bytes = bytes.fromhex(mac_address)
        except ValueError:
            print(f"Invalid MAC address: {mac_address}")
            return

        # Create the magic packet
        # Synchronization stream (6 * 0xFF)
        sync_stream = bytes([0xFF] * 6)
        
        # Repeat MAC address 16 times
        mac_repeated = mac_bytes * 16
        
        # Combine sync stream and repeated MAC
        magic_packet = sync_stream + mac_repeated

        # Create UDP socket
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
            sock.sendto(magic_packet, (broadcast_ip, port))
            print(f"Wake-on-LAN packet sent to {mac_address}")

    def add_machine(self, key_index: int, mac_address: str, broadcast_ip: str = '255.255.255.255'):
        """
        Configure a machine for a specific Stream Deck key
        
        :param key_index: Stream Deck key number to associate with the machine
        :param mac_address: MAC address of the machine
        :param broadcast_ip: Broadcast IP for the network (optional)
        """
        self.machines[key_index] = {
            'mac_address': mac_address,
            'broadcast_ip': broadcast_ip
        }

    def on_key_press(self, key, state):
        """
        Handle key press events
        
        :param key: The key that was pressed
        :param state: The state of the key (pressed/released)
        """
        if not state:  # Only trigger on key press down
            return

        # Check if the pressed key has a configured machine
        if key in self.machines:
            machine = self.machines[key]
            self.send_wol_packet(
                machine['mac_address'], 
                machine['broadcast_ip']
            )

def main():
    # Initialize the Stream Deck
    streamdecks = DeviceManager().enumerate()
    
    if not streamdecks:
        print("No Stream Deck found!")
        return

    # Use the first Stream Deck device
    deck = streamdecks[0]
    deck.open()
    deck.reset()

    # Create the WoL plugin
    wol_plugin = WakeOnLANPlugin(deck)

    # Example configuration - replace with your actual machine details
    # Add machines to specific keys
    wol_plugin.add_machine(0, '00:11:22:33:44:55')  # Key 0 wakes machine with this MAC
    wol_plugin.add_machine(1, 'AA:BB:CC:DD:EE:FF')  # Key 1 wakes another machine

    # Set up key press handling
    deck.set_key_callback(wol_plugin.on_key_press)

    # Keep the script running
    try:
        for i in range(deck.key_count()):
            # Optional: Set initial button images or colors
            # You could load custom images to indicate which buttons wake which machines
            pass

        # Keep the script running
        while True:
            deck.update()
    except KeyboardInterrupt:
        # Gracefully close the Stream Deck on interrupt
        deck.close()

if __name__ == "__main__":
    main()