#!/usr/bin/env python

import json
import os
import os.path
import subprocess
import sys

os.chdir(os.path.expanduser('~/.config/dotfiles/hypr'))
bars = "_▂▃▄▅▆▇█"

class Monitor:
    """
    A monitor on which many workspaces can appear.
    """
    def __init__(self, id: int, name: str):
        self.id = id
        self.name = name

    def __eq__(self, other):
        """
        Returns true if the monitors are the same.
        """
        self.id == other.id

    @staticmethod
    def active(monitors: list['Monitor']) -> 'Monitor':
        """
        Returns the active monitor.
        """
        proc = subprocess.run(['hyprctl', 'activeworkspace', '-j'], capture_output=True)
        result = json.loads(proc.stdout)
        return next((monitor for monitor in monitors if monitor.id == result['monitorID']))

    @staticmethod
    def all() -> list['Monitor']:
        """
        Returns the list of available monitors.
        """
        monitors: list['Monitor'] = []

        proc = subprocess.run(['hyprctl', 'monitors', '-j'], capture_output=True)
        result = json.loads(proc.stdout)

        for mon in result:
            monitors.append(Monitor(mon['id'], mon['name']))

        return monitors

    def previous(self, monitors: list['Monitor'], move: bool = False):
        """
        Moves the focus to the previous monitor.

        Params:
            move: whether you want to move the active window to the new monitor.
        """
        mon = monitors[(monitors.index(self) - 1) % len(monitors)]

        if move:
            subprocess.run(['hyprctl', 'dispatch', 'movewindow', 'mon:' + str(mon.id)])
        else:
            subprocess.run(['hyprctl', 'dispatch', 'focusmonitor', str(mon.id)])

    def next(self, monitors: list['Monitor'], move: bool = False):
        """
        Moves the focus to the next monitor.

        Params:
            move: whether you want to move the active window to the new monitor.
        """
        mon = monitors[(monitors.index(self) + 1) % len(monitors)]

        if move:
            subprocess.run(['hyprctl', 'dispatch', 'movewindow', 'mon:' + str(mon.name)])
        else:
            subprocess.run(['hyprctl', 'dispatch', 'focusmonitor', str(mon.name)])


class Workspace:
    """
    A workspace bound to a monitor.

    Workspaces with id from 10 * n + k belongs to the same monitor for k in [1, 10].
    """
    def __init__(self, id: int):
        self.id = id

    @staticmethod
    def active() -> 'Workspace':
        """
        Returns the active workspace.
        """
        proc = subprocess.run(['hyprctl', 'activeworkspace', '-j'], capture_output=True)
        result = json.loads(proc.stdout)
        return Workspace(result['id'])

    def nth(self, n: int, move: bool = False):
        """
        Goes to the nth workspace on the same monitor as the current workspace.

        Params:
            n: workspace to go to, between 1 and 10.
            move: whether you want to move the active window to the new workspace.
        """
        new_id = (self.id - 1) // 10 * 10 + n
        subprocess.run(['hyprctl', 'dispatch', 'movetoworkspace' if move else 'workspace', str(new_id)])

    def previous(self, move: bool = False):
        """
        Goes to the previous workspace on the same monitor, or the last if we're on the first.

        Params:
            move: whether you want to move the active window to the new workspace.
        """
        new_id = self.id - 1

        if new_id % 10 == 0:
            new_id += 10

        subprocess.run(['hyprctl', 'dispatch', 'movetoworkspace' if move else 'workspace', str(new_id)])

    def next(self, move: bool = False):
        """
        Goes to the next workspace on the same monitor, or the first if we're on the last.

        Params:
            move: whether you want to move the active window to the new workspace.
        """
        new_id = self.id + 1

        if new_id % 10 == 1:
            new_id -= 10

        subprocess.run(['hyprctl', 'dispatch', 'movetoworkspace' if move else 'workspace', str(new_id)])


def to_bar(x):
    return bars[round((len(bars) - 1) * x / 100)]


def monitor_stats():
    try:
        import psutil
    except:
        print('Stats requires psutil (run `pip install psutil`)',  file=sys.stderr)
        sys.exit(1)

    """
    Monitors CPU and MEM usage and prints info in file.
    """
    cpu_values = [0] * 10
    mem_values = [0] * 10

    while True:
        cpu_percent = psutil.cpu_percent(interval=2)
        mem_percent = psutil.virtual_memory().percent

        for i in range(len(cpu_values) - 1):
            cpu_values[i] = cpu_values[i+1]
            mem_values[i] = mem_values[i+1]

        cpu_values[-1] = cpu_percent
        mem_values[-1] = mem_percent

        with open('.stat.txt', 'w') as f:
            f.write(
                '  ' + ''.join([to_bar(x) for x in cpu_values]) + ' ' + "% 5.1f"%cpu_percent + '%\n' +
                '  ' + ''.join([to_bar(x) for x in mem_values]) + ' ' + "% 5.1f"%mem_percent + '%'
            )


def main():
    if len(sys.argv) < 2:
        return

    if sys.argv[1] == 'workspace' or sys.argv[1] == 'movetoworkspace':
        workspace = Workspace.active()
        move = sys.argv[1] == 'movetoworkspace'

        if sys.argv[2] == 'next':
            workspace.next(move)

        elif sys.argv[2] == 'previous':
            workspace.previous(move)

        else:
            new_workspace = int(sys.argv[2])
            workspace.nth(new_workspace, move)

    elif sys.argv[1] in ['movewindow', 'focusnextmonitor', 'focuspreviousmonitor']:
        monitors = Monitor.all()
        monitor = Monitor.active(monitors)

        if sys.argv[1] == 'movewindow':
            monitor.next(monitors, True)
        elif sys.argv[1] == 'focusnextmonitor':
            monitor.next(monitors)
        elif sys.argv[1] == 'focuspreviousmonitor':
            monitor.previous(monitors)

    elif sys.argv[1] == 'reload':
        subprocess.run(['systemctl', 'restart', 'waybar', '--user'])

    elif sys.argv[1] == 'stat':
        monitor_stats()

    else:
        print(f'Command not found: {sys.argv[1]}', file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main()