"""WizardKit: TUI functions"""
# vim: sts=2 sw=2 ts=2

import atexit
import logging
import time

from copy import deepcopy
from os import environ

from wk.exe import start_thread
from wk.std import sleep
from wk.ui import ansi, tmux

# STATIC VARIABLES
LOG = logging.getLogger(__name__)
TMUX_SIDE_WIDTH = 21
TMUX_TITLE_HEIGHT = 2
TMUX_LAYOUT = { # NOTE: This needs to be in order from top to bottom
  'Title':    {'Panes': [], 'height': TMUX_TITLE_HEIGHT},
  'Info':     {'Panes': []},
  'Current':  {'Panes': [environ.get('TMUX_PANE', None)]},
  'Workers':  {'Panes': []},
  'Started':  {'Panes': [], 'width': TMUX_SIDE_WIDTH},
  'Progress': {'Panes': [], 'width': TMUX_SIDE_WIDTH},
  }


# Classes
class TUI():
  """Object for tracking TUI elements."""
  def __init__(self, title_text=None) -> None:
    self.layout = deepcopy(TMUX_LAYOUT)
    self.side_width = TMUX_SIDE_WIDTH
    self.title_text = title_text if title_text else 'Title Text'
    self.title_text_line2 = ''
    self.title_colors = ['BLUE', None]

    # Init tmux and start a background process to maintain layout
    self.init_tmux()
    start_thread(self.fix_layout_loop)

    # Close all panes at exit
    atexit.register(tmux.kill_all_panes)

  def add_info_pane(
      self, lines=None, percent=None, update_layout=True, **tmux_args,
      ) -> None:
    """Add info pane."""
    if not (lines or percent):
      # Bail early
      raise RuntimeError('Neither lines nor percent specified.')

    # Calculate lines if needed
    if not lines:
      lines = int(tmux.get_pane_size()[1] * (percent/100))

    # Set tmux split args
    tmux_args.update({
      'behind': True,
      'lines': lines,
      'target_id': None,
      'vertical': True,
      })
    if self.layout['Info']['Panes']:
      tmux_args.update({
        'behind': False,
        'percent': 50,
        'target_id': self.layout['Info']['Panes'][-1],
        'vertical': False,
        })
      tmux_args.pop('lines')

    # Update layout
    if update_layout:
      self.layout['Info']['height'] = lines

    # Add pane
    self.layout['Info']['Panes'].append(tmux.split_window(**tmux_args))

  def add_title_pane(self, line1, line2=None, colors=None) -> None:
    """Add pane to title row."""
    lines = [line1, line2]
    colors = colors if colors else self.title_colors.copy()
    if not line2:
      lines.pop()
      colors.pop()
    tmux_args = {
      'behind': True,
      'lines': TMUX_TITLE_HEIGHT,
      'target_id': None,
      'text': ansi.color_string(lines, colors, sep='\n'),
      'vertical': True,
      }
    if self.layout['Title']['Panes']:
      tmux_args.update({
        'behind': False,
        'percent': 50,
        'target_id': self.layout['Title']['Panes'][-1],
        'vertical': False,
        })
      tmux_args.pop('lines')

    # Add pane
    self.layout['Title']['Panes'].append(tmux.split_window(**tmux_args))

  def add_worker_pane(
      self, lines=None, percent=None, update_layout=True, **tmux_args,
      ) -> None:
    """Add worker pane."""
    height = lines

    # Bail early
    if not (lines or percent):
      raise RuntimeError('Neither lines nor percent specified.')

    # Calculate height if needed
    if not height:
      height = int(tmux.get_pane_size()[1] * (percent/100))

    # Set tmux split args
    tmux_args.update({
      'behind': False,
      'lines': lines,
      'percent': percent,
      'target_id': None,
      'vertical': True,
      })

    # Update layout
    if update_layout:
      self.layout['Workers']['height'] = height

    # Add pane
    self.layout['Workers']['Panes'].append(tmux.split_window(**tmux_args))

  def clear_current_pane(self) -> None:
    """Clear screen and history for current pane."""
    tmux.clear_pane()

  def clear_current_pane_height(self) -> None:
    """Clear current pane height and update layout."""
    self.layout['Current'].pop('height', None)

  def fix_layout(self, forced=True) -> None:
    """Fix tmux layout based on self.layout."""
    try:
      tmux.fix_layout(self.layout, forced=forced)
    except RuntimeError:
      # Assuming self.panes changed while running
      pass

  def fix_layout_loop(self) -> None:
    """Fix layout on a loop.

    NOTE: This should be called as a thread.
    """
    while True:
      self.fix_layout(forced=False)
      sleep(1)

  def init_tmux(self) -> None:
    """Initialize tmux layout."""
    tmux.kill_all_panes()
    self.layout.clear()
    self.layout.update(deepcopy(TMUX_LAYOUT))

    # Title
    self.layout['Title']['Panes'].append(tmux.split_window(
      behind=True,
      lines=2,
      vertical=True,
      text=ansi.color_string(
        [self.title_text, self.title_text_line2],
        self.title_colors,
        sep = '\n',
        ),
      ))

    # Started
    self.layout['Started']['Panes'].append(tmux.split_window(
      lines=TMUX_SIDE_WIDTH,
      target_id=self.layout['Title']['Panes'][0],
      text=ansi.color_string(
        ['Started', time.strftime("%Y-%m-%d %H:%M %Z")],
        ['BLUE', None],
        sep='\n',
        ),
      ))

    # Progress
    self.layout['Progress']['Panes'].append(tmux.split_window(
      lines=TMUX_SIDE_WIDTH,
      text=' ',
      ))

  def remove_all_info_panes(self) -> None:
    """Remove all info panes and update layout."""
    self.layout['Info'].pop('height', None)
    panes = self.layout['Info']['Panes'].copy()
    self.layout['Info']['Panes'].clear()
    tmux.kill_pane(*panes)

  def remove_all_worker_panes(self) -> None:
    """Remove all worker panes and update layout."""
    self.layout['Workers'].pop('height', None)
    panes = self.layout['Workers']['Panes'].copy()
    self.layout['Workers']['Panes'].clear()
    tmux.kill_pane(*panes)

  def set_current_pane_height(self, height) -> None:
    """Set current pane height and update layout."""
    self.layout['Current']['height'] = height
    tmux.resize_pane(height=height)

  def set_progress_file(self, progress_file) -> None:
    """Set the file to use for the progresse pane."""
    tmux.respawn_pane(
      pane_id=self.layout['Progress']['Panes'][0],
      watch_file=progress_file,
      )

  def set_title(self, line1, line2=None, colors=None) -> None:
    """Set title text."""
    self.title_text = line1
    self.title_text_line2 = line2 if line2 else ''
    if colors:
      self.title_colors = colors

    # Update pane (if present)
    if self.layout['Title']['Panes']:
      tmux.respawn_pane(
        pane_id=self.layout['Title']['Panes'][0],
        text=ansi.color_string(
          [self.title_text, self.title_text_line2],
          self.title_colors,
          sep = '\n',
          ),
        )

  def update_clock(self) -> None:
    """Update 'Started' pane following clock sync."""
    tmux.respawn_pane(
      pane_id=self.layout['Started']['Panes'][0],
      text=ansi.color_string(
        ['Started', time.strftime("%Y-%m-%d %H:%M %Z")],
        ['BLUE', None],
        sep='\n',
        ),
      )


