from slapos.test.promise.plugin import TestPromisePluginMixin
import os
import sqlite3
import math
import random
from datetime import date, timedelta

def list_dates(start_date=None, end_date=None):
  if (start_date is None) or (end_date is None):
    start_date = date(2021,6,1) 
    end_date = date(2021,6,16) # until 06/15
  for n in range(int((end_date - start_date).days)):
    yield start_date + timedelta(n)

def update(add, used, free, all_partitions, name):
  if (all_partitions[name] + add) >= 0:
    used += add
    free -= add
    all_partitions[name] = int(all_partitions[name] + add)
  return int(used), int(free), all_partitions

class TestCheckFreeDiskSpace(TestPromisePluginMixin):

  INSERT_FOLDER_TEMPLATE = "insert into folder(" \
            "partition, disk_used, date, time) values " \
             "('%s', %s, '%s', '00:00:00')" 

  INSERT_DISK_TEMPLATE = "insert into disk("\
            " partition, used, free, mountpoint," \
            " date, time) "\
            "values ('/dev/sda1', %s, %s, '/', '%s', '00:00:00')"

  def setUp(self):
    TestPromisePluginMixin.setUp(self)
    log_folder = os.path.join(self.partition_dir, 'var/log')
    os.makedirs(log_folder)

    self.db_file = '/tmp/collector.db'

    # populate db
    self.conn = sqlite3.connect(self.db_file)
    f = open(self.base_path+"/partitiontest.sql")
    sql = f.read()
    self.conn.executescript(sql)
    self.insertTable(disk_size = 200e9)
    self.conn.commit()
    #self.conn.close()

    self.promise_name = "check-free-disk-space.py"

    content = """from slapos.promise.plugin.check_free_disk_space import RunPromise

extra_config_dict = {
  'collectordb': '%(collectordb)s',
  'test-check-date': '2021-06-15'
}
""" % {'collectordb': self.db_file}
    self.writePromise(self.promise_name, content)

  def tearDown(self):
    TestPromisePluginMixin.tearDown(self)
    if os.path.exists(self.db_file):
      os.rename(self.db_file, '/tmp/tests/collector.db')
      #os.remove(self.db_file)

  def insertDiskFolder(self, date, used, free, name, all_partitions):
    insertion_disk = self.INSERT_DISK_TEMPLATE % (used, free, date)
    insertion_folder = self.INSERT_FOLDER_TEMPLATE % (name, all_partitions[name]//1024, date)
    self.conn.execute(insertion_disk)
    self.conn.execute(insertion_folder)
  
  def insertTable(self, disk_size = 200e9):
    all_dates = list(list_dates())
    all_partitions = {'slapuser1':10e9, 'slapuser2':10e9, 'slapuser3':10e9, 'slapuser4':10e9}
    used = 0 + (all_partitions['slapuser1'] + all_partitions['slapuser3'] + \
      all_partitions['slapuser3'] + all_partitions['slapuser4'])
    free = disk_size - used
    i = 0

    for date in all_dates:
      name = 'slapuser1'
      used, free, all_partitions = update(2e9, used, free, all_partitions, name)
      self.insertDiskFolder(date, used, free, name, all_partitions)

      name = 'slapuser2'
      add = math.exp(i+8)*10e3 - math.exp(i+8-0.5)*10e3
      used, free, all_partitions = update(add, used, free, all_partitions, name)
      self.insertDiskFolder(date, used, free, name, all_partitions)

      name = 'slapuser3'
      add = math.atan(i+0.5)*10e9 -  math.atan(i)*10e9
      used, free, all_partitions = update(add, used, free, all_partitions, name)
      self.insertDiskFolder(date, used, free, name, all_partitions)

      name = 'slapuser4'
      add = random.randint(-3*10e8, 5*10e8) # between -3 giga and +5 giga
      used, free, all_partitions = update(add, used, free, all_partitions, name)
      self.insertDiskFolder(date, used, free, name, all_partitions)

      i += 0.5

  def test_slapuser1(self):
    content = """from slapos.promise.plugin.check_free_disk_space import RunPromise

extra_config_dict = {
  'collectordb': '%(collectordb)s',
  'test-check-date': '2021-06-15',
  'test-check-time': '00:00',
  'name-user': 'slapuser1'
}
""" % {'collectordb': self.db_file}
    self.writePromise(self.promise_name, content)

    self.configureLauncher()
    self.launcher.run()
    result = self.getPromiseResult(self.promise_name)
    print(result['result']['message'])

  def test_slapuser2(self):
    content = """from slapos.promise.plugin.check_free_disk_space import RunPromise

extra_config_dict = {
  'collectordb': '%(collectordb)s',
  'test-check-date': '2021-06-15',
  'test-check-time': '00:00',
  'name-user': 'slapuser2'
}
""" % {'collectordb': self.db_file}
    self.writePromise(self.promise_name, content)

    self.configureLauncher()
    self.launcher.run()
    result = self.getPromiseResult(self.promise_name)
    print(result['result']['message'])

  def test_slapuser3(self):
    content = """from slapos.promise.plugin.check_free_disk_space import RunPromise

extra_config_dict = {
  'collectordb': '%(collectordb)s',
  'test-check-date': '2021-06-15',
  'test-check-time': '00:00',
  'name-user': 'slapuser3'
}
""" % {'collectordb': self.db_file}
    self.writePromise(self.promise_name, content)

    self.configureLauncher()
    self.launcher.run()
    result = self.getPromiseResult(self.promise_name)
    print(result['result']['message'])

  def test_slapuser4(self):
    content = """from slapos.promise.plugin.check_free_disk_space import RunPromise

extra_config_dict = {
  'collectordb': '%(collectordb)s',
  'test-check-date': '2021-06-15',
  'test-check-time': '00:00',
  'name-user': 'slapuser4'
}
""" % {'collectordb': self.db_file}
    self.writePromise(self.promise_name, content)

    self.configureLauncher()
    self.launcher.run()
    result = self.getPromiseResult(self.promise_name)
    print(result['result']['message'])

if __name__ == '__main__':
  unittest.main()