#!/usr/bin/env python3
"""Control Gnoblin, its windows and live shell settings."""
import argparse
import json
import os
import re
from pathlib import Path
import shutil
import socket
import subprocess
import sys
import time
import uuid


class CommandError(Exception):
    pass


def number(low, high):
    def parse(value):
        try: result = int(value)
        except ValueError: raise argparse.ArgumentTypeError("must be an integer") from None
        if not low <= result <= high: raise argparse.ArgumentTypeError(f"must be between {low} and {high}")
        return result
    return parse


def parser():
    common = argparse.ArgumentParser(add_help=False)
    common.add_argument("-j", "--json", dest="format", action="store_const", const="json", default=argparse.SUPPRESS,
                        help="Print JSON, including when used in a terminal")
    common.add_argument("--format", choices=["auto", "json", "table"], default=argparse.SUPPRESS,
                        help="Output format (auto: tables in terminals, JSON in pipes)")
    common.add_argument("--timeout", type=number(1, 60), default=argparse.SUPPRESS, metavar="SECONDS", help="Request timeout (default: 5)")
    common.add_argument("--socket", default=argparse.SUPPRESS, metavar="PATH", help="Compositor socket path")
    root = argparse.ArgumentParser(prog="gnoblinctl", description=__doc__, parents=[common],
        epilog="Examples: gnoblinctl windows | gnoblinctl window focus 42 | gnoblinctl features --json\nUse COMMAND --help for arguments and examples.",
        formatter_class=argparse.RawDescriptionHelpFormatter)
    commands = root.add_subparsers(dest="command", metavar="COMMAND")

    def add(parent, name, help_text, **kwargs):
        return parent.add_parser(name, help=help_text, description=help_text, parents=[common], **kwargs)

    for name, description in [("ping", "Check the shell connection"), ("version", "Show shell and protocol version"),
            ("status", "Show compositor and shell status"), ("reload", "Soft-reload the shell, keeping open windows"),
            ("reload-config", "Reload gnoblin.toml"), ("features", "List shell feature switches"),
            ("extensions", "List GNOME extensions"), ("scripts", "List loaded user scripts"),
            ("reload-scripts", "Reload user scripts"), ("input-sources", "List configured keyboard sources"),
            ("input-source", "Show the active keyboard source"), ("privacy", "Show privacy indicators"),
            ("portal-grants", "List persistent portal permissions"), ("launch-state", "Show pending app launches"),
            ("monitors", "List logical monitors and geometry")]:
        add(commands, name, description)
    for name, description, arg in [("feature", "Show one feature switch", "id"), ("enable", "Enable a shell feature", "id"),
            ("disable", "Disable a shell feature", "id"), ("reload-ext", "Reload one GNOME extension", "uuid"),
            ("launch-end", "End a busy-cursor request", "token")]:
        add(commands, name, description).add_argument(arg)
    source = add(commands, "set-input-source", "Select a configured keyboard source")
    source.add_argument("type", help="Source type from input-sources, such as xkb or ibus")
    source.add_argument("id", help="Exact source ID")
    begin = add(commands, "launch-begin", "Start a global busy-cursor request")
    begin.add_argument("token")
    begin.add_argument("application", help="Desktop application ID")
    begin.add_argument("milliseconds", nargs="?", default=3000, type=number(1, 60000))
    revoke = add(commands, "revoke-grant", "Revoke one persistent portal permission")
    revoke.add_argument("kind", choices=["screen-cast", "remote-desktop"])
    revoke.add_argument("id")

    def window_filters(command):
        command.add_argument("--app-id", help="Match an exact application ID")
        command.add_argument("--title", help="Match title text, without case sensitivity")
        command.add_argument("--focused", action="store_true", help="Only show the focused window")
    window_filters(add(commands, "windows", "List open windows, IDs and geometry"))
    window = add(commands, "window", "Manage a window by ID; active means the focused window")
    actions = window.add_subparsers(dest="action", required=True, metavar="ACTION")
    window_filters(add(actions, "list", "List open windows"))
    for name in ("focus", "close", "minimize", "restore-or-minimize", "restore", "maximize", "unmaximize", "fullscreen", "unfullscreen"):
        action = add(actions, name, f"{name.capitalize()} one window")
        action.add_argument("id", nargs="?", default="active", help="Window ID from windows (default: active)")
    for name, fields in (("move", ["x", "y"]), ("resize", ["width", "height"]), ("workspace", ["workspace"]), ("monitor", ["monitor"])):
        action = add(actions, name, {"move": "Set position in logical screen coordinates", "resize": "Set frame size in logical pixels",
            "workspace": "Move to an existing workspace (one-based)", "monitor": "Move to a logical monitor (zero-based)"}[name])
        action.add_argument("id", help="Window ID, or active")
        for field in fields:
            bounds = (1, 32768) if field in ("width", "height") else (1, 1024) if field == "workspace" else (0, 1024) if field == "monitor" else (-100000, 100000)
            action.add_argument(field, type=number(*bounds))
    workspaces = add(commands, "workspaces", "List workspaces or switch to an existing workspace")
    workspaces.add_argument("action", nargs="?", default="list", choices=["list", "switch"])
    workspaces.add_argument("workspace", nargs="?", type=number(1, 1024), help="One-based workspace ID")
    completion = add(commands, "completion", "Print shell completion setup")
    completion.add_argument("shell", choices=["bash", "zsh", "fish"])
    return root


