#!/usr/bin/python3
import argparse
import os
import sys
import subprocess
import json

# Check if first arg is colour
waybar = {}
waybar["SELECTED"] = '<span underline="low">'
waybar["DEFAULT"] = '<span foreground="#ffffff">'
waybar["WHITE"]='<span foreground="#ffffff">'
waybar["RED"]='<span foreground="#ff0000">'
waybar["GREEN"]='<span foreground="#60d090">'
waybar["YELLOW"]='<span foreground="#ffff00">'
waybar["BLUE"]='<span foreground="#0000ff">'
waybar["NC"]='</span>'

shell = {}
shell["SELECTED"]='\033[0;32m'
shell["WHITE"]='\033[0;37m'
shell["RED"]='\033[0;31m'
shell["GREEN"]='\033[0;32m'
shell["YELLOW"]='\033[0;33m'
shell["BLUE"]='\033[0;34m'
shell["NC"]='\033[0m'



def get_info():
    state = subprocess.run(["hyprctl", "printstate"], capture_output=True, text=True)
    state = state.stdout.splitlines()
    state.pop(0)  # Remove the first line
    info = {}
    while True:
        if len(state) == 0:
            break

        line = state.pop(0)
        if line == "":
            continue
        if "-" not in line:
            break
        key, value = line.split(": ")
        key = key.strip().removeprefix("-").strip()
        value = value.strip() 
        desk_id = value
        info[desk_id] = {"name": key}
        for i in range(4):
            line = state.pop(0)
            if line == "":
                break
            line = line.split(": ")
            if len(line) == 2:
                key, value = line
                key = key.strip()
                value = value.strip()
                info[desk_id][key] = value
    return info

def get_icon(desk_name, focused=False):
    # Get the icon for the desk
    icon = ""
    desk_name = desk_name.strip().lower()
    if desk_name == "coding":
        return ""
    if desk_name == "web" or desk_name == "internet":
        return "󰖟"
    if desk_name == "social":
        return "󰭻"


    return "" if focused else ""

if __name__ == "__main__":
    # Check for arguments
    parser = argparse.ArgumentParser(description="hyprland desk script")    
    parser.add_argument("-d", "--desk", type=int, help="Desk number to switch to")
    parser.add_argument("-w", "--waybar", action="store_true", help="Waybar mode")
    parser.add_argument("-i", "--info", action="store_true", help="Info mode")
    parser.add_argument("-v", "--verbose", action="store_true", help="Verbose mode")
    args = parser.parse_args()

    if args.waybar or args.info:
        colours = waybar if args.waybar else shell
        # Get info from hyprctl printstate
        info = get_info()
        output = ""
        if args.desk:
            desk_num = args.desk
            if str(desk_num) not in info:
                print(f"Desk {desk_num} does not exist")
                sys.exit(1)
            desk = info[str(desk_num)]
            if args.verbose:
                print(desk)
            if desk["Populated"] == "false" and desk["Focused"] == "false":
                print(f"Desk {desk_num} is empty")
                sys.exit(1)

            if desk["Focused"] == "true":
                output += f"{colours['SELECTED']} {desk_num}: {get_icon(desk['name'],True)} {colours['NC']} "
            else:
                output += f"{colours["DEFAULT"]} {desk_num}: {get_icon(desk['name'])} {colours['NC']} "
            print(output)
            sys.exit(0)


        for i in range(1, int(max(info.keys()))+1):
            if str(i) not in info:
                continue
            desk = info[str(i)]
            if desk["Populated"] == "false" and desk["Focused"] == "false":
                continue
            if args.verbose:
                print(desk)

            if desk["Focused"] == "true":
                output += f"{colours['SELECTED']} {get_icon(desk['name'],True)} {colours['NC']} "
            else:
                output += f"{colours["DEFAULT"]} {get_icon(desk['name'])} {colours['NC']} "


        print(output)
        sys.exit(0)
    if args.desk:
        # Switch to desk
        if args.verbose:
            print(f"Switching to desk {args.desk}")
        
        subprocess.run(["hyprctl", "dispatch", "vdesk", str(args.desk)])
        sys.exit(0)