#!/usr/bin/python

import argparse

parser = argparse.ArgumentParser(description='Make custom Cyclecloud template.')
parser.add_argument('--name', nargs='?', required=True,
                    help='name of the template, creates <name>.txt file')
parser.add_argument('--scheduler', nargs='?', default='pbspro',
                    help='select the job scheduler pbspro(default),gridengine,slurm')
parser.add_argument('--array', nargs='+', default='execute',
                    help='specify the excute node arrays to create (default: execute)')
parser.add_argument('--nfs', nargs='+', help='connect to NFS export from CycleCloud NFS filer')
parser.add_argument('--anf', nargs='?', help='connect to ANF export from Azure')
parser.add_argument('--lustre', nargs='?', help='connect to ANF export from Azure')
parser.add_argument('--hpcbb', nargs='?', help='configure HPCBB alike cluster features')

args = parser.parse_args()
#print args


template = """# template generated with template generator #\n\n"""
template += """[cluster {}]\n""".format(args.name)
template += """FormLayout = selectionpanel\nCategory = My Clusters\nCategoryOrder=1\nAutoscale = $Autoscale\n"""
template += """
    [[node defaults]]
    UsePublicNetwork = $UsePublicNetwork
    Credentials = $Credentials
    ImageName = $ImageName
    SubnetId = $SubnetId
    Region = $Region
    KeyPairLocation = ~/.ssh/cyclecloud.pem
	[[[configuration]]]
    	cyclecloud.selinux.policy = permissive
"""
if args.scheduler == "slurm":
    template += """
	[[[configuration]]]
	slurm.version = $configuration_slurm_version
        [[[cluster-init cyclecloud/slurm:default]]]
  """
if args.nfs:
  template += """
        [[[configuration cyclecloud.mounts.other_cluster_fs]]]
        type = nfs
        mountpoint = $NFSMountPoint
        export_path = /mnt/exports/data
        cluster_name = $NFSServerName
  """

if args.anf:
  template += """
        [[[configuration cyclecloud.mounts.anf]]]
        type = nfs
        mountpoint = $ANFMountPoint
        export_path = $ANFExport
        address = $ANFServerName
"""

if args.lustre:
  template += """
        [[[cluster-init lustre:default:1.0.0]]]
        Optional = True
        [[[cluster-init lustre:client:1.0.0]]]
        Optional = True

        [[[configuration]]]
        lustre.client.cluster_name = $LustreClusterName
        lustre.client.mount_point = $LustreMountPoint

"""

template += """
    [[node master]]
    MachineType = $MasterMachineType
    IsReturnProxy = $ReturnProxy
    AdditionalClusterInitSpecs = $MasterClusterInitSpecs
        [[[configuration]]]
"""

if args.hpcbb == "true" :
    template += "        [[[cluster-init hpcbb:master]]]\n"

if args.scheduler == "pbspro":
    template += "        [[[cluster-init cyclecloud/pbspro:master]]]\n"
if args.scheduler == "gridengine":
    template += "        [[[cluster-init cyclecloud/gridengine:master]]]\n"
if args.scheduler == "slurm":
    template += "        [[[cluster-init cyclecloud/slurm:master]]]\n"

template += """
        [[[network-interface eth0]]]
        AssociatePublicIpAddress = $UsePublicNetwork

        [[[input-endpoint ganglia]]]
        PrivatePort = 8652
        PublicPort = 8652
"""
template += """
    [[nodearray execute]]
    MachineType = $ExecuteMachineType
    MaxCoreCount = $MaxExecuteCoreCount
    Interruptible = $UseLowPrio
    AdditionalClusterInitSpecs = $ExecuteClusterInitSpecs
        [[[configuration]]]
"""

if args.hpcbb == "true" :
    template += "        [[[cluster-init hpcbb:execute]]]\n"

if args.scheduler == "pbspro":
    template += "        [[[cluster-init cyclecloud/pbspro:execute]]]\n"
