#!/usr/bin/env python3

import json
import os
import signal
import socket
import traceback

from registry import Registry, RegistryError


SOCKET_PATH = os.environ.get(
    "MANAGED_CONTAINERD_SOCKET",
    "/run/managed-containerd/managed-containerd.sock",
)

server_socket = None
registry = None


def handle_request(req):
    action = req.get("action")

    if action == "ping":
        return {
            "ok": True,
            "daemon": "managed-containerd",
            "protocol_version": 1,
        }

    if action == "registry.list":
        return {
            "ok": True,
            "resources": registry.list_resources(),
        }

    resource = req.get("resource")

    if not resource:
        return {
            "ok": False,
            "error": "resource is required",
        }

    return registry.execute(
        resource,
        action,
    )


def cleanup_socket():
    try:
        os.unlink(SOCKET_PATH)
    except FileNotFoundError:
        pass


def shutdown(signum, frame):
    global server_socket

    print(
        f"managed-containerd received signal {signum}; shutting down",
        flush=True,
    )

    if server_socket is not None:
        try:
            server_socket.close()
        except OSError:
            pass

    cleanup_socket()
    raise SystemExit(0)


def serve():
    global server_socket
    global registry

    registry = Registry()

    os.makedirs(
        os.path.dirname(SOCKET_PATH),
        exist_ok=True,
    )

    cleanup_socket()

    signal.signal(signal.SIGTERM, shutdown)
    signal.signal(signal.SIGINT, shutdown)

    server_socket = socket.socket(
        socket.AF_UNIX,
        socket.SOCK_STREAM,
    )

    server_socket.bind(SOCKET_PATH)
    os.chmod(SOCKET_PATH, 0o660)
    server_socket.listen(16)

    print(
        f"managed-containerd listening on {SOCKET_PATH}",
        flush=True,
    )

    print(
        "registered resources: "
        + ", ".join(registry.resources),
        flush=True,
    )

    try:
        while True:
            conn, _ = server_socket.accept()

            with conn:
                try:
                    raw = b""

                    while b"\n" not in raw:
                        chunk = conn.recv(65536)

                        if not chunk:
                            break

                        raw += chunk

                    req = json.loads(
                        raw.decode("utf-8").strip()
                    )

                    response = handle_request(req)

                except (
                    RegistryError,
                    ValueError,
                    KeyError,
                ) as exc:
                    response = {
                        "ok": False,
                        "error": str(exc),
                    }

                except Exception as exc:
                    traceback.print_exc()

                    response = {
                        "ok": False,
                        "error": str(exc),
                    }

                conn.sendall(
                    json.dumps(
                        response,
                        default=str,
                    ).encode("utf-8")
                    + b"\n"
                )

    finally:
        cleanup_socket()


if __name__ == "__main__":
    serve()