#!/usr/bin/env python3
# Copyright 2023 Katteli Inc.
# TestFlows.com Open-Source Software Testing Framework (http://testflows.com)
#
# Licensed 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.
import os
import sys
import time
import hashlib
import logging
import threading

from concurrent.futures import ThreadPoolExecutor, Future

from hcloud import Client
from hcloud.ssh_keys.domain import SSHKey

from github import Github
from github.Repository import Repository

from argparse import ArgumentParser, RawTextHelpFormatter

from testflows.github.runners import __version__, __license__
from testflows.github.runners.actions import Action
from testflows.github.runners.scale_up import scale_up
from testflows.github.runners.scale_down import scale_down
from testflows.github.runners.scripts import Scripts, scripts

import testflows.github.runners.args as args
import testflows.github.runners.cloud as cloud
import testflows.github.runners.service as service

check_args = args.check
check_image = args.check_image

logger = logging.getLogger("testflows.github.runners")

description = """Auto-scaling GitHub Actions runners service using Hetzner Cloud.

    Service that starts and monitors queued up GitHub Actions workflows.
    When a new job is queued up, it creates new Hetzner Cloud server instance
    that provides an ephemeral GitHub Actions runner. Server is automatically
    powered off when job completes and then powered off servers are
    automatically deleted.

    By default, uses `$GITHUB_TOKEN`, `$GITHUB_REPOSITORY`, and `$HETZNER_TOKEN`
    environment variables or you can specify these values `--github-token`,
    `--github-repository`, and `--hetzner-token` options.
"""


