#!/usr/bin/env python2
#==============================================================
#
#  Abs:   Script for initializing a directory in $EPICS_IOCS
#         for a LinuxRT CPU
#
#  Name:  createLinuxRTCpu.py
#
#  Rem:   Some features of this script are already implemented
#         in iocIsBorn.py.
#
#  Usage: createLinuxRTCpu <facility> <cpu>
#
#  Facility: SLAC Controls (Development)
#
#  Auth:  20-Aug-2020, A. Le (anle)
#--------------------------------------------------------------
#  Mods: (Latest to oldest)
#     31-Aug-2020, A. Le     (anle)
#         Add switch for tag name
#         Slight refactoring
#
#==============================================================
#
import argparse
import os
import shutil
import sys
import re
import globals
from six.moves import input

# Constants
IOC_PATH = os.path.expandvars('$EPICS_IOCS')
CPU_PATTERN = r'^cpu-[a-z\d]+-([a-z]{2})\d{2}'
CPU_PLACEHOLDER = '<cpu>'
ST_CMD_FILE = 'startup.cmd'
IOC_SPEC_REL_LINK = 'iocSpecificRelease'
EPICS_SPEC_REL_LINK = 'epicsSpecificRelease'
NEW_ST_CMD_MARKER = 'EPICS_CPU_APP'
SUBSYSTEM_TO_DIR = {
    'ky': 'Klystron',
    'ls': 'Laser',
    'mc': 'motion',
    'mg': 'Magnet',
    'pm': 'ProfileMonitorAD',
    'sp': 'SharedPlatform',
    'ts': 'Tpg',
    'fb': 'FastFeedback',
    'rf': 'SAA'
}
CPU_SUBS_EXCEPTIONS = {
    'cpu-gunb-rf01': 'scllrf',
}

def die(message):
    """Print an error message and quit."""
    print message
    if __name__ == '__main__':
        sys.exit(-1)

def getSubsystem(cpu):
    """Extract the subsystem from a CPU name."""
    m = re.match(CPU_PATTERN, cpu)
    if not m:
        die('Error: {} is an invalid CPU name'.format(cpu))

    return m.group(1)

def templateReplace(templatePath, placeholder, value, outputPath):
    """Perform template substitution and write to a new file.

    Args:
        templatePath (string): Path to template file
        placeholder (string): Placeholder to be substituted
        value (string): String with which to substitute
        outputPath (string): Path to output file

    """
    with open(templatePath) as f:
        content = f.read()
        result = content.replace(placeholder, value)

    with open(outputPath, 'w') as f:
        f.write(result)

def confirm(message):
    """Ask user for confirmation."""
    prompt = '{} (y/N) '.format(message)
    answer = input(prompt)
    while answer.lower() not in ('y', 'n', ''):
        print('Invalid option {}'.format(answer))
        answer = input(prompt)

    if answer.lower() in ('n', ''):
        print('Aborted')
        sys.exit(0)

def debug(message):
    """Print message visible in verbose mode."""
    if globals.verbose:
        print(message)

def deleteIfExists(path):
    """Delete a file or symlink if it exists."""
    if os.path.exists(path):
        if os.path.isfile(path):
            debug('Deleting file {}'.format(path))
            os.remove(path)
        elif os.path.islink(path):
            debug('Deleting symlink {}'.format(path))
            os.unlink(path)
        else:
            die('Error: {} is a directory'.format(path))

    else:
        debug('Warning: {} does not exist'.format(path))