def dbus(method, signature="", arguments=(), service="org.gnoblin.Shell", timeout=5, run=subprocess.run):
    path = "/" + service.replace(".", "/")
    command = [os.environ.get("GNOBLIN_BUSCTL", "busctl"), "--user", "--json=short", f"--timeout={timeout}s",
               "call", service, path, service, method]
    if signature: command += [signature, *map(str, arguments)]
    result = run(command, capture_output=True, text=True, timeout=timeout + 1)
    if result.returncode:
        raise CommandError(result.stderr.strip() or "D-Bus request failed")
    if not result.stdout.strip(): return []
    try:
        value = json.loads(result.stdout)
        return value["data"]
    except (ValueError, KeyError, TypeError):
        raise CommandError("Invalid response from busctl") from None


def input_call(method, signature="", arguments=(), timeout=5, call=dbus):
    try: return call(method, signature, arguments, timeout=timeout)
    except CommandError as error:
        # Fall back only when the interface is absent, never after an uncertain mutation.
        text = str(error)
        if not re.match(r"^(?:Call failed: )?(?:Unknown method|Unknown interface|Unknown object|The name \S+ (?:was not provided by any \.service|is not activatable))", text):
            raise
        return call(method, signature, arguments, service="org.gnoblin.InputSources", timeout=timeout)


def compositor(request, path, timeout=5):
    identity = uuid.uuid4().hex
    deadline = time.monotonic() + timeout
    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
        connection.settimeout(timeout)
        try: connection.connect(path)
        except OSError as error:
            raise CommandError(f"Cannot connect to compositor socket {path}: {error.strerror}. Check the compositor-bridge script with gnoblinctl scripts.") from None
        connection.sendall((json.dumps(dict(request, op="command", id=identity)) + "\n").encode())
        buffer = b""
        while True:
            remaining = deadline - time.monotonic()
            if remaining <= 0: raise TimeoutError("Compositor request timed out; it was not retried")
            connection.settimeout(remaining)
            chunk = connection.recv(65536)
            if not chunk: raise CommandError("Compositor disconnected before replying; the request was not retried")
            buffer += chunk
            if len(buffer) > 4 * 1024 * 1024: raise CommandError("Compositor response is too large")
            while b"\n" in buffer:
                line, buffer = buffer.split(b"\n", 1)
                try: response = json.loads(line)
                except (ValueError, UnicodeDecodeError): raise CommandError("Invalid compositor response") from None
                if not isinstance(response, dict): raise CommandError("Invalid compositor response")
                if response.get("id") != identity: continue
                if response.get("event") == "error": raise CommandError(response.get("message", "Compositor rejected request"))
                if response.get("event") == "reply":
                    if not isinstance(response.get("result"), dict): raise CommandError("Invalid compositor response")
                    return response["result"]


def keyed(rows, fields):
    return [dict(zip(fields, row)) for row in rows]