def argparser():
    """Command line argument parser."""

    parser = ArgumentParser(
        "GitHub Actions runners scale up service",
        description=description,
        formatter_class=RawTextHelpFormatter,
    )

    parser.add_argument("-v", "--version", action="version", version=f"{__version__}")

    parser.add_argument(
        "--license",
        action="version",
        help="show program's license and exit",
        version=f"{__license__}",
    )

    parser.add_argument(
        "--github-token",
        type=args.env_var_type("GITHUB_TOKEN"),
        help="GitHub token, default: $GITHUB_TOKEN environment variable",
        default="",
    )

    parser.add_argument(
        "--github-repository",
        type=args.env_var_type("GITHUB_REPOSITORY"),
        help="GitHub repository, default: $GITHUB_REPOSITORY environment variable",
        default="",
    )

    parser.add_argument(
        "--hetzner-token",
        type=args.env_var_type("HETZNER_TOKEN"),
        help="Hetzner Cloud token, default: $HETZNER_TOKEN environment variable",
        default="",
    )

    parser.add_argument(
        "--ssh-key",
        metavar="path",
        type=str,
        help="path to public SSH key, default: ~/.ssh/id_rsa.pub",
        default=os.path.expanduser("~/.ssh/id_rsa.pub"),
    )

    parser.add_argument(
        "-m",
        "--max-runners",
        metavar="count",
        type=args.count_type,
        help="maximum number of active runners, default: unlimited",
    )

    parser.add_argument(
        "--default-image",
        metavar="type:name_or_description",
        type=args.image_type,
        help=(
            "default runner server image type and name or description,"
            "where the type is either: 'system','snapshot','backup','app',\n"
            "default: system:ubuntu-22.04"
        ),
        default="system:ubuntu-22.04",
    )

    parser.add_argument(
        "--default-type",
        metavar="name",
        type=args.server_type,
        help=("default runner server type name, default: cx11"),
        default="cx11",
    )

    parser.add_argument(
        "--default-location",
        metavar="name",
        type=args.location_type,
        help=("default runner server location name, by default not specified"),
    )

    parser.add_argument(
        "-w",
        "--workers",
        metavar="count",
        type=args.count_type,
        help="number of concurrent workers, default: 10",
        default=10,
    )

    parser.add_argument(
        "--logger-config",
        metavar="path",
        type=str,
        help="custom logger configuration file",
    )

    parser.add_argument(
        "--setup-script",
        metavar="path",
        type=str,
        help="path to custom server setup script",
    )

    parser.add_argument(
        "--startup-x64-script",
        metavar="path",
        type=str,
        help="path to custom x64 server startup script",
    )

    parser.add_argument(
        "--startup-arm64-script",
        metavar="path",
        type=str,
        help="path to custom ARM64 server startup script",
    )

    parser.add_argument(
        "--max-powered-off-time",
        metavar="sec",
        type=args.count_type,
        help="maximum time after which a powered off server is deleted, default: 20 sec",
        default=20,
    )

    parser.add_argument(
        "--max-idle-runner-time",
        metavar="sec",
        type=args.count_type,
        help="maximum time after which an idle runner is removed and its server deleted, default: 120 sec",
        default=120,
    )

    parser.add_argument(
        "--max-runner-registration-time",
        metavar="sec",
        type=args.count_type,
        help="maximum time after which the server will be deleted if it fails to register a runner, default: 60 sec",
        default=60,
    )

    parser.add_argument(
        "--scale-up-interval",
        metavar="sec",
        type=args.count_type,
        help="scale up service interval, default: 10 sec",
        default=10,
    )

    parser.add_argument(
        "--scale-down-interval",
        metavar="sec",
        type=args.count_type,
        help="scale down service interval, default: 10 sec",
        default=10,
    )

    parser.add_argument(
        "--debug",
        action="store_true",
        help="enable debugging mode, default: False",
        default=False,
    )

    commands = parser.add_subparsers(
        title="commands", metavar="command", description=None, help=None
    )

    cloud_parser = commands.add_parser(
        "cloud",
        help="cloud service commands",
        description="Deploying and running application as a service on a cloud instance.",
        formatter_class=RawTextHelpFormatter,
    )

    cloud_parser.add_argument(
        "-n",
        "--name",
        metavar="server",
        dest="server_name",
        type=str,
        help="deployment server name, default: github-runners",
        default="github-runners",
    )

    cloud_commands = cloud_parser.add_subparsers(
        title="commands", metavar="command", description=None, help=None
    )

    cloud_commands.required = True

    deploy_cloud_parser = cloud_commands.add_parser(
        "deploy",
        help="deploy cloud service",
        description="Deploy application as a service to a new server instance.",
        formatter_class=RawTextHelpFormatter,
    )

    deploy_cloud_parser.add_argument(
        "-f",
        "--force",
        action="store_true",
        help="force deployment if already exist",
    )

    deploy_cloud_parser.add_argument(
        "-l",
        "--location",
        metavar="name",
        type=args.location_type,
        help="deployment server location, default: ash",
        default="ash",
    )

    deploy_cloud_parser.add_argument(
        "-t",
        "--type",
        metavar="name",
        type=args.server_type,
        help="deployment server type, default: cpx11",
        default="cpx11",
    )

    deploy_cloud_parser.add_argument(
        "-i",
        "--image",
        metavar="type:name_or_description",
        type=args.image_type,
        help=(
            "deployment server image type and name or description,\n"
            "where the type is either: 'system','snapshot','backup','app',\n"
            "default: system:ubuntu-22.04"
        ),
        default="system:ubuntu-22.04",
    )

    deploy_cloud_parser.set_defaults(func=cloud.deploy)

    logs_cloud_parser = cloud_commands.add_parser(
        "logs",
        help="get cloud service logs",
        description="Get cloud service logs.",
        formatter_class=RawTextHelpFormatter,
    )

    logs_cloud_parser.add_argument(
        "-f",
        "--follow",
        action="store_true",
        help="follow the logs journal",
    )

    logs_cloud_parser.set_defaults(func=cloud.logs)

    status_cloud_parser = cloud_commands.add_parser(
        "status",
        help="get cloud service status",
        description="Get cloud service status.",
        formatter_class=RawTextHelpFormatter,
    )

    status_cloud_parser.set_defaults(func=cloud.status)

    start_cloud_parser = cloud_commands.add_parser(
        "start",
        help="start cloud service ",
        description="Start cloud service.",
        formatter_class=RawTextHelpFormatter,
    )

    start_cloud_parser.set_defaults(func=cloud.start)

    stop_cloud_parser = cloud_commands.add_parser(
        "stop",
        help="stop cloud service",
        description="Stop cloud service.",
        formatter_class=RawTextHelpFormatter,
    )

    stop_cloud_parser.set_defaults(func=cloud.stop)

    install_cloud_parser = cloud_commands.add_parser(
        "install",
        help="install cloud service",
        description=(
            "Install cloud service.\n\n"
            "The `github-runners <options> service install` command will be executed on the cloud server instance.\n\n"
            "Note: Just like with the `github-runners <options> service install` command,\n"
            "      the options that are passed to the `github-runners <options> cloud install` command\n"
            "      will be the same options with which the service will be executed."
        ),
        formatter_class=RawTextHelpFormatter,
    )

    install_cloud_parser.add_argument(
        "-f",
        "--force",
        action="store_true",
        help="force service install",
        default=False,
    )

    install_cloud_parser.set_defaults(func=cloud.install)

    uninstall_cloud_parser = cloud_commands.add_parser(
        "uninstall",
        help="uninstall cloud service",
        description="Uninstall cloud service.",
        formatter_class=RawTextHelpFormatter,
    )

    uninstall_cloud_parser.set_defaults(func=cloud.uninstall)

    upgrade_cloud_parser = cloud_commands.add_parser(
        "upgrade",
        help="upgrade cloud service",
        description=(
            "Upgrade cloud service application.\n\n"
            "If specific '--version' is specified then the `testflows.github.runners`\n"
            "package is upgraded to the specified version otherwise the version is\n"
            "upgraded to the latest available.\n\n"
            "Note: The service is not re-installed during the package upgrade process.\n"
            "      Instead, it is stopped before the upgrade and then started back up\n"
            "      after the package upgrade is complete."
        ),
        formatter_class=RawTextHelpFormatter,
    )

    upgrade_cloud_parser.add_argument(
        "--version",
        type=str,
        metavar="version",
        dest="upgrade_version",
        help="package version, default: the latest",
    )

    upgrade_cloud_parser.set_defaults(func=cloud.upgrade)

    delete_cloud_parser = cloud_commands.add_parser(
        "delete",
        help="delete cloud service",
        description=(
            "Delete cloud service.\n\n"
            "Deletes `github-runners` service running on the cloud server instance\n"
            "by first stopping the service and then deleting the server."
        ),
        formatter_class=RawTextHelpFormatter,
    )

    delete_cloud_parser.set_defaults(func=cloud.delete)

    service_parser = commands.add_parser(
        "service",
        help="service commands",
        description="Service commands.",
        formatter_class=RawTextHelpFormatter,
    )

    service_commands = service_parser.add_subparsers(
        title="commands", metavar="command", description=None, help=None
    )

    service_commands.required = True

    install_service_parser = service_commands.add_parser(
        "install",
        help="install",
        description=(
            "Install service.\n\n"
            "The `/etc/systemd/system/github-runners.service` file will be created and service will be started.\n\n"
            "Note: The options that are passed to the `github-runners <options> service install` command\n"
            "      will be the same options with which the service will be executed."
        ),
        formatter_class=RawTextHelpFormatter,
    )
    install_service_parser.add_argument(
        "-f",
        "--force",
        action="store_true",
        help="force service install",
        default=False,
    )
    install_service_parser.set_defaults(func=service.install)

    uninstall_service_parser = service_commands.add_parser(
        "uninstall",
        help="uninstall",
        description="Uninstall service.",
        formatter_class=RawTextHelpFormatter,
    )
    uninstall_service_parser.add_argument(
        "-f",
        "--force",
        action="store_true",
        help="force service uninstall",
        default=False,
    )
    uninstall_service_parser.set_defaults(func=service.uninstall)

    start_service_parser = service_commands.add_parser(
        "start",
        help="start",
        description="Start service.",
        formatter_class=RawTextHelpFormatter,
    )
    start_service_parser.set_defaults(func=service.start)

    stop_service_parser = service_commands.add_parser(
        "stop",
        help="stop",
        description="Stop service.",
        formatter_class=RawTextHelpFormatter,
    )
    stop_service_parser.set_defaults(func=service.stop)

    status_service_parser = service_commands.add_parser(
        "status",
        help="status",
        description="Get service status.",
        formatter_class=RawTextHelpFormatter,
    )
    status_service_parser.set_defaults(func=service.status)

    logs_service_parser = service_commands.add_parser(
        "logs",
        help="logs",
        description="Get service logs.",
        formatter_class=RawTextHelpFormatter,
    )
    logs_service_parser.add_argument(
        "-f",
        "--follow",
        action="store_true",
        help="follow the logs journal",
    )
    logs_service_parser.set_defaults(func=service.logs)

    return parser


