#!/usr/bin/python3
# -*- coding: utf-8 -*-
#
# Copyright (C) 2017 Linaro Limited
#
# Author: Remi Duraffort <remi.duraffort@linaro.org>
#
# This file is part of LAVA Dispatcher.
#
# LAVA Dispatcher is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# LAVA Dispatcher is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <http://www.gnu.org/licenses>.

# pylint: disable=wrong-import-order,superfluous-parens,too-many-statements
# pylint: disable=missing-docstring

import argparse
import errno
import logging
import os
from setproctitle import setproctitle
import signal
import sys
import traceback
import yaml

from lava_common.exceptions import InfrastructureError, JobCanceled, LAVABug, LAVAError
from lava_common.log import YAMLLogger
from lava_dispatcher.device import NewDevice
from lava_dispatcher.parser import JobParser


def parser():
    """ Configure the parser """
    p_obj = argparse.ArgumentParser()

    p_obj.add_argument(
        "--job-id",
        required=True,
        metavar="ID",
        help="Job identifier. " "This alters process name for easier debugging",
    )
    p_obj.add_argument(
        "--output-dir",
        required=True,
        metavar="DIR",
        help="Directory for temporary ressources",
    )
    p_obj.add_argument(
        "--validate",
        action="store_true",
        default=False,
        help="validate the job file, do not execute any steps. "
        "The description is saved into description.yaml",
    )

    group = p_obj.add_argument_group("logging")
    group.add_argument(
        "--logging-url",
        metavar="URL",
        default=None,
        help="URL of the ZMQ socket to send the logs to the master",
    )
    # zmq.auth.load_certificate expects a filename, avoid opening the file here
    group.add_argument(
        "--master-cert", default=None, metavar="PATH", help="Master certificate file"
    )
    # zmq.auth.load_certificate expects a filename, avoid opening the file here
    group.add_argument(
        "--slave-cert", default=None, metavar="PATH", help="Slave certificate file"
    )
    group.add_argument(
        "--socks-proxy", type=str, default=None, help="Connect using a socks proxy"
    )
    group.add_argument("--ipv6", action="store_true", default=False, help="Enable IPv6")

    group = p_obj.add_argument_group("configuration files")
    group.add_argument(
        "--device",
        metavar="PATH",
        type=argparse.FileType("r"),
        required=True,
        help="Device configuration",
    )
    group.add_argument(
        "--dispatcher",
        metavar="PATH",
        type=argparse.FileType("r"),
        default=None,
        help="Dispatcher configuration",
    )
    group.add_argument(
        "--env-dut",
        metavar="PATH",
        type=argparse.FileType("r"),
        default=None,
        help="DUT environment",
    )

    p_obj.add_argument(
        "--debug",
        action="store_true",
        default=False,
        help="Start remote pdb right before running the job, for debugging",
    )

    p_obj.add_argument("definition", type=argparse.FileType("r"), help="job definition")

    return p_obj


def setup_logger(options):
    # Pipeline always log as YAML so change the base logger.
    # Every calls to logging.getLogger will now return a YAMLLogger
    logging.setLoggerClass(YAMLLogger)

    # The logger can be used by the parser and the Job object in all phases.
    logger = logging.getLogger("dispatcher")
    if options.logging_url is not None:
        if options.master_cert and options.slave_cert:
            if not os.path.exists(options.master_cert) or not os.path.exists(
                options.slave_cert
            ):
                return None
        # pylint: disable=no-member
        logger.addZMQHandler(
            options.logging_url,
            options.master_cert,
            options.slave_cert,
            options.job_id,
            options.socks_proxy,
            options.ipv6,
        )
    else:
        logger.addHandler(logging.StreamHandler())

    return logger


def parse_job_file(logger, options):
    """
    Uses the parsed device_config instead of the old Device class
    so it can fail before the Pipeline is made.
    Avoids loading all configuration for all supported devices for every job.
    """
    # Prepare the pipeline from the file using the parser.
    device = None  # secondary connections do not need a device
    device = NewDevice(options.device)
    parser = JobParser()

    # Load the configuration files (this should *not* fail)
    env_dut = None
    if options.env_dut is not None:
        env_dut = options.env_dut.read()
    dispatcher_config = None
    if options.dispatcher is not None:
        dispatcher_config = options.dispatcher.read()

    # Generate the pipeline
    return parser.parse(
        options.definition.read(),
        device,
        options.job_id,
        logger=logger,
        dispatcher_config=dispatcher_config,
        env_dut=env_dut,
    )


