import serial
import time

class WLEDController:
    def __init__(self, port: str, baud_rate: int = 115200, timeout: float = 1.0):
        """
        Initialize the WLEDController.

        :param port: The serial port to which the ESP32 is connected (e.g., '/dev/ttyUSB0').
        :param baud_rate: The baud rate for the serial connection.
        :param timeout: The timeout for serial communication in seconds.
        """
        self.port = port
        self.baud_rate = baud_rate
        self.timeout = timeout
        self.connection = None

    def connect(self):
        """Establish the serial connection."""
        try:
            self.connection = serial.Serial(
                self.port, self.baud_rate, timeout=self.timeout
            )
            print(f"Connected to {self.port}")
        except serial.SerialException as e:
            print(f"Failed to connect to {self.port}: {e}")

    def disconnect(self):
        """Close the serial connection."""
        if self.connection and self.connection.is_open:
            self.connection.close()
            print(f"Disconnected from {self.port}")

    def send_command(self, command: str):
        """
        Send a command to the WLED device.

        :param command: The JSON command string to send.
        """
        if not self.connection or not self.connection.is_open:
            raise ConnectionError("Serial connection is not open.")
        try:
            self.connection.write(command.encode('utf-8'))
            print(f"Sent command: {command}")
        except serial.SerialException as e:
            print(f"Failed to send command: {e}")

    def switch_to_red(self):
        """Switch the light to red."""
        self.send_command('{"ps":"1"}')

    def switch_to_yellow(self):
        """Switch the light to yellow."""
        self.send_command('{"ps":"2"}')

    def switch_to_green(self):
        """Switch the light to green."""
        self.send_command('{"ps":"3"}')

    def turn_off_all(self):
        """Turn off all lights."""
        self.send_command('{"ps":"4"}')

    def change_effect(self):
        """Change the color effect."""
        self.send_command('{"ps":"5"}')

# Example usage:
if __name__ == "__main__":
    controller = WLEDController(port="/dev/ttyUSB0")
    controller.connect()
    
    try:
        controller.switch_to_red()
        time.sleep(2)  # Wait for 2 seconds

        controller.switch_to_yellow()
        time.sleep(2)  # Wait for 2 seconds

        controller.switch_to_green()
        time.sleep(2)  # Wait for 2 seconds

        controller.change_effect()
        time.sleep(2)  # Wait for 2 seconds

        controller.turn_off_all()
        time.sleep(2)  # Wait for 2 seconds

    finally:
        controller.disconnect()