def main(args, scripts: Scripts, worker_pool: ThreadPoolExecutor, terminate_timeout=30):
    """Auto-scale runners service."""

    terminate = threading.Event()

    try:
        with Action("Logging in to Hetzner Cloud"):
            client = Client(token=args.hetzner_token)

        with Action("Logging in to GitHub"):
            github = Github(login_or_token=args.github_token)

        with Action(f"Getting repository {args.github_repository}"):
            repo: Repository = github.get_repo(args.github_repository)

        with Action("Checking if default image exists"):
            args.default_image = check_image(client=client, image=args.default_image)

        with Action(f"Checking if SSH key exists"):
            with open(args.ssh_key, "r", encoding="utf-8") as ssh_key_file:
                public_key = ssh_key_file.read()
            key_name = hashlib.md5(public_key.encode("utf-8")).hexdigest()
            ssh_key = SSHKey(name=key_name, public_key=public_key)

            if not client.ssh_keys.get_by_name(name=ssh_key.name):
                with Action(f"Creating SSH key {ssh_key.name}"):
                    client.ssh_keys.create(
                        name=ssh_key.name, public_key=ssh_key.public_key
                    )

        try:
            with Action("Creating scale up service"):
                scale_up_service: Future = worker_pool.submit(
                    scale_up,
                    terminate=terminate,
                    repo=repo,
                    client=client,
                    scripts=scripts,
                    ssh_key=ssh_key,
                    default_type=args.default_type,
                    default_image=args.default_image,
                    default_location=args.default_location,
                    worker_pool=worker_pool,
                    github_token=args.github_token,
                    github_repository=args.github_repository,
                    interval=args.scale_up_interval,
                    max_servers=args.max_runners,
                )

            with Action("Creating scale down service"):
                scale_down_service: Future = worker_pool.submit(
                    scale_down,
                    terminate=terminate,
                    repo=repo,
                    client=client,
                    max_powered_off_time=args.max_powered_off_time,
                    max_idle_runner_time=args.max_idle_runner_time,
                    max_runner_registration_time=args.max_runner_registration_time,
                    interval=args.scale_down_interval,
                )

            while True:
                time.sleep(1)

                if scale_up_service.done():
                    raise RuntimeError("scale-up service exited")

                if scale_down_service.done():
                    raise RuntimeError("scale-down service exited")

        except BaseException:
            with Action("Requesting all services to terminate"):
                terminate.set()
            raise

        finally:
            with Action("Waiting for scale up service to terminate", ignore_fail=True):
                scale_down_service.result(timeout=terminate_timeout)

            with Action(
                "Waiting for scale down service to terminate", ignore_fail=True
            ):
                scale_up_service.result(timeout=terminate_timeout)

    except KeyboardInterrupt as exc:
        msg = "❗ KeyboardInterrupt"
        if args.debug:
            logger.exception(f"{msg}\n{exc}")
        else:
            logger.error(msg)
        sys.exit(1)
    except Exception as exc:
        msg = f"❗ Error: {type(exc).__name__} {exc}"
        if args.debug:
            logger.exception(f"{msg}\n{exc}")
        else:
            logger.error(msg)
        sys.exit(1)