def setupCpuDir(facility, cpu, appName, tagName):
    """Set up directory in $EPICS_IOCS for a LinuxRT CPU.
    
    Args:
        facility (string): The facility where the CPU is to be installed
        cpu (string): The LinuxRT CPU to be installed
        appName (string): The name of the application
        tagName (string): The tag name of the application

    Returns:
        bool: True if successful, False otherwise

    """
    # Navigate to CPU directory
    enterCpuDir(cpu)

    # Copy startup.cmd from template
    stCmdTemplate = '{}.linuxRT.cpu.{}'.format(ST_CMD_FILE, facility)
    stCmdTempPath = os.path.join(IOC_PATH, 'template', stCmdTemplate)
    if not os.path.isfile(stCmdTempPath):
        die('Error: {} does not exist'.format(stCmdTempPath))
        return False

    if os.path.islink(ST_CMD_FILE):
        debug('Warning: {} already exists as a symlink, removing'.format(ST_CMD_FILE))
        os.unlink(ST_CMD_FILE)

    debug('Copying {} into {}'.format(stCmdTempPath, ST_CMD_FILE))
    templateReplace(stCmdTempPath, CPU_PLACEHOLDER, cpu, ST_CMD_FILE)
    os.chmod(ST_CMD_FILE, 0o775)

    # Get subsystem directory
    if appName:
        # Get it from argname argument
        appDir = appName
    elif cpu in CPU_SUBS_EXCEPTIONS:
        # Or from exception list
        appDir = CPU_SUBS_EXCEPTIONS[cpu]
    else:
        # Or from CPU name
        subsystem = getSubsystem(cpu)
        if subsystem not in SUBSYSTEM_TO_DIR:
            die('Error: cannot determine app name for subsystem {}. Please specify the app name with --appname'.format(subsystem))
            return False

        appDir = SUBSYSTEM_TO_DIR[subsystem]

    # Specify app version
    if tagName:
        tag = tagName
    else:
        tag = 'current'

    debug('Using version {} of {}'.format(tag, appDir))

    # Create symbolic link to iocSpecificRelease
    iocSpecificReleaseDir = os.path.join('..', '..', 'iocTop', appDir, tag)
    if not os.path.isdir(iocSpecificReleaseDir):
        die('Error: {} does not exist'.format(iocSpecificReleaseDir))
        return False

    if os.path.islink(IOC_SPEC_REL_LINK):
        debug('Warning: {} already exists, replacing'.format(IOC_SPEC_REL_LINK))
        debug('Removing {}'.format(IOC_SPEC_REL_LINK))
        os.unlink(IOC_SPEC_REL_LINK)

    debug('Creating symlink to {}'.format(iocSpecificReleaseDir))
    os.symlink(iocSpecificReleaseDir, IOC_SPEC_REL_LINK)

    # Exit CPU directory
    exitCpuDir()
    print('Successfully set up boot links and startup script for {}'.format(cpu))
    return True

def cleanupCpuDir(cpu, fullClean):
    """Clean up a LinuxRT CPU's directory in $EPICS_IOCS.

    epicsSpecificRelease and startup.cmd are removed. kernel-modules.cmd is
    also deleted unless fullDelete is False

    Args:
        fullClean (bool): Also remove kernel-modules.cmd if True

    """
    # Navigate to CPU directory
    enterCpuDir(cpu)

    # Delete epicsSpecificRelease and startup.cmd
    deleteIfExists(EPICS_SPEC_REL_LINK)
    if os.path.isfile(ST_CMD_FILE):
        # Special case: only delete startup.cmd if it contains EPICS_CPU_APP
        with open(ST_CMD_FILE) as f:
            if NEW_ST_CMD_MARKER not in f.read():
                debug('Warning: {} does not contain {}, not deleted'.format(ST_CMD_FILE, NEW_ST_CMD_MARKER))
            else:
                deleteIfExists(ST_CMD_FILE)

    if fullClean:
        # Also delete kernel-modules.cmd
        deleteIfExists('kernel-modules.cmd')

    # Exit CPU directory
    exitCpuDir()

def enterCpuDir(cpu):
    """Change directory to $EPICS_IOCS/<cpu>."""
    globals.initialDir = os.getcwd()
    cpuPath = os.path.join(IOC_PATH, cpu)
    if not os.path.isdir(cpuPath):
        die('Error: {} does not exist'.format(cpuPath))
        return False

    debug('Entering {}'.format(cpuPath))
    os.chdir(cpuPath)

def exitCpuDir():
    """Return to initial directory when script was called."""
    debug('Exiting {}'.format(os.getcwd()))
    os.chdir(globals.initialDir)

if __name__ == '__main__':
    # Parse arguments
    parser = argparse.ArgumentParser(description='Set up directory in $EPICS_IOCS for a LinuxRT CPU')
    parser.add_argument('facility', help='The facility where the CPU is to be installed')
    parser.add_argument('cpu', help='The LinuxRT CPU to be installed')
    parser.add_argument('-v', '--verbose', action='store_true', help='Show verbose output')
    parser.add_argument('-f', '--force', action='store_true', help='Force the script to proceed. By default, ask user for confirmation')
    parser.add_argument('-a', '--appname', help='Specify the application name. If omitted, infer app name from CPU name')
    parser.add_argument('-t', '--tagname', help='Specify the tag name of your application. If omitted, use "current"')
    group = parser.add_mutually_exclusive_group()
    group.add_argument('-c', '--cleanup', action='store_true', help='Remove epicsSpecificRelease, startup.cmd, and kernel-modules.cmd if they exist')
    group.add_argument('-p', '--partial-cleanup', action='store_true', help='Remove epicsSpecificRelease and startup.cmd but not kernel-modules.cmd')
    
    args = parser.parse_args()
    globals.verbose = args.verbose

    if args.cleanup or args.partial_cleanup:
        if not args.force:
            confirm('Are you sure you want to clean up scripts for {}?'.format(args.cpu))

        # Clean up CPU directory
        cleanupCpuDir(args.cpu, args.cleanup)

    else:
        if not args.force:
            confirm('Are you sure you want to modify scripts for {} in {}?'.format(args.cpu, args.facility))

        # Set up CPU directory
        setupCpuDir(args.facility, args.cpu, args.appname, args.tagname)