def dispatch(args, cli):
    cmd = args.command
    call = lambda method, signature="", arguments=(), service="org.gnoblin.Shell": dbus(method, signature, arguments, service, args.timeout)
    rpc = lambda request: compositor(request, args.socket, args.timeout)
    if cmd == "windows" or cmd == "window" and args.action == "list":
        result = rpc({"command": "windows"})
        result["windows"] = [window for window in result["windows"] if
            (not args.app_id or window["appId"] == args.app_id) and
            (not args.title or args.title.casefold() in window["title"].casefold()) and
            (not args.focused or window["focused"])]
        return result
    if cmd == "window":
        request = {"command": "window", "action": args.action, "window": args.id}
        for field in ("x", "y", "width", "height", "workspace", "monitor"):
            if hasattr(args, field): request[field] = getattr(args, field)
        return rpc(request)
    if cmd == "monitors": return rpc({"command": "monitors"})
    if cmd == "workspaces":
        if (args.action == "switch") != (args.workspace is not None): cli.error("workspaces switch requires an ID; workspaces list takes no ID")
        return rpc({"command": "workspaces"} if args.action == "list" else {"command": "workspace-switch", "workspace": args.workspace})
    if cmd == "status":
        result = {"version": call("GetVersion")[0], "connected": call("Ping")[0] == "pong"}
        try:
            windows = rpc({"command": "windows"})["windows"]
            result.update(windows=len(windows), focused=next((window["id"] for window in windows if window["focused"]), None))
        except (CommandError, TimeoutError) as error: result["windowControlError"] = str(error)
        return result
    if cmd == "ping": return call("Ping")[0]
    if cmd == "version": return call("GetVersion")[0]
    if cmd == "features": return {"features": keyed(call("ListFeatures")[0], ["id", "description", "enabled"])}
    if cmd == "feature": return {"id": args.id, "enabled": call("GetFeature", "s", [args.id])[0]}
    if cmd in ("enable", "disable"):
        call("SetFeature", "sb", [args.id, str(cmd == "enable").lower()])
        return {"ok": True, "id": args.id, "enabled": cmd == "enable"}
    if cmd == "extensions": return {"extensions": keyed(call("ListExtensions")[0], ["id", "state"])}
    if cmd == "scripts": return {"scripts": call("ListScripts")[0]}
    if cmd in ("input-sources", "input-source", "set-input-source"):
        fields = ["type", "id", "shortName", "name"]
        if cmd == "input-sources": return {"sources": keyed(input_call("ListInputSources", timeout=args.timeout)[0], fields)}
        if cmd == "input-source": return dict(zip(fields, input_call("GetCurrentInputSource", timeout=args.timeout)))
        input_call("SetInputSource", "ss", [args.type, args.id], timeout=args.timeout)
        return {"ok": True, "type": args.type, "id": args.id}
    if cmd == "privacy": return dict(zip(["screenSharing", "microphoneInUse", "locationInUse"], call("GetPrivacyState")))
    if cmd == "portal-grants": return {"grants": keyed(call("ListPortalGrants")[0], ["id", "kind", "requester", "devices", "clipboard", "screenStreams"])}
    if cmd.startswith("launch-"):
        service = "org.gnoblin.LaunchFeedback"
        if cmd == "launch-state":
            result = call("GetState", service=service)[0]
            return json.loads(result) if isinstance(result, str) else result
        if cmd == "launch-begin": call("Begin", "ssu", [args.token, args.application, args.milliseconds], service)
        else: call("End", "s", [args.token], service)
        return {"ok": True, "token": args.token}
    operations = {"reload": ("Reload", "", []), "reload-config": ("ReloadConfig", "", []),
        "reload-scripts": ("ReloadScripts", "", []), "reload-ext": ("ReloadExtension", "s", [getattr(args, "uuid", "")]),
        "revoke-grant": ("RevokePortalGrant", "ss", [getattr(args, "kind", ""), getattr(args, "id", "")])}
    method, signature, values = operations[cmd]
    call(method, signature, values)
    return {"ok": True, "action": cmd}


def text(value):
    if value is None: return "-"
    if isinstance(value, bool): return "yes" if value else "no"
    if isinstance(value, (list, dict)): value = json.dumps(value, ensure_ascii=False)
    return "".join(character if character.isprintable() else " " for character in str(value))