if __name__ == "__main__":
    args = argparser().parse_args()

    logging_level = logging.INFO

    if args.debug:
        Action.debug = True
        logging_level = logging.DEBUG

    if args.logger_config:
        logging.config.fileConfig(args.logger_config)
    else:
        logger.setLevel(logging_level)
        handler = logging.StreamHandler(sys.stdout)
        formatter = logging.Formatter(
            "%(asctime)s %(levelname)8s %(threadName)16s %(funcName)15s %(message)s",
            datefmt="%m/%d/%Y %I:%M:%S %p",
        )
        handler.setFormatter(formatter)
        logger.addHandler(handler)

    if args.setup_script:
        scripts.setup = args.setup_script

    if args.startup_x64_script:
        scripts.startup_x64 = args.startup_x64_script

    if args.startup_arm64_script:
        scripts.startup_arm64 = args.startup_arm64_script

    if hasattr(args, "func"):
        try:
            args.func(args=args)
        except KeyboardInterrupt:
            if args.debug:
                raise
        except BaseException as exc:
            if args.debug:
                raise
            logger.fatal(f"❗ Error: {exc}")

    else:
        check_args(args)

        with ThreadPoolExecutor(
            max_workers=args.workers + 2, thread_name_prefix="worker"
        ) as worker_pool:
            main(args=args, scripts=scripts, worker_pool=worker_pool)
