'''
'''
#  Licensed to the Apache Software Foundation (ASF) under one
#  or more contributor license agreements.  See the NOTICE file
#  distributed with this work for additional information
#  regarding copyright ownership.  The ASF licenses this file
#  to you under the Apache License, Version 2.0 (the
#  "License"); you may not use this file except in compliance
#  with the License.  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

from __future__ import print_function
import os
from ports import get_port


def make_id(s):
    return s.replace(".", "_").replace('-', '_')


# A mapping from log type to the log name. 'stdout' and 'stderr' are handled
# specially and are used to indicate the stdout and stderr streams,
# respectively.
default_log_data = {'diags': 'diags.log', 'error': 'error.log', 'manager': 'manager.log'}

# 'block_for_debug', if True, causes traffic_server to run with the --block option enabled, and effectively
# disables timeouts that could be triggered by running traffic_server under a debugger.


def MakeATSProcess(
        obj,
        name,
        command='traffic_server',
        select_ports=True,
        enable_tls=False,
        enable_cache=True,
        enable_quic=False,
        block_for_debug=False,
        log_data=default_log_data,
        use_traffic_out=True, enable_proxy_protocol=False):
    #####################################
    # common locations

    # directory we will setup for the ATS to run under
    ts_dir = os.path.join(obj.RunDirectory, name)
    # common bin directory
    bin_dir = 'bin'

    # manually set up all directory for the test

    # configuration directory
    config_dir = os.path.join(ts_dir, 'config')

    # directory contains the html response templates
    template_dir = os.path.join(config_dir, 'body_factory')

    # contains plugins
    plugin_dir = os.path.join(ts_dir, 'plugin')

    # the log directory
    log_dir = os.path.join(ts_dir, 'log')

    # runtime dir
    runtime_dir = os.path.join(ts_dir, 'runtime')

    #ssl, storage & cache
    ssl_dir = os.path.join(ts_dir, 'ssl')
    storage_dir = os.path.join(ts_dir, 'storage')
    cache_dir = os.path.join(ts_dir, 'cache')

    ts_args = ''
    if use_traffic_out:
        # Bind stdout/err to traffic.out. This allows tests to wait upon
        # content from the traffic.out file, something that cannot be done on
        # process stdout/stderr in AuTest.
        traffic_out = os.path.join(log_dir, 'traffic.out')
        ts_args += f' --bind_stderr {traffic_out}'
        ts_args += f' --bind_stdout {traffic_out}'

    if block_for_debug:
        ts_args += ' --block'

    if 'traffic_manager' in command:
        if '--tsArgs' in command:
            # None of our tests do this, but if we eventually pass in --tsArgs
            # at some point from a test, we'll need to figure out how to mix
            # that with the arguments we're assembling for --tsArgs here.
            host.WriteError('Do not know how to assemble --tsArgs.')
        if ts_args:
            command += f' --tsArgs "{ts_args}"'
    elif 'traffic_server' in command:
        command += ts_args

    # create process
    p = obj.Processes.Process(name, command)
    #p_debug = obj.Processes.Process("port-debug", "ss --listen --tcp --process")
    #p_debug.Env['PATH'] = "/usr/sbin" + os.pathsep + p.ComposeEnv()['PATH']
    # p.StartBefore(p_debug)
    # we want to have a few directories more fixed
    # this helps with debugging as location are common
    # we do this by overriding locations from the "layout"
    # used as part of build. This means loctaion such as
    # PROXY_CONFIG_BIN_PATH with always be $root/bin
    # not something else such as bin64
    #####

    # set root for this test
    p.Env['TS_ROOT'] = ts_dir
    p.Setup.MakeDir(ts_dir)

    # set bin location

    p.Env['PROXY_CONFIG_BIN_PATH'] = bin_dir
    bin_path = os.path.join(ts_dir, bin_dir)
    p.Env['PATH'] = bin_path + os.pathsep + p.ComposeEnv()['PATH']
    p.Setup.Copy(p.Variables.BINDIR, bin_path, CopyLogic.SoftFiles)

    # setup config directory
    AddMethodToInstance(p, chownForATSProcess)

    # copy all basic config files we need to get this to work
    cfg_dir = os.path.join(AUTEST_SITE_PATH, "min_cfg")

    p.Setup.MakeDir(config_dir)
    p.chownForATSProcess(config_dir)

    for f in os.listdir(cfg_dir):
        p.Setup.CopyAs(os.path.join(cfg_dir, f), config_dir)

    #########################################################
    # setup config directory
    p.Env['PROXY_CONFIG_CONFIG_DIR'] = config_dir
    p.Variables.CONFIGDIR = config_dir
    #########################################################
    # setup read-only data directory in config. Needed for response body
    # responses

    p.Env['PROXY_CONFIG_BODY_FACTORY_TEMPLATE_SETS_DIR'] = template_dir
    p.Variables.BODY_FACTORY_TEMPLATE_DIR = template_dir
    p.Setup.Copy(os.path.join(p.Variables.SYSCONFDIR, 'body_factory'), template_dir)

    #########################################################
    # setup cache directory
    p.Setup.MakeDir(cache_dir)
    p.Env['PROXY_CONFIG_CACHE_DIR'] = cache_dir
    p.Variables.CACHEDIR = cache_dir

    #########################################################
    # setup read-only data directory for plugins

    p.Env['PROXY_CONFIG_PLUGIN_PLUGIN_DIR'] = plugin_dir
    p.Setup.Copy(p.Variables.PLUGINDIR, plugin_dir, CopyLogic.SoftFiles)

    #########################################################
    # create subdirectories that need to exist (but are empty)
    # log directory has to be created with correct permissions
    p.Setup.MakeDir(log_dir)  # log directory has to be created
    p.chownForATSProcess(log_dir)

    # set env so traffic server uses correct locations
    p.Env['PROXY_CONFIG_LOG_LOGFILE_DIR'] = log_dir
    p.Variables.LOGDIR = log_dir

    # this is needed for cache and communication sockets
    # Below was to make shorter paths but the code in
    # set runtime directory(local state dir)
    p.Env['PROXY_CONFIG_LOCAL_STATE_DIR'] = runtime_dir
    p.Env['PROXY_CONFIG_HOSTDB_STORAGE_PATH'] = runtime_dir
    p.Variables.RUNTIMEDIR = runtime_dir
    p.Variables.LOCALSTATEDIR = runtime_dir

    p.Setup.MakeDir(runtime_dir)
    p.chownForATSProcess(runtime_dir)

    # will need this for traffic_manager is it runs
    p.Setup.MakeDir(os.path.join(config_dir, 'snapshots'))
    p.Env['PROXY_CONFIG_SNAPSHOT_DIR'] = os.path.join(config_dir, 'snapshots')
    ##########################################################
    # create subdirectories that need to exist (but are empty)
    # ssl directory has to be created for keeping certs and keys
    p.Setup.MakeDir(ssl_dir)
    p.chownForATSProcess(ssl_dir)

    # set env so traffic server uses correct locations
    p.Env['PROXY_CONFIG_SSL_DIR'] = ssl_dir
    p.Variables.SSLDir = ssl_dir
    AddMethodToInstance(p, addSSLfile)

    # Add default cert folder & register handy method.
    AddMethodToInstance(p, addDefaultSSLFiles)
    AddMethodToInstance(p, addSSLFileFromDefaultTestFolder)
    ########################################################
    # cache.db directory
    p.Setup.MakeDir(storage_dir)
    p.chownForATSProcess(storage_dir)

    # set env so traffic server uses correct locations
    p.Env['PROXY_CONFIG_STORAGE_DIR'] = storage_dir
    p.Variables.STORAGEDIR = storage_dir
    #########################################################
    # define the basic file for a given test run
    # traffic.out ?? # cannot find it at the moment...
    # squid.log
    fname = "squid.log"
    tmpname = os.path.join(log_dir, fname)
    p.Disk.File(tmpname, id=make_id(fname))
    # error.log
    fname = log_data['error']
    if fname == 'stdout':
        p.Disk.error_log = p.Streams.stdout
    elif fname == 'stderr':
        p.Disk.error_log = p.Streams.stderr
    else:
        tmpname = os.path.join(log_dir, fname)
        p.Disk.File(tmpname, id='error_log')
    # diags.log
    fname = log_data['diags']
    if fname == 'stdout':
        p.Disk.diags_log = p.Streams.stdout
    elif fname == 'stderr':
        p.Disk.diags_log = p.Streams.stderr
    else:
        tmpname = os.path.join(log_dir, fname)
        p.Disk.File(tmpname, id='diags_log')
    # add this test back once we have network namespaces working again
    p.Disk.diags_log.Content = Testers.ExcludesExpression("ERROR:", f"Diags log file {fname} should not contain errors")
    p.Disk.diags_log.Content += Testers.ExcludesExpression("FATAL:", f"Diags log file {fname} should not contain errors")
    p.Disk.diags_log.Content += Testers.ExcludesExpression(
        "Unrecognized configuration value",
        f"Diags log file {fname} should not contain a warning about an unrecognized configuration")

    # traffic.out
    fname = "traffic.out"
    tmpname = os.path.join(log_dir, fname)
    p.Disk.File(tmpname, id='traffic_out')

    if "traffic_manager" in command:
        fname = log_data['manager']
        if fname == 'stdout':
            p.Disk.manager_log = p.Streams.stdout
        elif fname == 'stderr':
            p.Disk.manager_log = p.Streams.stderr
        else:
            tmpname = os.path.join(log_dir, fname)
            p.Disk.File(tmpname, id='manager_log')

    # config files
    def MakeConfigFile(self, fname):
        tmpname = os.path.join(config_dir, fname)
        return self.File(tmpname, id=make_id(fname), typename="ats:config")

    AddMethodToInstance(p.Disk, MakeConfigFile)

    # "Core" config files are pre-defined as variables.
    fname = "records.config"
    tmpname = os.path.join(config_dir, fname)
    p.Disk.File(tmpname, id=make_id(fname), typename="ats:config:records")

    fname = "cache.config"
    tmpname = os.path.join(config_dir, fname)
    p.Disk.File(tmpname, id=make_id(fname), typename="ats:config")

    fname = "hosting.config"
    tmpname = os.path.join(config_dir, fname)
    p.Disk.File(tmpname, id=make_id(fname), typename="ats:config")

    fname = "ip_allow.yaml"
    tmpname = os.path.join(config_dir, fname)
    p.Disk.File(tmpname, id=make_id(fname), typename="ats:config")

    fname = "logging.yaml"
    tmpname = os.path.join(config_dir, fname)
    p.Disk.File(tmpname, id=make_id(fname), typename="ats:config")

    fname = "parent.config"
    tmpname = os.path.join(config_dir, fname)
    p.Disk.File(tmpname, id=make_id(fname), typename="ats:config")

    fname = "plugin.config"
    tmpname = os.path.join(config_dir, fname)
    p.Disk.File(tmpname, id=make_id(fname), typename="ats:config")

    fname = "remap.config"
    tmpname = os.path.join(config_dir, fname)
    p.Disk.File(tmpname, id=make_id(fname), typename="ats:config")

    fname = "socks.config"
    tmpname = os.path.join(config_dir, fname)
    p.Disk.File(tmpname, id=make_id(fname), typename="ats:config")

    fname = "splitdns.config"
    tmpname = os.path.join(config_dir, fname)
    p.Disk.File(tmpname, id=make_id(fname), typename="ats:config")

    fname = "ssl_multicert.config"
    tmpname = os.path.join(config_dir, fname)
    p.Disk.File(tmpname, id=make_id(fname), typename="ats:config")

    fname = "sni.yaml"
    tmpname = os.path.join(config_dir, fname)
    p.Disk.File(tmpname, id=make_id(fname), typename="ats:config")

    fname = "storage.config"
    tmpname = os.path.join(config_dir, fname)
    p.Disk.File(tmpname, id=make_id(fname), typename="ats:config")

    fname = "volume.config"
    tmpname = os.path.join(config_dir, fname)
    p.Disk.File(tmpname, id=make_id(fname), typename="ats:config")

    ##########################################################
    # set up default ports
    # get some ports  TODO make it so we can hold on to the socket
    if select_ports:

        # some system have a bug in which ipv4 and ipv6 share port space
        # Make two different ports to avoid this
        get_port(p, "port")
        get_port(p, "portv6")

        if enable_tls:
            get_port(p, "ssl_port")
            get_port(p, "ssl_portv6")

        if enable_proxy_protocol:
            get_port(p, "proxy_protocol_port")
            get_port(p, "proxy_protocol_portv6")

            if enable_tls:
                get_port(p, "proxy_protocol_ssl_port")
                get_port(p, "proxy_protocol_ssl_portv6")
    else:
        p.Variables.port = 8080
        p.Variables.portv6 = 8080

        if enable_tls:
            p.Variables.ssl_port = 4443
            p.Variables.ssl_portv6 = 4444

    get_port(p, "manager_port")
    get_port(p, "admin_port")

    if enable_tls or enable_quic:
        fname = "tls_session_keys.txt"
        tmpname = os.path.join(log_dir, fname)
        p.Disk.File(tmpname, id='tls_session_keys')
        p.Disk.records_config.update({
            'proxy.config.ssl.keylog_file': tmpname,
        })

    if enable_cache:
        # In records.config, the cache is enabled by default so there's nothing
        # we have to do here to functionally enable it. However, the tests that
        # rely upon the cache will not function correctly if ATS starts
        # processing traffic before the cache is ready. Thus we set the
        # wait_for_cache configuration.
        p.Disk.records_config.update(
            {
                # Do not accept connections from clients until cache subsystem is
                # operational.
                'proxy.config.http.wait_for_cache': 1,
            })
    else:
        # The user wants the cache to be disabled.
        p.Disk.records_config.update({'proxy.config.http.cache.http': 0})

    if enable_quic:
        p.Disk.records_config.update({
            'proxy.config.udp.threads': 1,
        })

    # The following message was added so that tests and users can know when
    # Traffic Server is ready to both receive and optimize traffic.
    p.Ready = When.FileContains(p.Disk.diags_log.AbsPath, "NOTE: Traffic Server is fully initialized")

    if select_ports:
        # default config
        port_str = "{port} {v6_port}:ipv6".format(port=p.Variables.port, v6_port=p.Variables.portv6)
        if enable_tls:
            port_str += " {ssl_port}:ssl {ssl_portv6}:ssl:ipv6".format(
                ssl_port=p.Variables.ssl_port, ssl_portv6=p.Variables.ssl_portv6)
        if enable_quic:
            port_str += " {ssl_port}:quic {ssl_portv6}:quic:ipv6".format(
                ssl_port=p.Variables.ssl_port, ssl_portv6=p.Variables.ssl_portv6)
        if enable_proxy_protocol:
            port_str += f" {p.Variables.proxy_protocol_port}:pp {p.Variables.proxy_protocol_portv6}:pp:ipv6"
            if enable_tls:
                port_str += f" {p.Variables.proxy_protocol_ssl_port}:pp:ssl {p.Variables.proxy_protocol_ssl_portv6}:pp:ssl:ipv6"
        #p.Env['PROXY_CONFIG_HTTP_SERVER_PORTS'] = port_str
        p.Disk.records_config.update({
            'proxy.config.http.server_ports': port_str,
        })

    p.Env['PROXY_CONFIG_PROCESS_MANAGER_MGMT_PORT'] = str(p.Variables.manager_port)
    p.Env['PROXY_CONFIG_ADMIN_SYNTHETIC_PORT'] = str(p.Variables.admin_port)
    p.Env['PROXY_CONFIG_ADMIN_AUTOCONF_PORT'] = str(p.Variables.admin_port)  # support pre ATS 6.x

    if 'traffic_manager' in command:
        p.ReturnCode = 2
    else:
        p.ReturnCode = 0

    if block_for_debug:
        p.Env['PROXY_BLOCK'] = '1'
        # Cause traffic_server to wait effectviely indefinitely (10 hours) for the debugger to attach.
        p.StartupTimeout = 10 * 60 * 60
        # Hitting breakpoints may cause long delays in traffic_server transaction processing, so make
        # the timeout for test processes (that may interact with traffic_server) effectively forever
        # (10 hours).
        obj.Variables.Autest.Process.TimeOut = 10 * 60 * 60

    return p


