#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# =============================================================================
## @file pplaunch
#  Helper script to launch pp-server on remote hosts
#   - builds ssh tunnel
#   - sets the remote environment
#   - starts the remote pp-server 
#  @author Vanya BELYAEV Ivan.Belyaev@itep.ru
#  @date   2011-06-07
# =============================================================================
"""Launch pp-server on remote hosts via ssh-tunnel
- SSH tunnels to remote hosts are created
- remote hosts are configured
- pp-servers start at remote hosts
"""
# =============================================================================
__all__ = ()
# =============================================================================
import ROOT, os
os.environ['OSTAP_ARGPARSE'] = 'False'

ROOT.PyConfig.IgnoreCommandLineOptions = True
# =============================================================================
## create the argument parser 
def make_parser () :
    """Create the argument parser
    """
    from argparse import ArgumentParser 
    parser = ArgumentParser ( prog = 'pplaunch' ,
                              description = __doc__ )
    
    group1 = parser.add_mutually_exclusive_group()    
    group1.add_argument ( 
        "-q" , "--quiet"       ,
        dest    = 'Quiet'      , 
        action  = 'store_true' ,
        help    = "Quite processing [default: %(default)s]" ,
        )    
    group1.add_argument ( 
        "-d" , "--debug"     ,
        dest    = 'Debug'    , 
        action  = 'store_true' ,
        help    = "Debug processing [default: %(default)s]" ,
        )
    group1.add_argument ( 
        "--verbose"     ,
        dest    = 'Verbose'    , 
        action  = 'store_true' ,
        help    = "Verbose processing [default: %(default)s]" ,
        )
    
    group2  = parser.add_argument_group ( "Configurtaion" )
    group2.add_argument (
        "-c" , "--config"  ,
        dest    = 'Config' ,
        type    =  str     , 
        help    = 'Configuration script to be transferred to remote host and sourced',
        )
    
    group2.add_argument (
        "-p" , "--profile"  ,
        dest    = 'Profile' ,
        type    =  str     , 
        help    = 'Remote script to be sourced',
        )
    
    parser.add_argument (
        "-s" , "--secret" ,
        dest    = 'Secret' ,
        type    =  str     , 
        help    = 'Secret password used for connection' ,
        )

    parser.add_argument (
        "-t" , "--timeout"  ,
        dest    = 'Timeout' ,
        type    =  int      , 
        help    = 'Timeout for SSH connection [default: %(default)s]' ,
        default = 7200 
        )
    
    parser.add_argument (
        "remotes" ,
        metavar = 'Remotes' ,
        nargs   = '+'       , 
        help    = "Remote hosts [*/auto/config] -> take from ostap-configuration" )
    
    return parser 