def cancelling_handler(*_):
    """
    Catches most signals and raise JobCanceled (inherit from LAVAError).
    The exception will go through all the stack frames cleaning and reporting
    the error.
    """
    signal.signal(signal.SIGHUP, terminating_handler)
    signal.signal(signal.SIGINT, terminating_handler)
    signal.signal(signal.SIGQUIT, terminating_handler)
    signal.signal(signal.SIGTERM, terminating_handler)
    signal.signal(signal.SIGUSR1, terminating_handler)
    signal.signal(signal.SIGUSR2, terminating_handler)
    raise JobCanceled("The job was canceled")


def terminating_handler(*_):
    """
    Second signal handler to notify to the user that the job was canceled twice
    """
    raise JobCanceled("The job was canceled again (too long to cancel)")


def dump_as_safe_yaml(obj):
    """
    Dump as yaml without using any tags.
    Breaks the implicit dependency between lava_scheduler_app and
    lava_dispatcher python modules.
    Note that the resulting description will differ from the pipeline
    references used in the dispatcher test suite.
    We have to use a specific emitter that will remove tags on the fly.
    """
    process_tag = yaml.emitter.Emitter.process_tag

    def replace_tags(self, *args, **kwargs):
        if self.event.tag.startswith("tag:yaml.org,2002:python/name:"):
            self.event.value = self.event.tag.replace(
                "tag:yaml.org,2002:python/name:", ""
            )
            self.event.tag = "tag:yaml.org,2002:str"
        return process_tag(self, *args, **kwargs)

    yaml.emitter.Emitter.process_tag = replace_tags
    data = yaml.dump(obj)
    yaml.emitter.Emitter.process_tag = process_tag
    return data


def main():
    # Parse the command line
    options = parser().parse_args()

    # Check that we are running as root
    if os.geteuid() != 0:
        print("lava-run should be executed as root")
        return 1

    # Set process title for easier debugging
    setproctitle("lava-run [job: %s]" % options.job_id)

    # Setup the logger as early as possible
    logger = setup_logger(options)
    if not logger:
        print("lava-run failed to setup logging")
        return 1

    # By default, that's a failure
    description = ""
    success = False
    try:
        # Set the signal handler
        signal.signal(signal.SIGHUP, cancelling_handler)
        signal.signal(signal.SIGINT, cancelling_handler)
        signal.signal(signal.SIGQUIT, cancelling_handler)
        signal.signal(signal.SIGTERM, cancelling_handler)
        signal.signal(signal.SIGUSR1, cancelling_handler)
        signal.signal(signal.SIGUSR2, cancelling_handler)

        # Should be an absolute directory
        options.output_dir = os.path.abspath(options.output_dir)
        # Create the output directory
        try:
            os.makedirs(options.output_dir, mode=0o755)
        except OSError as exc:
            if exc.errno != errno.EEXIST:
                raise InfrastructureError("Unable to create %s" % options.output_dir)

        # Parse the definition and create the job object
        job = parse_job_file(logger, options)
        description = dump_as_safe_yaml(job.describe())

        job.validate()
        if not options.validate:
            if options.debug:
                from remote_pdb import set_trace

                set_trace()
            job.run()

    except LAVAError as exc:
        error_help = exc.error_help
        error_msg = str(exc)
        error_type = exc.error_type
    except BaseException as exc:  # pylint: disable=broad-except
        logger.exception(traceback.format_exc())
        error_help = LAVABug.error_help
        error_msg = str(exc)
        error_type = LAVABug.error_type
    else:
        success = True
    finally:
        result_dict = {"definition": "lava", "case": "job"}
        if success:
            result_dict["result"] = "pass"
            logger.info("Job finished correctly")
        else:
            result_dict["result"] = "fail"
            result_dict["error_msg"] = error_msg
            result_dict["error_type"] = error_type
            logger.error(error_help)
        logger.results(result_dict)  # pylint: disable=no-member

    # Closing the socket. We are now sure that all messages where sent.
    logger.close()  # pylint: disable=no-member

    # Generate the description
    # A missing description.yaml indicate that lava-run crashed
    description_file = os.path.join(options.output_dir, "description.yaml")
    with open(description_file, "w") as f_describe:
        f_describe.write(description)

    return 0 if success else 1


if __name__ == "__main__":
    sys.exit(main())