##################################
# added to ats process object to help deal with config files


class Config(File):
    '''
    Class to represent a config file
    '''

    def __init__(self, runable, name, exists=None, size=None, content_tester=None, execute=False, runtime=True, content=None):
        super(Config, self).__init__(runable, name, exists=None, size=None, content_tester=None, execute=False, runtime=True)

        self.content = content
        self._added = False

    def AddLines(self, lines):
        for line in lines:
            self.AddLine(line)

    def _do_write(self, name):
        '''
        Write contents to disk
        '''
        host.WriteVerbosef('ats-config-file', "Writing out file {0}", self.Name)
        if self.content is not None:
            with open(name, 'w') as f:
                f.write(self.content)
        return (True, "Appended file {0}".format(self.Name), "Success")

    def AddLine(self, line):
        if not self._added:
            self.WriteCustomOn(self._do_write)
            self._added = True
        if self.content is None:
            self.content = ""
        if not line.endswith('\n'):
            line += '\n'
        self.content += line


class RecordsConfig(Config, dict):
    '''
    Create a "dict" representation of records.config

    This can be accessed as a multi-level dictionary

    such as:
    rc['CONFIG']['proxy.config.log.hostname']
    '''

    reverse_kind_map = {
        str: 'STRING',
        int: 'INT',
        float: 'FLOAT',
    }

    line_template = 'CONFIG {name} {kind} {val}\n'

    def __init__(self, runable, name, exists=None, size=None, content_tester=None, execute=False, runtime=True):
        super(RecordsConfig, self).__init__(runable, name, exists=None, size=None, content_tester=None, execute=False, runtime=True)
        self.WriteCustomOn(self._do_write)

    def _do_write(self, name):
        host.WriteVerbosef('ats-config-file', "Writing out file {0}", name)
        if len(self) > 0:
            with open(name, 'w') as f:
                for name, val in self.items():
                    f.write(self.line_template.format(name=name, kind=self.reverse_kind_map[type(val)], val=val))
        return (True, "Writing config file {0}".format(os.path.split(self.Name)[-1]), "Success")


##########################################################################


def chownForATSProcess(self, path):
    self.Setup.Chown(path, "nobody", "nobody", ignore=True)
    # covers ubuntu's unprivileged group
    self.Setup.Chown(path, "nobody", "nogroup", ignore=True)


def addSSLfile(self, filename):
    self.Setup.CopyAs(filename, self.Variables.SSLDir)


def addDefaultSSLFiles(self):
    addSSLfile(self, os.path.join(self.Variables.AtsTestToolsDir, "ssl", "server.pem"))
    addSSLfile(self, os.path.join(self.Variables.AtsTestToolsDir, "ssl", "server.key"))


def addSSLFileFromDefaultTestFolder(self, filename):
    addSSLfile(self, os.path.join(self.Variables.AtsTestToolsDir, "ssl", filename))


RegisterFileType(Config, "ats:config")
RegisterFileType(RecordsConfig, "ats:config:records")
ExtendTestRun(chownForATSProcess, name="ChownForATSProcess")
ExtendTest(MakeATSProcess, name="MakeATSProcess")
ExtendTestRun(MakeATSProcess, name="MakeATSProcess")