if args.scheduler == "gridengine":
    template += "        [[[cluster-init cyclecloud/gridengine:execute]]]\n"
if args.scheduler == "slurm":
    template += "        [[[cluster-init cyclecloud/slurm:execute]]]\n"

template += """
        [[[network-interface eth0]]]
        AssociatePublicIpAddress = $ExecuteNodesPublic
"""

if args.array == "execute":
    pass 
else:
    for array in args.array:
        template += "\n# specified array {}\n".format(array)
        template += "    [[nodearray {}]]".format(array)
	template += """
    Extends = execute
    [[[configuration]]]\n"""
        if args.scheduler == "pbspro":
            template += "    pbspro.slot_type = {}\n".format(array) 

template += """
[parameters Required Settings]
Order = 1

    [[parameters Template Description]]
"""
template += "    Description = \"This this cluster is set up to use the {} scheduler\"\n".format(args.scheduler)
template += """
    Order = 10

    [[parameters Location and network ]]
    Description = "The cluster, in this case, has two roles: the scheduler master-node with shared filer and the execute hosts. Configure which VM types to use based on the requirements of your application."
    Order = 20

        [[[parameter Region]]]
        Label = Region
        Description = Deployment Location
        ParameterType = Cloud.Region
        DefaultValue = westus2

        [[[parameter SubnetId]]]
        Label = Subnet ID
        Description = Subnet Resource Path (ResourceGroup/VirtualNetwork/Subnet)
        ParameterType = Azure.Subnet
        Required = True
"""
template += """
[parameters Node Settings]
Order = 20

    [[parameters Virtual Machines ]]
    Description = "The cluster, in this case, has two roles: the scheduler master-node with shared filer and the execute hosts. Configure which VM types to use based on the requirements of your application."
    Order = 20

        [[[parameter MasterMachineType]]]
        Label = Master VM Type
        Description = The VM type for scheduler master and shared filer.
        ParameterType = Cloud.MachineType
        DefaultValue = Standard_D12_v2

        [[[parameter ExecuteMachineType]]]
        Label = Execute VM Type
        Description = The VM type for execute nodes
        ParameterType = Cloud.MachineType
        DefaultValue = Standard_F2s_v2
        Config.Multiselect = true
"""
template += """
    [[parameters Software]]
    Description = "Specify the scheduling software, and base OS installed on all nodes, and optionally the cluster-init and chef versions from your Locker."
    Order = 10                                                                                                                         

        [[[parameter ImageName]]]
        Label = Base OS
        ParameterType = Cloud.Image
        Config.OS = linux
        DefaultValue = cycle.image.centos7
        Config.Filter := Package === "cycle.image.centos7"

        [[[parameter MasterClusterInitSpecs]]]
        Label = Master Cluster-Init
        DefaultValue = =undefined
        Description = Cluster init specs to apply to the master node
        ParameterType = Cloud.ClusterInitSpecs

        [[[parameter ExecuteClusterInitSpecs]]]
        Label = Execute Cluster-Init
        DefaultValue = =undefined
        Description = Cluster init specs to apply to execute nodes
        ParameterType = Cloud.ClusterInitSpecs
"""

