#!/usr/bin/env python3
# -*- coding: utf-8 -*-

##############################################################################
# Switch noisy drives out of the default storage groups during sleep period,
# and set them to go to sleep.
#
# Intended to be run from a systemd timer unit.
#
# Author: J S Worthington
# Created: 2019-08-30
#
##############################################################################


##############################################################################
# Imports
##############################################################################

from enum import Enum
import asyncio
import lxml
import os
import shlex
import subprocess
import sys
import time
import MythTV
import MythTV.services_api.send as api
import MythTV.services_api.utilities as util


##############################################################################
# Version 1.0 - First version.
# Version 1.1 - Convert to Python 3.
# Version 1.2 - Tidy up.
# Version 1.3 - Add rec5.
# Version 1.4 - Remove Seagate drives, add new WD drives with 5 second PWL.
#             - Spin down noisy drives during quiet period to prevent PWL
#               noise.
# Version 1.5 - Add Noisy and Quiet commands.  Add Seagate drives again.
# Version 1.6 - Add spinning up of drives.
# Version 1.7 - Ensure partitions are mounted.
# Version 1.8 - Fix for start/stop times on either side of midnight.
#             - Update default start/stop times.
# Version 1.9 - Run drive operations in parallel.
##############################################################################

VERSION = 1.9


##############################################################################
# Configuration
##############################################################################

# Start and stop times for only using the quiet drives.
DEFAULT_QUIET_START = '23:03:00'
DEFAULT_QUIET_STOP =  '11:00:00'

# Left end of mythconverg.storagegroup.dirname string to match for noisy drives.
#dirname_match = ['/mnt/rec1', '/mnt/rec2', '/mnt/rec3', '/mnt/rec5', '/mnt/rec7']
dirname_match = ['/mnt/rec2', '/mnt/rec3', '/mnt/rec5', '/mnt/rec7']

# The IP address of the mythbackend to talk to.  Default: 127.0.0.1
HOST = '127.0.0.1'

# hdparm -S sleep timeout value (12 => 12*5 = 60 seconds).
SLEEP_TIMEOUT = 12


##############################################################################
# Constants
##############################################################################

TIME_FORMAT = '%H:%M:%S'


##############################################################################
# Debug ouput support.
##############################################################################

DEBUG_OUTPUT = True

program_name = os.path.basename(sys.argv[0])

if DEBUG_OUTPUT:
    import datetime

    global noisy_drives_debug
    noisy_drives_debug = None

    def dprint_init():
        global noisy_drives_debug
        if noisy_drives_debug == None:
            noisy_drives_debug = open('/var/log/mythtv/noisy-drives-debug.log', 'w', 1)
            dprint(program_name + ' debug output started')

    def dprint(s):
        global noisy_drives_debug
        if noisy_drives_debug != None:
            print(datetime.datetime.now().strftime('%H:%M:%S.%f')[:-3] + ' ' + s, file=noisy_drives_debug)

else:

    def dprint_init():
        pass
    def dprint():
        pass


##############################################################################
# Read config.xml file to get database settings only (other settings ignored).
##############################################################################

class ConfigXml:

    DBHostName = 'localhost'
    DBUsername = 'mythtv'
    DBPassword = 'mythtv'
    DBName = 'mythconverg'
    DBPort = 3306

    _conf_trans = {
        'Host':'DBHostName',
        'UserName':'DBUserName',
        'Password':'DBPassword',
        'DatabaseName':'DBName',
        'Port':'DBPort'
        }


    def readXML(self, filename):
        if not os.access(filename, os.R_OK):
            dprint('File ' + filename + ' not accessible!')
            return False

        try:
            config = lxml.etree.parse(filename)

            for child in config.xpath('/Configuration/Database')[0].getchildren():
                if child.tag in self._conf_trans:
                    #dprint('child.tag=' + str(child.tag))
                    setattr(self, self._conf_trans[child.tag], child.text)
        except Exception as e:
            dprint(str(e))
            dprint('lxml.etree.parse failed')
            return False
        return True


##############################################################################
# Backend class.  Handles communication with a backend and database selected
# using a config.xml file.  If no config.xml file is provided, the local
# database and backend are used.
##############################################################################

class Backend():

    def __init__(self):
        pass

    def init(self, config_xml=''):
        dprint('config_xml="' + config_xml + '"')
        self.config_xml = config_xml
        if self.config_xml != '':
            self.remote_config = ConfigXml()
            if not self.remote_config.readXML(self.config_xml):
                dprint('Failed to read or correctly parse ' + self.config_xml + ' file, using ' + HOST)
                self.config_xml = ''
                self.host = HOST
            # dprint(str(self.remote_config.DBHostName))
            # dprint(str(self.remote_config.DBUsername))
            # dprint(str(self.remote_config.DBPassword))
            # dprint(str(self.remote_config.DBName))
            # dprint(str(self.remote_config.DBPort))
            self.host = self.remote_config.DBHostName
            dprint('self.host = ' + self.host)
        else:
            self.host = HOST

        self.backend = api.Send(self.host)

        # Initialise the UTCOffset so that the conversion to local time works.
        api.Send(self.host)
        util.get_utc_offset(self.backend)

        # Connect to the database.
        if self.config_xml != '':
            self.db = MythTV.MythDB(
                DBHostName = self.remote_config.DBHostName,
                DBUsername = self.remote_config.DBUsername,
                DBPassword = self.remote_config.DBPassword,
                DBName =     self.remote_config.DBName,
                DBPort =     self.remote_config.DBPort
                )
        else:
            self.db = MythTV.MythDB()
        self.dbc = self.db.cursor()
        self.dbc1 = self.db.cursor()