# =================================================================================
if '__main__' == __name__ :
    
    # =============================================================================
    from ostap.logger.logger import getLogger
    logger = getLogger ( 'ostap.pplaunch' )
    # =============================================================================
        
    parser = make_parser() 
    
    import sys, os 
    args = sys.argv[1:]

    config = parser.parse_args ( args )

    if not config.Quiet :
        logger.info ( "Launch remote pp-server via ssh tunnel:\n%s" % __doc__ ) 
    else : 
        logger.info ( "Launch remote pp-server via ssh tunnel")
   
    logger.debug ( 'Configuration:%s ' % config) 

    remotes = config.remotes 
    if 1 == len ( remotes ) :
        rems = remotes [ 0 ]
        rems = rems.lower()
        if rems in ( 'config' , '*' , 'auto' ) :
            from ostap.parallel.utils import get_ppservers
            import socket
            local_host = socket.getfqdn ().lower()  
            remotes    = get_ppservers ( local_host )
            assert remotes, 'No remotes are extracted from configuration'

    remotes = tuple ( remotes )
    if config.Verbose or config.Debug :
        logger.info ( "List of remote hosts: %s" % list ( remotes ) )
                      
    ## check alive remote hosts 
    from ostap.parallel.utils import good_pings 
    alive     = good_pings ( *remotes )
    if len ( alive ) != len ( remotes ) :
        dead = list ( set ( remotes ) - set ( alive ) )
        logger.warning ("Dead remote hosts: %s" % dead ) 
        remotes = tuple ( alive ) 

    assert remotes, 'No valid remote hosts are defined!'
    if config.Verbose or config.Debug :
        logger.info ( "List of good remote hosts: %s" % list ( remotes ) )

    kwargs = {} 
    if config.Config :
        assert os.path.exist  ( config.Config ) and \
               os.path.isfile ( config.Config ), \
               "Invalid config scrips %s is specifid " % config.Config        
        kwargs [ 'script' ] = config.Config
        if config.Verbose or config.Debug :
            logger.info ( "Configuration script: %s" % config.Script )
        
    if config.Profile :
        kwargs [ 'profile' ] = config.Profile
        if config.Verbose or config.Debug :
            logger.info ( "Configuration script: %s" % config.Script )

    secret = config.Secret
    if not secret :
        from ostap.utils.utils import gen_password 
        secret = gen_password ( 6 )
        logger.debug ( 'Secret if auto-generated: %s' % secret ) 
        kwargs [ 'secret' ] = secret 
        if config.Verbose or config.Debug:
            logger.info ( "Secret: %s" %  secret  )
        
    if config.Timeout and 0 < config.Timeout :
        kwargs [ 'timeout' ] = config.Timeout
        if config.Verbose or config.Debug:
            logger.info ( "Timeout: %s" % config.Timeout )
            
    kwargs [ 'silent' ] = config.Quiet
    if config.Verbose or config.Debug:
        logger.info ( "Configuration: %s" % kwargs  )
    
    from ostap.parallel.pptunnel import ppServer, show_tunnels 
    
    ppservers = []

    if   config.Verbose : 
        from ostap.logger.utils import logVerbose as Context 
    elif config.Debug   : 
        from ostap.logger.utils import logDebug   as Context 
    else :
        from ostap.logger.utils import NoContext  as Context 

    from ostap.logger.colorized import attention
    message = '\n' + attention( "Press <ENTER> to terminate pp-tunnels> " ) + '\n'

    ## show progress bar? 
    progress = 4 < len ( remotes ) and config.Quiet and ( not congif.Verbose and not config.Debug ) 
    silent   = not progress
    
    from ostap.utils.progress_bar import progress_bar 
    
    if not silent :
        strrem = str ( remotes )
        if len ( strrem ) < 50 :
            logger.info ( "Launching %s pp-tunnels: %s" % ( len  ( remotes ) , list ( remotes ) ) ) 
        else : 
            logger.info ( "Launching %s pp-tunnels at remote hosts:" % len  ( remotes ) )
            logger.info ( "%s"                                       % list ( remotes ) )
            
    if 3 <= sys.version_info.major :

        from contextlib              import ExitStack
        with Context() , ExitStack() as stack :
            for remote in progress_bar ( remotes , silent = silent ) : 
                ppsvc = ppServer ( remote , **kwargs )
                if not ppsvc.pid :
                    del ppsvc
                    continue 
                stack.enter_context ( ppsvc ) 
                ppservers.append ( ppsvc )

            if not ppservers : logger.fatal ("No valid pp-tunnels created!")  
            else : 
                show_tunnels ( ppservers ) 
                local_ports = [ p.local_port  for p in ppservers ]

                logger.info (80*'*')
                logger.info ("Tunnels: secret='%s', ppservers=%s" % ( secret , local_ports ) )
                logger.info (80*'*')
                
                input ( message )
                    
    else :

        with Context() :                
            for remote in progress_bar ( remotes , silent = silent ) : 
                ppsvc = ppServer ( remote , **kwargs )
                if not ppsvc.pid :
                    del ppsvc
                    continue 
                ppservers.append ( ppsvc )
                
            if not ppservers : logger.fatal ("No valid pp-tunnels created!")  
            else : 

                show_tunnels ( ppservers ) 
                local_ports = [ p.local_port  for p in ppservers ]
                
                logger.info (80*'*')
                logger.info ("Tunnels: secret='%s', ppservers=%s" % ( secret , local_ports ) )
                logger.info (80*'*')
                
                raw_input ( message )

            if not silent :
                logger.info ('Terminating %d pp-tunnels' % len ( ppservers ) )
                logger.info ("%s" % [ "%s->%s" % ( p.local_port , p.remote ) for p in ppservers ] ) 
            for ppsvc in progress_bar ( ppservers , silent = silent ) :
                ppsvc.end()
                
    
    while ppservers : ppservers.pop()
    logger.info (80*'*')

# =============================================================================
##                                                                      The END 
# =============================================================================
