#!/usr/bin/env python3

import argparse
import collections
import curses
import itertools
import os
import sys
import math

from enum import Enum
from math import floor
from threading import Thread, Lock

from cubestat.readers.linux_reader import LinuxReader
from cubestat.readers.macos_reader import AppleReader

class EnumLoop(Enum):
    def next(self):
        values = list(self.__class__)
        return values[(values.index(self) + 1) % len(values)]
    
class EnumStr(Enum):
    def __str__(self):
        return self.value

class Legend(EnumLoop, EnumStr):
    hidden = 'off'
    last = 'last'
    
class CPUMode(EnumLoop, EnumStr):
    all = 'all'
    by_cluster = 'by_cluster'
    by_core = 'by_core'

class PowerMode(EnumLoop, EnumStr):
    combined = 'combined'
    all = 'all'
    off = 'off'

class GPUMode(EnumLoop, EnumStr):
    collapsed = 'collapsed'
    load_only = 'load_only'
    load_and_vram = 'load_and_vram'
    
class Color(EnumStr):
    red = 'red'
    green = 'green'
    blue = 'blue'
    pink = 'pink'
    olive = 'olive'
    navy = 'navy'
    blue_dark = 'blue_dark'
    purple = 'purple'
    mixed = 'mixed'
    dark = 'dark'

def auto_cpu_mode():
     return CPUMode.all if os.cpu_count() < 40 else CPUMode.by_cluster

parser = argparse.ArgumentParser("cubestat")
parser.add_argument('--refresh_ms', '-i', type=int, default=1000, help='Update frequency, milliseconds')
parser.add_argument('--buffer_size', type=int, default=500, help='How many datapoints to store. Having it larger than screen width is a good idea as terminal window can be resized')
parser.add_argument('--cpu', type=CPUMode, default=auto_cpu_mode(), choices=list(CPUMode), help='CPU mode - showing all cores, only cumulative by cluster or both. Can be toggled by pressing c.')
parser.add_argument('--gpu', type=GPUMode, default=GPUMode.load_only, choices=list(GPUMode), help='GPU mode - hidden, showing all GPUs load, or showing load and vram usage. Can be toggled by pressing g.')
parser.add_argument('--power', type=PowerMode, default=PowerMode.combined, choices=list(PowerMode), help='Power mode - off, showing breakdown CPU/GPU/ANE load, or showing combined usage. Can be toggled by pressing p.')
parser.add_argument('--color', type=Color, default=Color.mixed, choices=list(Color))
parser.add_argument('--legend', type=Legend, default=Legend.last, choices=list(Legend), help='Show/hide numeric utilization percentage. Can be toggled by pressing l.')
parser.add_argument('--disk', action="store_true", default=True, help="Show disk read/write. Can be toggled by pressing d.")
parser.add_argument('--swap', action="store_true", default=True, help="Show swap . Can be toggled by pressing s.")
parser.add_argument('--network', action="store_true", default=True, help="Show network io. Can be toggled by pressing n.")
parser.add_argument('--no-disk', action="store_false", dest="disk", help="Hide disk read/write. Can be toggled by pressing d.")
parser.add_argument('--no-swap', action="store_false", default=True, help="Hide swap. Can be toggled by pressing s.")
parser.add_argument('--no-network', action="store_false", dest="network", help="Hide network io. Can be toggled by pressing n.")

args = parser.parse_args()

