import json
import subprocess
import os
import sys
import socket
import threading
import time

class UnifiedWorkspaceManager:
    def __init__(self):
        self.last_sent_json = ""
        self.urgent_window_addresses = set()
        self.script_directory = os.path.dirname(os.path.realpath(__file__))
        
        self.clients_cache = {}
        self.monitors_cache = {}
        
        self.load_icon_configuration()
        self.clients_refresh()
        self.monitors_refresh()

    def load_icon_configuration(self):
        # тянем иконки из конфига
        config_path = os.path.join(self.script_directory, "workspaces_map.json")
        try:
            with open(config_path, 'r') as file:
                self.icon_map = json.load(file)
        except:
            self.icon_map = {"default": "", "no-icon": ""}

    def monitors_refresh(self):
        # обновляем инфу по мониторам
        try:
            raw = subprocess.check_output(["hyprctl", "-j", "monitors"], stderr=subprocess.DEVNULL)
            monitors = json.loads(raw)
            self.monitors_cache = {
                m['name']: {
                    "id": m['id'],
                    "active_workspace_id": m['activeWorkspace']['id'],
                    "focused": m.get('focused', False)
                } for m in monitors
            }
        except:
            self.monitors_cache = {}

    def clients_refresh(self):
        # полный сброс и пересборка кэша окон
        try:
            raw = subprocess.check_output(["hyprctl", "-j", "clients"], stderr=subprocess.DEVNULL)
            clients = json.loads(raw)
            self.clients_cache = {
                c['address'].strip().replace('0x', ''): {
                    "class": c['class'],
                    "workspace_id": c['workspace']['id'],
                    "monitor_id": c['monitor']
                } for c in clients
            }
        except:
            self.clients_cache = {}

    def get_current_state_json(self):
        # собираем финальный json
        try:
            output = {}

            for mon_name, mon_data in self.monitors_cache.items():
                m_id = mon_data["id"]
                active_id = mon_data["active_workspace_id"]
                focused = mon_data["focused"]

                # готовим 10 слотов
                workspaces = {
                    str(i): {"exists": False, "active": False, "focused": False, "urgent": False, "icons": []} 
                    for i in range(1, 11)
                }
                
                # раскидываем окна по воркспейсам
                for addr, client in self.clients_cache.items():
                    if client['monitor_id'] == m_id:
                        ws_id = str(client['workspace_id'])
                        if ws_id in workspaces:
                            slot = workspaces[ws_id]
                            slot["exists"] = True
                            if addr in self.urgent_window_addresses:
                                slot["urgent"] = True
                            
                            icon = self.icon_map.get(client['class'], self.icon_map.get("default", ""))
                            if icon not in slot["icons"]:
                                slot["icons"].append(icon)

                # помечаем текущий активный воркспейс
                curr_ws = str(active_id)
                if curr_ws in workspaces:
                    workspaces[curr_ws].update({"active": True, "focused": focused, "exists": True})

                # рисуем no-icon
                for ws in workspaces.values():
                    if ws["exists"] and not ws["icons"]:
                        ws["icons"] = [self.icon_map.get("no-icon", "")]
                
                output[mon_name.replace("-", "_")] = workspaces

            return json.dumps(output)
        except:
            return "{}"

    def push_update(self):
        # отдаём json в stdout при изменении
        current = self.get_current_state_json()
        if current != self.last_sent_json:
            self.last_sent_json = current
            sys.stdout.write(current + '\n')
            sys.stdout.flush()

    def update_client(self, window_address):
        # обновляем окно
        try:
            clean_addr = window_address.replace('0x', '')
            search_addr = f"0x{clean_addr}"
            
            raw = subprocess.check_output(["hyprctl", "-j", "clients"], stderr=subprocess.DEVNULL)
            clients = json.loads(raw)
            
            target = next((c for c in clients if c['address'] == search_addr), None)
            
            if target:
                self.clients_cache[clean_addr] = {
                    "class": target.get('class', 'unknown'),
                    "workspace_id": target['workspace']['id'],
                    "monitor_id": target['monitor']
                }
            else:
                self.clients_cache.pop(clean_addr, None)
                
        except:
            # если что-то отвалилось, то просто пересобираем всё
            self.clients_refresh()

    def run(self):
        self.push_update()

        def event_socket_listener():
            # слушаем сокет хайпера
            signature = os.environ.get('HYPRLAND_INSTANCE_SIGNATURE')
            socket_path = f"/run/user/{os.getuid()}/hypr/{signature}/.socket2.sock"

            while True:
                try:
                    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
                        sock.connect(socket_path)
                        buf = ""
                        while True:
                            data = sock.recv(4096).decode('utf-8', errors='ignore')
                            if not data: break
                            buf += data
                            
                            if '\n' in buf:
                                lines = buf.split('\n')
                                buf = lines.pop()
                                
                                trigger = False
                                for line in lines:
                                    if ">>" not in line: continue
                                    event, args = line.split('>>')
                                    
                                    # ивенты с окнами
                                    if event in ["openwindow", "movewindow"]:
                                        addr = args.split(',')[0].replace('0x', '')
                                        self.update_client(addr)
                                        trigger = True
                                    
                                    elif event == "closewindow":
                                        addr = args.replace('0x', '')
                                        self.clients_cache.pop(addr, None)
                                        self.urgent_window_addresses.discard(addr)
                                        trigger = True

                                    # ивенты с мониторами
                                    elif event in ["workspace", "workspacev2", "monitoradded", "monitorremoved"]:
                                        self.monitors_refresh()
                                        trigger = True

                                    elif event in ["focusedmon", "focusedmonv2"]:
                                        mon_name = args.split(',')[0]
                                        for name in self.monitors_cache:
                                            self.monitors_cache[name]["focused"] = (name == mon_name)
                                        trigger = True

                                    # алерты
                                    elif event == "urgent":
                                        self.urgent_window_addresses.add(args.replace('0x', ''))
                                        trigger = True
                                        
                                    elif event in ["activewindow", "activewindowv2"]:
                                        addr = args.split(',')[0].replace('0x', '')
                                        if addr in self.urgent_window_addresses:
                                            self.urgent_window_addresses.remove(addr)
                                            trigger = True

                                if trigger:
                                    self.push_update()
                except:
                    time.sleep(2) # реконнект, если хайпер отвалился

        # слушаем поток в фоне
        threading.Thread(target=event_socket_listener, daemon=True).start()
        
        try:
            while True: time.sleep(100)
        except KeyboardInterrupt: sys.exit(0)

if __name__ == "__main__":
    UnifiedWorkspaceManager().run()