#!/usr/bin/env python

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

os.chdir(os.path.expanduser('~/.config/hypr'))
bars = "_▂▃▄▅▆▇█"
batteries = "󰂎󰁺󰁻󰁼󰁽󰁾󰁿󰂀󰂁󰂁󰁹"
batteries_charging = "󰢟󰢜󰂆󰂇󰂈󰢝󰂉󰢞󰂊󰂋󰂅"


class Monitor:
    """
    A monitor on which many workspaces can appear.
    """

    def __init__(self, id: int, name: str, special_workspace: typing.Optional['Workspace']):
        """
        Constructs a monitor from its id and its name.
        """
        self.id = id
        self.name = name
        self.special_workspace = special_workspace

    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:
            if mon['specialWorkspace']['id'] == 0:
                monitors.append(Monitor(mon['id'], mon['name'], None))
            else:
                special_workspace = Workspace(mon['specialWorkspace']['id'], mon['specialWorkspace']['name'])
                monitors.append(Monitor(mon['id'], mon['name'], special_workspace))

        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, name: str):
        """
        Constructs a workspace from its id and its name.
        """
        self.id = id
        self.name = name

    def is_special(self) -> bool:
        """
        Returns true if the workspace is a magic workspace.
        """
        return self.id < 0

    def special_name(self) -> typing.Optional[str]:
        """
        Returns the part after the : if the workspace is special, none otherwise.
        """
        if self.is_special():
            return ':'.join(self.name.split(':')[1:])
        else:
            return None

    @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'], result['name'])

    def nth(self, n: int, monitor: Monitor, 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.
            monitor: active monitor.
            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)])

        current = monitor.special_workspace
        if current is not None:
            subprocess.run(['hyprctl', 'dispatch', 'togglespecialworkspace', current.special_name()])

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

        Params:
            monitor: active monitor.
            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)])

        current = monitor.special_workspace
        if current is not None:
            subprocess.run(['hyprctl', 'dispatch', 'togglespecialworkspace', current.special_name()])

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

        Params:
            monitor: active monitor.
            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)])

        current = monitor.special_workspace
        if current is not None:
            subprocess.run(['hyprctl', 'dispatch', 'togglespecialworkspace', current.special_name()])


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


def monitor_stats():
    """
    Monitors CPU and MEM usage and prints info in file.
    """

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

    cpu_values = [0] * 10
    mem_values = [0] * 10
    temp_values = [0] * 10

    min_temp = 45
    max_temp = 95

    while True:
        cpu_percent = psutil.cpu_percent(interval=2)
        mem_percent = psutil.virtual_memory().percent
        battery = psutil.sensors_battery()
        all_temps = psutil.sensors_temperatures(fahrenheit=False)

        temps = all_temps.get('k10temp', None)
        if temps is None:
            temps = all_temps.get('coretemp')

        temp = max([x.current for x in temps if x.label != 'Tctl'])

        temp_percent = 100 * min(max([temp - min_temp, 0]) / (max_temp - min_temp), 1)

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

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

        # Notify when memory greater than 75%
        if mem_percent >= 75 and not any(map(lambda x: x > 75, mem_values[:-1])):
            subprocess.run(['notify-send', '-u', 'critical', 'High memory usage', "% 5.1f" % mem_percent + '% of RAM used'])

        if battery is None:
            battery_json = '{"text": "", "class": "hidden"}'
        else:
            battery_index = round(battery.percent / 10)
            battery_icon = batteries_charging[battery_index] if battery.power_plugged else batteries[battery_index]

            if battery.power_plugged:
                battery_class = 'charging'
            elif battery.percent < 30.0:
                battery_class = 'low'
            elif battery.percent < 15.0:
                battery_class = 'critical'
            else:
                battery_class = None

            battery_text = "% 4.0f" % battery.percent if battery.percent <= 99.95 else "  100"
            battery_json = f', "class": "{battery_class}"' if battery_class else ''
            battery_json = '{"text": "' + battery_icon + ' ' + battery_text + '%"' + battery_json + '}'

        cpu_percent_text = "% 5.1f" % cpu_percent if cpu_percent <= 99.95 else "  100"
        mem_percent_text = "% 5.1f" % mem_percent if mem_percent <= 99.95 else "  100"

        with open('.stat.txt', 'w') as f:
            f.write(
                '  ' + ''.join([to_bar(x) for x in cpu_values]) + ' ' + cpu_percent_text + '%\n' +
                '  ' + ''.join([to_bar(x) for x in mem_values]) + ' ' + mem_percent_text + '%\n' +
                '  ' + ''.join([to_bar(x) for x in temp_values]) + ' ' + "% 5.1f" % temp + '°\n' +
                battery_json
            )


def start_hyprlock():
    proc = subprocess.run(['pidof', 'hyprlock'])
    if proc.returncode != 0:
        # We need to start this in the background so that the screen turns off afterwards
        subprocess.Popen(['hyprlock'], start_new_session=True)


def lock_session():
    if os.environ.get('HYPR_DISABLE_LOCK', None) is None:
        subprocess.run(['loginctl', 'lock-session'])


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

    if sys.argv[1] == 'workspace' or sys.argv[1] == 'movetoworkspace':
        monitors = Monitor.all()
        monitor = Monitor.active(monitors)

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

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

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

        else:
            new_workspace = int(sys.argv[2])
            workspace.nth(new_workspace, monitor, 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()

    elif sys.argv[1] == 'start-hyprlock':
        start_hyprlock()

    elif sys.argv[1] == 'lock-session':
        lock_session()

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


if __name__ == '__main__':
    main()