##############################################################################
# Process a time argument.
##############################################################################

def time_arg(arg):
    return time.strftime(TIME_FORMAT, time.strptime(arg, TIME_FORMAT))


##############################################################################
# Run a system command asyncronously.
##############################################################################

async def run(cmd):
    proc = await asyncio.create_subprocess_shell(
        cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE
        )

    out, err = await proc.communicate()

    return proc, out, err


##############################################################################
# Set the hdparm -S sleep timeout settings for a drive.  If waking up the
# drive, also ensure it is spun up and the partition is mounted.
##############################################################################

async def set_sleep_timeout(drive, hdparm_S):
    if not os.path.ismount(drive):
        result, out, err = await run('mount ' + drive)
        if result.returncode != 0:
            dprint('Error: mount ' + drive + ' failed, return code ' + str(result.returncode) + '  Output: ' +
                out.decode() + ' ' + err.decode())
        else:
            dprint('Partition ' + drive + ' was mounted correctly')
    command = 'dev-mounted ' + drive
    result, out, err = await run(command)
    device = out.decode().rstrip('\n')
    if device == '':
        dprint('Error ' + drive + ' not mounted ')
    else:
        command = 'hdparm -S ' + str(hdparm_S) + ' ' + device
        result, out, err = await run(command)
        if result.returncode != 0:
            dprint('Error: hdparm -S failed, return code ' + str(result.returncode) + '  Output: ' +
                out.decode() + ' ' + err.decode())
        else:
            dprint('Sleep setting for ' + drive + ' (' + device + ') was set correctly')
            dprint('Output: ' + out.decode() + ' ' + err.decode())
        if hdparm_S == 0:
            # Do an uncached raw read from the drive to spin it up.
            command = 'dd if=' + device + ' bs=4096 count=1 of=/dev/null iflag=direct'
            result, out, err = await run(command)
            if result.returncode != 0:
                dprint('Error: dd command to spin up drive failed, return code ' + str(result.returncode) + '  Output: ' +
                    out.decode() + ' ' + err.decode())
            else:
                dprint('Drive ' + drive + ' (' + device + ') was spun up correctly')
                dprint('Output: ' + out.decode() + ' ' + err.decode())


##############################################################################
# Set the hdparm -S sleep timeout settings for the drives mounted on list of
# mount points, in parallel.
##############################################################################

async def set_sleep_timeouts(hdparm_S):
    tasks = []
    for drive in dirname_match:
        task = asyncio.create_task(set_sleep_timeout(drive, hdparm_S))
        tasks.append(task)

    await asyncio.gather(*tasks)


##############################################################################
# Check for a match up to the length of the left argument.
##############################################################################

def match_left(left: str, right: str):
    match_len = len(right)
    left_len = len(left)
    if left_len < match_len:
        match_len = left_len
    return left[0:match_len].lower() == right[0:match_len].lower()


##############################################################################
# Main
##############################################################################

class Force(Enum):
    Do_Not_Force = 0
    Quiet = 1
    Noisy = 2

dprint_init()

force: Force = Force.Do_Not_Force
quiet_start = time_arg(DEFAULT_QUIET_START)
quiet_stop = time_arg(DEFAULT_QUIET_STOP)

args = len(sys.argv)
if args == 2:
    if match_left(sys.argv[1], 'noisy'):
        force = Force.Noisy
    elif match_left(sys.argv[1], 'quiet'):
        force = Force.Quiet
if force == Force.Do_Not_Force:
    if args > 3:
        print('Error: Too many arguments!')
        exit(1)
    if args > 1:
        quiet_start = time_arg(sys.argv[1])
    if args == 3:
        quiet_stop = time_arg(sys.argv[2])

now = time.strftime(TIME_FORMAT)
dprint(f'{now=}')
dprint(f'{quiet_start=}')
dprint(f'{quiet_stop= }')
dprint(f'{force=}')

be = Backend()
be.init('/etc/mythtv/config.xml')

if force == Force.Quiet or force == Force.Do_Not_Force and (
    (quiet_start <= quiet_stop and now >= quiet_start and now < quiet_stop) or
    (quiet_start > quiet_stop and ((now >= quiet_start and now <= '23:59:59') or (now < quiet_stop)))
    ):
    dprint('Quiet')

    sql = "SELECT * FROM storagegroup WHERE LEFT(groupname, 7) = 'Default' and ("
    for match in dirname_match:
        sql += "LEFT(dirname, " + str(len(match)) + ") = '" + match + "' or "
    dprint('sql=' + sql)
    sql = sql[:-4:] + ')'
    dprint('sql=' + sql)
    result = be.dbc.execute(sql)
    if result == 0:
        dprint('Noisy drives not in default')
    else:
        for fetched in be.dbc:
            sql = "UPDATE storagegroup SET groupname='Noisy' WHERE id=" + str(fetched[0])
            dprint(sql)
            be.dbc1.execute(sql)
        asyncio.run(set_sleep_timeouts(SLEEP_TIMEOUT))

elif force != Force.Quiet:
    dprint('Noisy')
    result = be.dbc.execute("SELECT * from storagegroup WHERE groupname='Noisy'")
    if result == 0:
        dprint('Noisy group is empty')
    else:
        sql = "UPDATE storagegroup SET groupname='Default' WHERE groupname='Noisy' and not dirname like '%_crw-pvr%'"
        dprint(sql)
        be.dbc1.execute(sql)
        sql = "UPDATE storagegroup SET groupname='Default_crw-pvr' WHERE groupname='Noisy' and dirname like '%_crw-pvr%'"
        dprint(sql)
        be.dbc1.execute(sql)
        asyncio.run(set_sleep_timeouts(0))