template += """
    [[parameters Auto-Scaling]]
    Description = "The cluster can autoscale to the workload, adding execute hosts as jobs are queued. To enable this check the box below and choose the initial and maximum core counts for the cluster"
    Order = 30

        [[[parameter Autoscale]]]
        Label = Autoscale
        DefaultValue = true
        Widget.Plugin = pico.form.BooleanCheckBox
        Widget.Label = Start and stop execute instances automatically

        [[[parameter MaxExecuteCoreCount]]]
        Label = Max Cores
        Description = The total number of execute cores to start
        DefaultValue = 100
        Config.Plugin = pico.form.NumberTextBox
        Config.MinValue = 1
        Config.IntegerOnly = true

        [[[parameter UseLowPrio]]]
        Label = Low Priority
        DefaultValue = false
        Widget.Plugin = pico.form.BooleanCheckBox
        Widget.Label = Use low priority instances for execute hosts
"""
template += """
[parameters Advanced Settings]
Order = 20

    [[parameters Subscription Settings]]
    Order = 10

        [[[parameter Credentials]]]
        Description = The credentials for the cloud provider
        ParameterType = Cloud.Credentials
"""
template += """
    [[parameters Advanced Networking]]
    Description = Advanced networking settings

        [[[parameter ReturnProxy]]]
        Label = Return Proxy
        DefaultValue = true
        ParameterType = Boolean
        Config.Label = Use SSH tunnel to connect to CycleCloud (required if direct access is blocked)

        [[[parameter UsePublicNetwork]]]
        Label = Public Head Node
        DefaultValue = true
        ParameterType = Boolean
        Config.Label = Access master node from the Internet

        [[[parameter ExecuteNodesPublic]]]
        Label = Public Execute
        DefaultValue = false
        ParameterType = Boolean
        Config.Label = Access execute nodes from the Internet
        Conditions.Excluded := UsePublicNetwork isnt true
"""
if args.nfs or args.anf or args.lustre:
  template += """
[parameters External Filesystems]
Order = 30
"""

if args.nfs:
  template += """
    [[parameters NFS Settings]]
    Order = 26
    Description = "Use a NFS server as a NAS. Settings for defining the NFS"

        [[[parameter NFSServerName]]]
        Label = NFS server
        Description = Name of the NFS server to connect to. The NFS server should be orchestrated by the same CycleCloud Server
        Required = True
        Config.Plugin = pico.form.QueryDropdown
        Config.Query = select ClusterName as Name from Cloud.Node where Cluster().IsTemplate =!= True && ClusterInitSpecs["nfs:default"] isnt undefined
        Config.SetDefault = false

        [[[parameter NFSMountPoint]]]
        Label = NFS MountPoint
        Description = The mount point to mount the NFS file server on.
        DefaultValue = /nfsdata
        Required = True
"""
if args.anf:
  template += """
    [[parameters ANF Settings]]
    Order = 28
    Description = "Use a ANF server as a NAS. Settings for defining the ANF"

        [[[parameter ANFServerName]]]
        Label = ANF Server name or IP address
        Description = The ANF server name or ip address.
        DefaultValue = 10.1.3.4
        Required = True

        [[[parameter ANFExport]]]
        Label = ANF Export
        Description = The export on the ANF.
        DefaultValue = /netappdata
        Required = True

        [[[parameter ANFMountPoint]]]
        Label = ANF MountPoint
        Description = The mount point to mount the ANF file server on.
        DefaultValue = /netappdata
        Required = True
"""
if args.lustre:
  template += """
    [[parameters Lustre Settings]]
    Order = 28
    Description = "Use a Lustre cluster as a PFS. Settings for defining the Lustre cluster"

        [[[parameter LustreClusterName]]]
        Label = Lustre Cluster
        Description = Name of the Lustre cluster to connect to.
        Required = True
        Config.Plugin = pico.form.QueryDropdown
        Config.Query = select ClusterName as Name from Cloud.Node where Cluster().IsTemplate =!= True && ClusterInitSpecs["lustre:mds"] isnt undefined
        Config.SetDefault = false                                                                                    

        [[[parameter LustreMountPoint]]]
        Label = Lustre MountPoint
        Description = The mount point to mount the Lustre file server on.
        DefaultValue = /lustre
        Required = True

"""
if args.scheduler == "slurm":
    template += """
   [[parameters Slurm Settings]]
   Order = 40
   Description = "Settings specific to Slurm"

	[[[parameter configuration_slurm_version]]]
	required = True
	label = Slurm Version
	description = Version of Slurm to install on the cluster
	defaultvalue = 17.11.12-1
  """

filename = '{}.txt'.format(args.name)
f = open( filename, 'w' )
f.write( template )
f.close()

print ("template {}.txt written, to import:\ncyclecloud import_template -f {}.txt [--force]\n").format(args.name, args.name) 

