#!/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

PROTOCOL_VERSION = 1


def protocol_response(payload):
    """
    Return a daemon response carrying the wire-protocol version.

    protocol_version is daemon-owned metadata and cannot be overridden by
    registry or plugin output.
    """
    response = dict(payload)
    response["protocol_version"] = PROTOCOL_VERSION
    return response


def handle_request(req):
    if not isinstance(req, dict):
        return protocol_response({
            "ok": False,
            "error_code": "invalid_request",
            "error": "request must be a JSON object",
        })

    requested_version = req.get("protocol_version")

    # Backward compatibility for the pre-versioned proxy:
    # absence is accepted during the rolling upgrade, but an explicitly
    # incompatible version is rejected.
    if (
        requested_version is not None
        and requested_version != PROTOCOL_VERSION
    ):
        return protocol_response({
            "ok": False,
            "error_code": "unsupported_protocol_version",
            "error": (
                "unsupported protocol version: "
                f"{requested_version}; "
                f"supported={PROTOCOL_VERSION}"
            ),
        })

    action = req.get("action")

    if action == "ping":
        return protocol_response({
            "ok": True,
            "daemon": "managed-containerd",
        })

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

    resource = req.get("resource")

    if not resource:
        return protocol_response({
            "ok": False,
            "error_code": "invalid_request",
            "error": "resource is required",
        })

    return protocol_response(
        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 = protocol_response({
                        "ok": False,
                        "error_code": "request_error",
                        "error": str(exc),
                    })

                except Exception as exc:
                    traceback.print_exc()

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

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

    finally:
        cleanup_socket()


if __name__ == "__main__":
    serve()