# Functions
def fix_layout(layout, forced=False):
  """Fix pane sizes based on layout."""
  resize_kwargs = []

  # Bail early
  if not (forced or layout_needs_fixed(layout)):
    # Layout should be fine
    return

  # Remove closed panes
  for data in layout.values():
    data['Panes'] = [pane for pane in data['Panes'] if tmux.poll_pane(pane)]

  # Update main panes
  for section, data in layout.items():
    if section == 'Workers':
      # Skip for now
      continue

    if 'height' in data:
      for pane_id in data['Panes']:
        resize_kwargs.append({'pane_id': pane_id, 'height': data['height']})
    if 'width' in data:
      for pane_id in data['Panes']:
        resize_kwargs.append({'pane_id': pane_id, 'width': data['width']})
  for kwargs in resize_kwargs:
    try:
      tmux.resize_pane(**kwargs)
    except RuntimeError:
      # Assuming pane was closed just before resizing
      pass

  # Update "group" panes widths
  for group in ('Title', 'Info'):
    num_panes = len(layout[group]['Panes'])
    if num_panes <= 1:
      continue
    width = int( (tmux.get_pane_size()[0] - (1 - num_panes)) / num_panes )
    for pane_id in layout[group]['Panes']:
      tmux.resize_pane(pane_id, width=width)
    if group == 'Title':
      # (re)fix Started pane
      tmux.resize_pane(layout['Started']['Panes'][0], width=TMUX_SIDE_WIDTH)

  # Bail early
  if not (layout['Workers']['Panes'] and layout['Workers']['height']):
    return

  # Update worker heights
  worker_height = layout['Workers']['height']
  workers = layout['Workers']['Panes'].copy()
  num_workers = len(workers)
  avail_height = sum(tmux.get_pane_size(pane)[1] for pane in workers)
  avail_height += tmux.get_pane_size()[1] # Current pane
  # Check if window is too small
  if avail_height < (worker_height*num_workers) + 3:
    # Just leave things as-is
    return
  # Resize current pane
  tmux.resize_pane(height=avail_height-(worker_height*num_workers))
  # Resize bottom pane
  tmux.resize_pane(workers.pop(0), height=worker_height)
  # Resize the rest of the panes by adjusting the ones above them
  while len(workers) > 1:
    next_height = sum(tmux.get_pane_size(pane)[1] for pane in workers[:2])
    next_height -= worker_height
    tmux.resize_pane(workers[1], height=next_height)
    workers.pop(0)

def layout_needs_fixed(layout):
  """Check if layout needs fixed, returns bool."""
  needs_fixed = False

  # Check panes
  for data in layout.values():
    if 'height' in data:
      needs_fixed = needs_fixed or any(
        tmux.get_pane_size(pane)[1] != data['height'] for pane in data['Panes']
        )
    if 'width' in data:
      needs_fixed = needs_fixed or any(
        tmux.get_pane_size(pane)[0] != data['width'] for pane in data['Panes']
        )

  # Done
  return needs_fixed

def test():
  """TODO: Deleteme"""
  ui = TUI()
  ui.add_info_pane(lines=10, text='Info One')
  ui.add_info_pane(lines=10, text='Info Two')
  ui.add_info_pane(lines=10, text='Info Three')
  ui.add_worker_pane(lines=3, text='Work One')
  ui.add_worker_pane(lines=3, text='Work Two')
  ui.add_worker_pane(lines=3, text='Work Three')
  ui.fix_layout()
  return ui


if __name__ == '__main__':
  print("This file is not meant to be called directly.")