class Horizon:
    def __init__(self, stdscr, reader):
        stdscr.nodelay(False)
        stdscr.timeout(50)
        curses.curs_set(0)
        curses.start_color()
        curses.use_default_colors()

        self.spacing_width = 1
        self.filling = '.'

        self.cells = self.prepare_cells()
        self.stdscr = stdscr

        # all of the fields below are mutable and can be accessed from 2 threads
        self.lock = Lock()
        self.data = {k: collections.defaultdict(lambda: collections.deque(maxlen=args.buffer_size)) for k in ['cpu', 'ram', 'swap', 'gpu', 'ane', 'disk', 'network', 'power']}
        dark_colormap = {
            'cpu': Color.purple,
            'ram': Color.navy,
            'gpu': Color.blue_dark,
            'ane': Color.blue_dark,
            'disk': Color.olive,
            'network': Color.olive,
            'swap': Color.navy,
            'power': Color.blue_dark,
        }
        light_colormap = {
            'cpu': Color.green,
            'ram': Color.pink,
            'gpu': Color.red,
            'ane': Color.red,
            'disk': Color.blue,
            'network': Color.blue,
            'swap': Color.pink,
            'power': Color.red,
        }

        if args.color == Color.mixed:
            self.colormap = light_colormap
        elif args.color == Color.dark:
            self.colormap = dark_colormap
        else:
            self.colormap = {k: args.color for k, _ in light_colormap.items()}
        self.snapshots_observed = 0
        self.snapshots_rendered = 0
        self.legend_mode = args.legend
        self.cpumode = args.cpu
        self.show_disk = args.disk
        self.show_swap = args.swap
        self.show_network = args.network
        self.power_mode = args.power
        self.gpumode = args.gpu
        self.settings_changed = False
        self.reader = reader
        self.vertical_shift = 0
        self.horizontal_shift = 0

    def prepare_cells(self):
        chrs = [' ', '▁', '▂', '▃', '▄', '▅', '▆', '▇', '█']
        colorschemes = {
            Color.green: [-1, 150, 107, 22],
            Color.red: [-1, 224, 181, 138],
            Color.blue: [-1, 189, 146, 103],
            Color.pink: [-1, 223, 180, 137],
            Color.olive: [-1, 58, 101, 144],
            Color.navy: [-1, 18, 61, 105],
            Color.blue_dark: [-1, 23, 66, 109],
            Color.purple: [-1, 53, 96, 138]
        }
        cells = {}
        colorpair = 1
        for name, colors in colorschemes.items():
            cells[name] = []
            for fg, bg in zip(colors[1:], colors[:-1]):
                curses.init_pair(colorpair, fg, bg)
                cells[name].extend((chr, colorpair) for chr in chrs)
                colorpair += 1
        return cells

    def process_snapshot(self, data, cpu_clusters):
        for group, vals in data:
            for title, value in vals.items():
                with self.lock:
                    self.data[group][title].append(value)

        with self.lock:
            self.cpu_clusters = cpu_clusters
            self.snapshots_observed += 1
            if self.horizontal_shift > 0:
                self.horizontal_shift += 1

    def write_string(self, row, col, s, color=0):
        if row < 0 or row >= self.rows or col < 0:
            return
        if col + len(s) > self.cols:
            s = s[:self.cols - col]
        try:
            self.stdscr.addstr(row, col, s, color)
        except:
            # TODO: log something
            pass

    def write_char(self, row, col, chr, color=0):
        if row < 0 or row >= self.rows or col < 0 or col >= self.cols:
            return
        try:
            self.stdscr.addch(row, col, chr, color)
        except:
            # TODO: log something
            pass

    # buckets is a list of factor/label, e.g. [(1024*1024, 'Mb'), (1024, 'Kb'), (1, 'b')]
    def format_measurement(self, spacing, curr, mx, buckets):
        if self.legend_mode != Legend.last:
            return f'{spacing}╗'
        for lim, unit in buckets[:-1]:
            if mx > lim:
                return f'last:{curr / lim :3.0f}|{int(mx / lim)}{unit}{spacing}╗'
        return f'last:{curr :3.0f}|{int(mx)}{buckets[-1][1]}{spacing}╗'

    def render(self):
        with self.lock:
            if self.snapshots_observed == self.snapshots_rendered and not self.settings_changed:
                return
        self.stdscr.erase()
        self.rows, self.cols = self.stdscr.getmaxyx()
        spacing = ' ' * self.spacing_width

        # Each chart takes two lines, with format roughly
        # ╔ GPU util %........................................................................last:  4% ╗
        # ╚ ▁▁▁  ▁    ▁▆▅▄ ▁▁▁      ▂ ▇▃▃▂█▃▇▁▃▂▁▁▂▁▁▃▃▂▁▂▄▄▁▂▆▁▃▁▂▃▁▁▁▂▂▂▂▂▂▁▁▃▂▂▁▂▁▃▄▃ ▁▁▃▁▄▂▃▂▂▂▃▃▅▅ ╝

        with self.lock:
            i = 0
            skip = self.vertical_shift
            is_multigpu = len(self.data['gpu']) > 1
            for group_name, group in self.data.items():
                if group_name == 'disk' and not self.show_disk:
                    continue
                if group_name == 'network' and not self.show_network:
                    continue
                if group_name == 'swap' and not self.show_swap:
                    continue

                cells = self.cells[self.colormap[group_name]]
                range = len(cells)
                for title, series in group.items():
                    indent = ''

                    if group_name == 'cpu':
                        if self.cpumode == CPUMode.by_cluster and title not in self.cpu_clusters:
                            continue
                        if self.cpumode == CPUMode.by_core and title in self.cpu_clusters:
                            continue
                        if self.cpumode == CPUMode.all and title not in self.cpu_clusters:
                            indent = '  '

                    if group_name == 'gpu':
                        if is_multigpu and self.gpumode == GPUMode.collapsed and "Total GPU" not in title:
                            continue
                        if self.gpumode == GPUMode.load_only and "vram" in title:
                            continue
                        if is_multigpu and "Total GPU" not in title:
                            indent = '  '

                    if group_name == 'power':
                        if self.power_mode == PowerMode.off:
                            continue
                        if self.power_mode == PowerMode.combined and 'total' not in title:
                            continue
                        if 'total' not in title:
                            indent = '  '
                    
                    if skip > 0:
                        skip -= 1
                        continue

                    # render title and left border, for example
                    #
                    # ╔ GPU util %
                    # ╚
                    titlestr = f'{indent}╔{spacing}{title}'
                    self.write_string(i * 2, 0, titlestr)
                    self.write_string(i * 2 + 1, 0, f'{indent}╚')

                    # data slice size
                    length = len(series) - self.horizontal_shift if self.horizontal_shift > 0 else len(series)
                    
                    # chart area width
                    width = self.cols - 2 * self.spacing_width - 2 - len(indent)
                    index = max(0, length - width)
                    data_slice = list(itertools.islice(series, index, min(index + width, length)))

                    # for percentage-like measurements
                    B = 100.0
                    strvalue = f'last:{data_slice[-1]:3.0f}%{spacing}╗' if self.legend_mode == Legend.last else f'{spacing}╗'
                    
                    if group_name == 'disk' or group_name == 'network':
                        B = max(data_slice)
                        B = float(1 if B == 0 else 2 ** (int((B - 1)).bit_length()))
                        strvalue = self.format_measurement(spacing, data_slice[-1], B, [(1024 * 1024, 'MB/s'), (1024, 'KB/s'), (1, 'Bytes/s')])

                    if group_name == 'swap':
                        B = max(data_slice)
                        B = float(1 if B == 0 else 2 ** (int((B - 1)).bit_length()))
                        strvalue = self.format_measurement(spacing, data_slice[-1], B, [(1024 * 1024, 'MB'), (1024, 'KB'), (1, 'Bytes')])

                    if group_name == 'power':
                        B = max(data_slice)
                        B = float(1 if B <= 0 else 10 ** math.ceil(math.log10(B)))
                        strvalue = self.format_measurement(spacing, data_slice[-1], B, [(1000 * 1000, 'kW'), (1000, 'W'), (1, 'mW')])

                    # render the rest of title row
                    #
                    # ╔ GPU util %........................................................................last:  4% ╗
                    # ╚
                    title_filling = self.filling * (self.cols - len(strvalue) - len(titlestr))
                    self.write_string(i * 2, len(titlestr), title_filling)
                    self.write_string(i * 2, self.cols - len(strvalue), strvalue)

                    # render the right border
                    #
                    # ╔ GPU util %........................................................................last:  4% ╗
                    # ╚                                                                                             ╝
                    border = f'{spacing}╝'
                    self.write_string(i * 2 + 1, self.cols - len(border), border)

                    # Render the chart itself
                    #
                    # ╔ GPU util %........................................................................last:  4% ╗
                    # ╚ ▁▁▁  ▁    ▁▆▅▄ ▁▁▁      ▂ ▇▃▃▂█▃▇▁▃▂▁▁▂▁▁▃▃▂▁▂▄▄▁▂▆▁▃▁▂▃▁▁▁▂▂▂▂▂▂▁▁▃▂▂▁▂▁▃▄▃ ▁▁▃▁▄▂▃▂▂▂▃▃▅▅ ╝
                    scaler = range / B
                    col = self.cols - (len(data_slice) + self.spacing_width) - 2
                    for v in data_slice:
                        col += 1
                        cell_index = floor(v * scaler)
                        if cell_index <= 0:
                            continue
                        if cell_index >= range:
                            cell_index = range - 1
                        chr, color_pair = cells[cell_index]
                        self.write_char(i * 2 + 1, col, chr, curses.color_pair(color_pair))

                    i += 1

                self.snapshots_rendered += 1
                
        self.stdscr.refresh()

    def render_loop(self):
        self.stdscr.keypad(True)
        while True:
            self.render()
            key = self.stdscr.getch()
            if key == ord('q') or key == ord('Q'):
                exit(0)
            if key == ord('l'):
                with self.lock:
                    self.legend_mode = self.legend_mode.next()
                    self.settings_changed = True
            if key == ord('c'):
                with self.lock:
                    self.cpumode = self.cpumode.next()
                    self.settings_changed = True
            if key == ord('p'):
                with self.lock:
                    self.power_mode = self.power_mode.next()
                    self.settings_changed = True
            if key == ord('g'):
                with self.lock:
                    self.gpumode = self.gpumode.next()
                    self.settings_changed = True
            if key == ord('s'):
                with self.lock:
                    self.show_swap = not self.show_swap
                    self.settings_changed = True
            if key == ord('d'):
                with self.lock:
                    self.show_disk = not self.show_disk
                    self.settings_changed = True
            if key == ord('n'):
                with self.lock:
                    self.show_network = not self.show_network
                    self.settings_changed = True
            if key == curses.KEY_UP:
                with self.lock:
                    if self.vertical_shift > 0:
                        self.vertical_shift -= 1
                        self.settings_changed = True
            if key == curses.KEY_DOWN:
                with self.lock:
                    self.vertical_shift += 1
                    self.settings_changed = True
            if key == curses.KEY_LEFT:
                with self.lock:
                    if self.horizontal_shift + 1 < self.snapshots_observed:
                        self.horizontal_shift += 1
                        self.settings_changed = True
            if key == curses.KEY_RIGHT:
                with self.lock:
                    if self.horizontal_shift > 0:
                        self.horizontal_shift -= 1
                        self.settings_changed = True
            if key == ord('0'):
                with self.lock:
                    if self.horizontal_shift > 0:
                        self.horizontal_shift = 0
                        self.settings_changed = True

    def loop(self):
        reader_thread = Thread(target=self.reader.loop, daemon=True, args=[self.process_snapshot])
        reader_thread.start()
        self.render_loop()

def start(stdscr, reader):
    h = Horizon(stdscr, reader)
    h.loop()

def main():
    if sys.platform == "darwin":
        curses.wrapper(start, AppleReader(args.refresh_ms))
    if sys.platform == "linux" or sys.platform == "linux2":
        curses.wrapper(start, LinuxReader(args.refresh_ms))
    # TODO: write something about platform not supported

if __name__ == '__main__':
    main()
