#!/usr/bin/env python3
"""Refresh Lua geometry at login and after display/config changes."""
import argparse
import fcntl
import json
import os
from pathlib import Path
import socket
import subprocess
import sys
import time

config = Path(os.environ.get('XDG_CONFIG_HOME', Path.home() / '.config'))
root = config / 'hypr/lua'
sys.path.insert(0, str(root / 'bin'))
from hypr_state import read_state, atomic_write
from hypr_bridge import lua, dispatch, checked_run, require_ok
from display_geometry import window_rules


def read_settings():
    settings_path = config / 'quickshell/default/settings.txt'
    if not settings_path.exists():
        return {}
    return dict(line.split('=', 1) for line in settings_path.read_text().splitlines() if '=' in line)


def choose_primary(monitors, settings=None):
    settings = read_settings() if settings is None else settings
    preferred = settings.get('primaryMonitor') or read_state(root).get('primary_monitor')
    primary = next((m for m in monitors if m['name'] == preferred), None)
    primary = primary or next((m for m in monitors if m.get('focused')), monitors[0])
    return primary, settings


def shell_layers():
    """Monitors that currently carry a Quickshell layer surface (the taskbar)."""
    result = subprocess.run(['/usr/bin/hyprctl', 'layers', '-j'], capture_output=True, text=True)
    if result.returncode:
        return None
    try:
        layers = json.loads(result.stdout)
    except json.JSONDecodeError:
        return None
    return {name for name, data in layers.items()
            if any(item.get('namespace') == 'quickshell' for items in data.get('levels', {}).values() for item in items)}


def shell_is_placed(monitors, with_shell, preferred):
    """The taskbar belongs on the primary monitor. If that monitor is gone, any monitor will do."""
    names = {m['name'] for m in monitors}
    if preferred in names:
        return preferred in with_shell
    return bool(with_shell & names)


_last_restart = 0.0


def ensure_shell():
    """After a monitor event: bring the Quickshell panel back if it vanished or sits on the wrong monitor.

    Quickshell keeps running when its monitor is unplugged but loses the panel's layer surface and does not
    create a new one when the monitor returns. A clean restart is the reliable fix.
    """
    global _last_restart
    result = subprocess.run(['/usr/bin/hyprctl', 'monitors', '-j'], capture_output=True, text=True)
    try:
        monitors = json.loads(result.stdout)
    except json.JSONDecodeError:
        return
    if not monitors:
        return
    with_shell = shell_layers()
    if with_shell is None:
        return
    primary, _ = choose_primary(monitors)
    if shell_is_placed(monitors, with_shell, primary['name']):
        return
    if time.monotonic() - _last_restart < 20:
        return
    _last_restart = time.monotonic()
    subprocess.Popen(['windoesnot-restart-shell'], start_new_session=True,
                     stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def refresh():
    result = subprocess.run(['/usr/bin/hyprctl', 'monitors', '-j'], capture_output=True, text=True)
    if result.returncode:
        return
    try:
        monitors = json.loads(result.stdout)
    except json.JSONDecodeError:
        return
    if not monitors:
        return
    primary, settings = choose_primary(monitors)
    geometry = window_rules(primary, settings)
    content = '-- Generated from current monitor geometry; do not edit.\nreturn ' + lua(geometry) + '\n'
    path = root / 'geometry.lua'
    if path.exists() and path.read_text() == content:
        return
    atomic_write(path, content)
    subprocess.run(['/usr/bin/hyprctl', 'reload'], check=True, capture_output=True)
    # Static placement rules only affect the next map. Also fit already-open
    # shell windows after a scale/orientation change, in global logical pixels.
    titles = {'quickshell-' + name: name + '-config' for name in
              ['startmenu','settings','calendar','audio','bluetooth','network','nearby','docker','ollama','mediaplayer','colorpalette','search']}
    titles['quickshell-does-pinned'] = 'copilot-config-pinned'
    titles['quickshell-copilot-pinned'] = 'copilot-config-pinned'
    clients = json.loads(subprocess.check_output(['/usr/bin/hyprctl','clients','-j'], text=True))
    for client in clients:
        rule = geometry['rules'].get(titles.get(client.get('title')))
        if not rule:
            continue
        selector = 'address:' + client['address']
        w, h = rule['size']
        x, y = rule['move']
        commands = [('resizewindowpixel', f'exact {w} {h},{selector}'),
                    ('movewindowpixel', f'exact {primary["x"] + x} {primary["y"] + y},{selector}')]
        for name, arguments in commands:
            require_ok(checked_run(['eval', 'return hl.dispatch(' + dispatch(name, arguments) + ')']))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--watch', action='store_true')
    args = parser.parse_args()
    runtime = Path(os.environ['XDG_RUNTIME_DIR'])
    signature = os.environ['HYPRLAND_INSTANCE_SIGNATURE']
    with (runtime / 'windoesnot-layout.lock').open('a') as lock:
        try:
            fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            return
        refresh()
        if not args.watch:
            return
        connection = socket.socket(socket.AF_UNIX)
        connection.connect(str(runtime / 'hypr' / signature / '.socket2.sock'))
        stream = connection.makefile('r')
        last_check = 0.0
        for event in stream:
            name = event.split('>>', 1)[0]
            if name in {'monitoradded','monitoraddedv2','monitorremoved','monitorremovedv2','configreloaded'}:
                # The compositor may need a moment to settle output/reservation state.
                time.sleep(0.15)
                refresh()
            if name in {'monitoradded','monitoraddedv2','monitorremoved','monitorremovedv2'} and time.monotonic() - last_check > 8:
                # Let Quickshell react first (it often recovers alone), then check once; a cable plugged in
                # produces a burst of events, so only the first one of a burst waits and checks.
                time.sleep(3)
                ensure_shell()
                last_check = time.monotonic()

if __name__ == '__main__':
    main()
