#Pscreens v2.1 - tool for quickly changing resolutions, turning screens on/off, in a handy terminal
#!/usr/bin/env python3
import subprocess
import json

def run_command(command):
    try:
        output = subprocess.check_output(command, shell=True, text=True)
        return output.strip()
    except subprocess.CalledProcessError as e:
        print(f"Command failed: {e}")
        return ""

def list_monitors():
    command = "wlr-randr"
    output = run_command(command)
    if output:
        monitors = output.split('\n')
        print("\nAvailable Monitors:")
        for i, monitor in enumerate(monitors, 1):
            print(f"{i}: {monitor}")
        return monitors
    else:
        print("\nFailed to list monitors.")
        return []

def toggle_monitor():
    command = "swaymsg -t get_outputs"
    output = run_command(command)
    if output:
        monitors = json.loads(output)
        print("\nAvailable Monitors:")
        for i, monitor in enumerate(monitors, 1):
            name = monitor.get("name")
            make = monitor.get("make", "Unknown")
            model = monitor.get("model", "Unknown")
            active = "Active" if monitor.get("active") else "Inactive"
            print(f"{i}: {name} - {make} {model} [{active}]")
        
        monitor_index = int(input("\nEnter the monitor number to toggle (on/off): ")) - 1
        if 0 <= monitor_index < len(monitors):
            monitor = monitors[monitor_index]["name"]
            status = input("Enter 'on' to enable or 'off' to disable the monitor: ").strip().lower()
            if status == "off":
                command = f"wlr-randr --output {monitor} --off"
            elif status == "on":
                command = f"wlr-randr --output {monitor} --auto"
            else:
                print("\nInvalid input. Please enter 'on' or 'off'.")
                return
            result = run_command(command)
            if result == "":
                print(f"\nMonitor {monitor} turned {status} successfully.")
            else:
                print(f"\nFailed to toggle monitor {monitor}.")
        else:
            print("\nInvalid monitor number.")
    else:
        print("\nFailed to list monitors.")

def list_available_resolutions():
    monitors = list_monitors()
    if not monitors:
        return
    monitor_index = int(input("\nEnter the monitor number to list resolutions: ")) - 1
    if 0 <= monitor_index < len(monitors):
        monitor = monitors[monitor_index].split()[0]
        command = f"wlr-randr --output {monitor}"
        output = run_command(command)
        if output:
            resolutions = output.split('\n')
            print(f"\nAvailable Resolutions for {monitor}:")
            for i, resolution in enumerate(resolutions, 1):
                print(f"{i}: {resolution}")
        else:
            print(f"\nFailed to list resolutions for {monitor}.")
    else:
        print("\nInvalid monitor number.")
#kind of sloppy but it works!
def change_resolution():
    monitors = list_monitors()
    if not monitors:
        return
    monitor_index = int(input("\nEnter the monitor number to change resolution: ")) - 1
    if 0 <= monitor_index < len(monitors):
        monitor = monitors[monitor_index].split()[0]
        command = f"wlr-randr --output {monitor}"
        output = run_command(command)
        if output:
            resolutions = output.split('\n')
            print(f"\nAvailable Resolutions for {monitor}:")
            for i, resolution in enumerate(resolutions, 1):
                print(f"{i}: {resolution}")
            resolution_index = int(input("\nEnter the desired resolution number: ")) - 1
            if 0 <= resolution_index < len(resolutions):
                resolution = resolutions[resolution_index].split()[0]
                command = f"wlr-randr --output {monitor} --mode {resolution}"
                result = run_command(command)
                if result == "":
                    print(f"\nResolution for {monitor} changed to {resolution} successfully.")
                else:
                    print(f"\nFailed to change resolution for {monitor}.")
            else:
                print("\nInvalid resolution number.")
        else:
            print(f"\nFailed to list resolutions for {monitor}.")
    else:
        print("\nInvalid monitor number.")
#So this works, but.. it ends up being virtually the same as the resolution option
def choose_refresh_rate():
    monitors = list_monitors()
    if not monitors:
        return
    monitor_index = int(input("\nEnter the monitor number to change refresh rate: ")) - 1
    if 0 <= monitor_index < len(monitors):
        monitor = monitors[monitor_index].split()[0]
        command = f"wlr-randr --output {monitor}"
        output = run_command(command)
        if output:
            modes = output.split('\n')
            print(f"\nAvailable Modes for {monitor} (including refresh rates):")
            for i, mode in enumerate(modes, 1):
                print(f"{i}: {mode}")
            mode_index = int(input("\nEnter the desired mode number: ")) - 1
            if 0 <= mode_index < len(modes):
                mode = modes[mode_index].split()[0]
                command = f"wlr-randr --output {monitor} --mode {mode}"
                result = run_command(command)
                if result == "":
                    print(f"\nMode for {monitor} changed to {mode} successfully.")
                else:
                    print(f"\nFailed to change mode for {monitor}.")
            else:
                print("\nInvalid mode number.")
        else:
            print(f"\nFailed to list modes for {monitor}.")
    else:
        print("\nInvalid monitor number.")

def choose_multi_display_mode():
    print("\nMulti-Display Modes:")
    print("1: Mirrored")
    print("2: Extended")
    mode_choice = input("\nEnter multi-display mode number: ")
    if mode_choice == "1":
        command = "wlr-randr --mirror"
    elif mode_choice == "2":
        command = "wlr-randr --auto"
    else:
        print("\nInvalid mode selected.")
        return
    result = run_command(command)
    if result == "":
        mode = "mirrored" if mode_choice == "1" else "extended"
        print(f"\nMulti-display mode set to {mode} successfully.")
    else:
        print(f"\nFailed to set multi-display mode.")

def main():
    while True:
        print("\nMonitor Configuration Menu:")
        print("1: List Monitors")
        print("2: Toggle Monitor On/Off")
        print("3: List Available Resolutions")
        print("4: Change Resolution")
        print("5: Choose Refresh Rate")
        print("6: Choose Multi-Display Mode")
        print("7: Exit")

        choice = input("\nEnter your choice: ")

        if choice == '1':
            list_monitors()
        elif choice == '2':
            toggle_monitor()
        elif choice == '3':
            list_available_resolutions()
        elif choice == '4':
            change_resolution()
        elif choice == '5':
            choose_refresh_rate()
        elif choice == '6':
            choose_multi_display_mode()
        elif choice == '7':
            break
        else:
            print("\nInvalid choice. Please try again.")

if __name__ == "__main__":
    main()