def table(rows):
    if not rows:
        print("No items.")
        return
    if not isinstance(rows[0], dict):
        for row in rows: print(text(row))
        return
    fields = list(rows[0])
    if "appId" in fields:
        fields = ["id", "focused", "workspace", "monitorIndex", "appId", "title"]
    values = [[text(row.get(field)) for field in fields] for row in rows]
    widths = [max(len(field), *(len(row[index]) for row in values)) for index, field in enumerate(fields)]
    available = max(40, shutil.get_terminal_size((120, 24)).columns) - 2 * (len(fields) - 1)
    while sum(widths) > available and max(widths) > 10:
        index = widths.index(max(widths)); widths[index] -= 1
    def line(row):
        return "  ".join((value if len(value) <= width else value[:max(0, width - 3)] + "...").ljust(width) for value, width in zip(row, widths)).rstrip()
    print(line([field.upper() for field in fields]))
    print(line(["-" * width for width in widths]))
    for row in values: print(line(row))


def render(result, output_format):
    if output_format == "json" or output_format == "auto" and not sys.stdout.isatty() and not isinstance(result, str):
        print(json.dumps(result, ensure_ascii=False))
    elif isinstance(result, str): print(text(result))
    elif isinstance(result, dict) and len(result) == 1 and isinstance(next(iter(result.values())), list): table(next(iter(result.values())))
    elif isinstance(result, dict):
        for key, value in result.items(): print(f"{key}: {text(value)}")
    else: print(text(result))


def completions(cli, shell):
    choices = next(action.choices for action in cli._actions if isinstance(action, argparse._SubParsersAction))
    names = " ".join(choices)
    actions = "list focus close minimize restore-or-minimize restore maximize unmaximize fullscreen unfullscreen move resize workspace monitor"
    if shell == "bash":
        print("""_gnoblinctl() {
    local cur=${COMP_WORDS[COMP_CWORD]} words='--help --json --format --timeout --socket'
    case "${COMP_WORDS[1]}:$COMP_CWORD" in
        *:1) words='NAMES' ;;
        window:2) words='ACTIONS' ;;
        workspaces:2) words='list switch' ;;
        completion:2) words='bash zsh fish' ;;
    esac
    case "${COMP_WORDS[COMP_CWORD-1]}" in
        --format) words='auto json table' ;;
        --socket) COMPREPLY=($(compgen -f -- "$cur")); return ;;
    esac
    COMPREPLY=($(compgen -W "$words" -- "$cur"))
}
complete -F _gnoblinctl gnoblinctl""".replace("NAMES", names).replace("ACTIONS", actions))
    elif shell == "zsh":
        print("""#compdef gnoblinctl
_gnoblinctl() {
    if (( CURRENT == 2 )); then
        compadd NAMES
    elif (( CURRENT == 3 )); then
        case $words[2] in
            window) compadd ACTIONS ;;
            workspaces) compadd list switch ;;
            completion) compadd bash zsh fish ;;
        esac
    else
        _arguments '--json[JSON output]' '--format[Output format]:format:(auto json table)' '--socket[Compositor socket]:path:_files' '--timeout[Request timeout]:seconds:'
    fi
}
compdef _gnoblinctl gnoblinctl""".replace("NAMES", names).replace("ACTIONS", actions))
    else:
        print("complete -c gnoblinctl -f")
        for name in choices: print(f"complete -c gnoblinctl -n '__fish_use_subcommand' -a '{name}'")
        print(f"complete -c gnoblinctl -n '__fish_seen_subcommand_from window' -a '{actions}'")
        print("complete -c gnoblinctl -l json -s j -d 'JSON output'")


def main(argv=None):
    cli = parser()
    argv = list(sys.argv[1:] if argv is None else argv)
    if argv and argv[0] == "help": argv = [*argv[1:], "--help"]
    args = cli.parse_args(argv)
    defaults = {"format": "auto", "timeout": 5, "socket": os.environ.get("GNOBLIN_COMPOSITOR_SOCKET") or str(Path(os.environ.get("XDG_RUNTIME_DIR", f"/run/user/{os.getuid()}")) / "gnoblin/compositor-v1.sock")}
    for key, value in defaults.items():
        if not hasattr(args, key): setattr(args, key, value)
    if not args.command:
        cli.print_help()
        return 0
    if args.command == "completion":
        completions(cli, args.shell)
        return 0
    try:
        render(dispatch(args, cli), args.format)
        return 0
    except BrokenPipeError:
        return 0
    except (CommandError, OSError, subprocess.TimeoutExpired, TimeoutError) as error:
        message = str(error)
        if isinstance(error, (TimeoutError, subprocess.TimeoutExpired)): message = "Request timed out; it was not retried. Check current state before repeating an action."
        print(f"gnoblinctl: {message}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())
