#!/usr/bin/python3
#
# Copyright 2022 Red Hat Inc.
#
# 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 json
import os
import re
import shutil
import subprocess
import sys

from pkg_resources import packaging


SOCKET = "unix:/run/podman/podman.sock"
PS_FORMAT = ('{"service": "{{.Names}}", "container": "{{.ID}}", '
             '"status": "{{.State}}", "healthy": "{{.Status}}"},')
BASE_FORMAT = "{service: .Name, container: .Id, status: .State.Running, "
RUNNING_REX = re.compile(r"Up .*(?P<health>\(u?n?healthy\))")
SKIP_LIST = ['.*_bootstrap', 'container-puppet-.*', '.*_db_sync',
             '.*_cron', 'create_.*_wrapper', '.*_wait_bundle',
             'configure_.*', '.*_fix_perms', '.*_init_logs?',
             '.*_init_perm', '(nova|placement)_wait_for_.*', 'nova_.*ensure_',
             '(swift_)?setup_.*', 'mysql_data_ownership', 'swift_copy_rings',
             'nova_statedir_owner']


class ExecuteError(Exception):
    def __init__(self, rc, msg):
        self.rc = rc
        self.msg = msg
        super().__init__()


def execute(cmd, workdir: str = None,
            prev_proc: subprocess.Popen = None) -> subprocess.Popen:
    # Note(mmagr): When this script is executed by collectd-sensubility started
    #              via collectd the script has non-root permission but inherits
    #              environment from collectd with root permission. We need
    #              to avoid sensubility access /root when using podman-remote.
    #              See https://bugzilla.redhat.com/show_bug.cgi?id=2091076 for
    #              more info.
    proc_env = os.environ.copy()
    proc_env["HOME"] = "/tmp"
    if type(cmd[0]) is list:  # multiple piped commands
        last = prev_proc
        for c in cmd:
            last = execute(c, workdir, last)
        return last
    else:  # single command
        inpipe = prev_proc.stdout if prev_proc is not None else None
        proc = subprocess.Popen(cmd, cwd=workdir, env=proc_env, stdin=inpipe,
                                stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if prev_proc is not None:
            prev_proc.stdout.close()
            prev_proc.stderr.close()
        return proc


def get_health_from_status(status):
    healthy = status in ('healthy', 'running')
    return int(healthy)


def fetch_state_from_inspect(cont, fmt):
    proc = execute([
        [shutil.which('podman-remote'),
           '--url', SOCKET, 'inspect', cont["container"]],
        [shutil.which('jq'), '.[] | %s' % fmt]
    ])
    o, e = proc.communicate()
    if proc.returncode != 0:
        msg = "Failed to fetch status of %s: %s" % (cont, e.decode())
        raise ExecuteError(proc.returncode, msg)

    item = json.loads(o.decode())
    item['status'] = 'running' if item['status'] else 'stopped'
    if len(item['healthy']) > 0 and item['status'] != 'stopped':
        item['status'] = item['healthy']

    item['healthy'] = get_health_from_status(item['status'])
    return item


def fetch_state(item):
    if 'up ' in item['status'].lower():
        item['status'] = 'running'
    if item['status'] != 'running':
        item['status'] = 'stopped'

    match = RUNNING_REX.match(item['healthy'])
    if match:
        health = match.group('health')
    else:
        health = ''
    if health:
        item['status'] = health.strip("()")

    item['healthy'] = get_health_from_status(item['status'])
    return item


def fetch_container_health(ps_result, inspect_fmt):

    def container_filter(item):
        for rx in [re.compile(i) for i in SKIP_LIST]:
            if rx.match(item["service"]):
                return False
        return True
    try:
        containers = filter(container_filter, json.loads(ps_result))
    except json.decoder.JSONDecodeError:
        return 1, "%s\nFailed to parse above output." % ps_result

    out = []
    for cont in containers:
        try:
            out.append(fetch_state(cont))
        except Exception:
            # fallback to slow inspect way in case of any error
            try:
                out.append(fetch_state_from_inspect(cont, inspect_fmt))
            except ExecuteError as ex:
                return ex.rc, ex.msg
    return 0, out


if __name__ == "__main__":
    proc = execute([shutil.which('podman-remote'),
                    '--url', SOCKET, 'version',
                    '--format', r'{{.Server.Version}}'])
    o, e = proc.communicate()
    try:
        if packaging.version.parse(o.decode().strip()) >= packaging.version.parse("4.0.0"):
            inspect_fmt = BASE_FORMAT + "healthy: .State.Health.Status}"
        else:
            inspect_fmt = BASE_FORMAT + "healthy: .State.Healthcheck.Status}"
    except Exception:
        # keep podman-4.0.0+ format in case of version decoding error
        inspect_fmt = BASE_FORMAT + "healthy: .State.Health.Status}"

    proc = execute([shutil.which('podman-remote'), '--url', SOCKET,
                    'ps', '--all', '--format', PS_FORMAT])
    o, e = proc.communicate()
    if proc.returncode != 0:
        print("Failed to list containers:\n%s\n%s" % (o.decode(), e.decode()))
        sys.exit(1)

    rc, status = fetch_container_health("[%s]" % o.decode().strip("\n,"), inspect_fmt)
    if rc != 0:
        print("Failed to inspect containers:\n%s" % status)
        sys.exit(rc)
    print(json.dumps(status))