#!/usr/bin/env python3

import asyncio
import json
import os
import sys

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()

    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": ""}

    async def monitors_refresh(self):
        try:
            proc = await asyncio.create_subprocess_exec(
                "hyprctl", "-j", "monitors",
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.DEVNULL
            )
            stdout, _ = await proc.communicate()
            if proc.returncode == 0:
                monitors = json.loads(stdout.decode())
                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 = {}

    async def clients_refresh(self):
        try:
            proc = await asyncio.create_subprocess_exec(
                "hyprctl", "-j", "clients",
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.DEVNULL
            )
            stdout, _ = await proc.communicate()
            if proc.returncode == 0:
                clients = json.loads(stdout.decode())
                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):
        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"]

                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})

                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):
        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()

    async def update_client(self, window_address):
        try:
            clean_addr = window_address.replace('0x', '')
            search_addr = f"0x{clean_addr}"
            
            proc = await asyncio.create_subprocess_exec(
                "hyprctl", "-j", "clients",
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.DEVNULL
            )
            stdout, _ = await proc.communicate()
            
            if proc.returncode == 0:
                clients = json.loads(stdout.decode())
                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']
                    }
                    return
            
            self.clients_cache.pop(clean_addr, None)
        except:
            await self.clients_refresh()

    async def run(self):
        await self.clients_refresh()
        await self.monitors_refresh()
        self.push_update()

        signature = os.environ.get('HYPRLAND_INSTANCE_SIGNATURE')
        socket_path = f"/run/user/{os.getuid()}/hypr/{signature}/.socket2.sock"

        while True:
            try:
                reader, writer = await asyncio.open_unix_connection(socket_path)
                
                while True:
                    line_bytes = await reader.readline()
                    if not line_bytes:
                        break
                    
                    line = line_bytes.decode('utf-8', errors='ignore').strip()
                    if ">>" not in line:
                        continue
                        
                    event, args = line.split('>>')
                    trigger = False
                    
                    if event in ["openwindow", "movewindow"]:
                        addr = args.split(',')[0].replace('0x', '')
                        await 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"]:
                        await 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 (asyncio.CancelledError, KeyboardInterrupt):
                break
            except Exception as e:
                await asyncio.sleep(2)

if __name__ == "__main__":
    try:
        asyncio.run(UnifiedWorkspaceManager().run())
    except KeyboardInterrupt:
        sys.exit(0)