#!/usr/bin/python3
# SPDX-License-Identifier: GPL-2.0-only
# Copyright (C) 2025-2026 Guillaume Tucker

"""Containerized builds"""

import abc
import argparse
import dataclasses
import logging
import os
import pathlib
import shutil
import subprocess
import sys
import tomllib
import uuid


class ContainerRuntime(abc.ABC):
    """Base class for a container runtime implementation"""

    name = None  # Property defined in each implementation class

    def __init__(self, args, config, logger):
        def _first_not_none(*args):
            return next((item for item in args if item is not None))

        self._uid = _first_not_none(
            args.uid, config.uid, os.getuid()
        )
        self._gid = _first_not_none(
            args.gid, config.gid, args.uid, config.uid, os.getgid()
        )
        self._env_file = args.env_file or config.env_file
        self._shell = args.shell
        self._logger = logger

    @classmethod
    def is_present(cls):
        """Determine whether the runtime is present on the system"""
        return shutil.which(cls.name) is not None

    @abc.abstractmethod
    def _do_run(self, image, cmd, container_name):
        """Runtime-specific handler to run a command in a container"""

    @abc.abstractmethod
    def _do_abort(self, container_name):
        """Runtime-specific handler to abort a running container"""

    def run(self, image, cmd):
        """Run a command in a runtime container"""
        container_name = str(uuid.uuid4())
        self._logger.debug("container: %s", container_name)
        try:
            return self._do_run(image, cmd, container_name)
        except KeyboardInterrupt:
            self._logger.error("user aborted")
            self._do_abort(container_name)
            return 1


class CommonRuntime(ContainerRuntime):
    """Common logic for Docker and Podman"""

    def _do_run(self, image, cmd, container_name):
        cmdline = [self.name, 'run']
        cmdline += self._get_opts(container_name)
        cmdline.append(image)
        cmdline += cmd
        self._logger.debug('command: %s', ' '.join(cmdline))
        return subprocess.call(cmdline)

    def _get_opts(self, container_name):
        opts = [
            '--name', container_name,
            '--rm',
            '--volume', f'{pathlib.Path.cwd()}:/src',
            '--workdir', '/src',
        ]
        if self._env_file:
            opts += ['--env-file', self._env_file]
        if self._shell:
            opts += ['--interactive', '--tty']
        return opts

    def _do_abort(self, container_name):
        subprocess.call([self.name, 'kill', container_name])


class DockerRuntime(CommonRuntime):
    """Run a command in a Docker container"""

    name = 'docker'

    def _get_opts(self, container_name):
        return super()._get_opts(container_name) + [
            '--user', f'{self._uid}:{self._gid}'
        ]


class PodmanRuntime(CommonRuntime):
    """Run a command in a Podman container"""

    name = 'podman'

    def _get_opts(self, container_name):
        return super()._get_opts(container_name) + [
            '--userns', f'keep-id:uid={self._uid},gid={self._gid}',
        ]


class Runtimes:
    """List of all supported runtimes"""

    runtimes = [PodmanRuntime, DockerRuntime]

    @classmethod
    def get_names(cls):
        """Get a list of all the runtime names"""
        return list(runtime.name for runtime in cls.runtimes)

    @classmethod
    def get(cls, name):
        """Get a single runtime class matching the given name"""
        for runtime in cls.runtimes:
            if runtime.name == name:
                if not runtime.is_present():
                    raise ValueError(f"runtime not found: {name}")
                return runtime
        raise ValueError(f"unknown runtime: {name}")

    @classmethod
    def find(cls):
        """Find the first runtime present on the system"""
        for runtime in cls.runtimes:
            if runtime.is_present():
                return runtime
        raise ValueError("no runtime found")


@dataclasses.dataclass
class Config:
    """Container configuration"""
    image: str = None
    runtime: str = None
    env_file: str = None
    uid: int = None
    gid: int = None

    @classmethod
    def from_toml(cls, config_file_path, profile_name):
        """Create a config object from a TOML file"""
        if not config_file_path:
            config_file_path = '.container.toml'
            if not os.path.exists(config_file_path):
                return cls()
        elif not os.path.exists(config_file_path):
            raise ValueError(f"config file not found: {config_file_path}")
        with open(config_file_path, 'rb') as config_file:
            config = tomllib.load(config_file)
        default = config.get('DEFAULT', {})
        if not profile_name:
            profile = {}
        else:
            profile = config.get(profile_name)
            if profile is None:
                raise ValueError(f"unknown profile: {profile_name}")
        kwargs = {
            name: type(value) for (name, type, value) in (
                (op.name, op.type, profile.get(op.name, default.get(op.name)))
                for op in dataclasses.fields(cls)
            ) if value is not None
        }
        return cls(**kwargs)


def _get_logger(verbose):
    """Set up a logger with the appropriate level"""
    logger = logging.getLogger('container')
    handler = logging.StreamHandler()
    handler.setFormatter(logging.Formatter(
        fmt='[container {levelname}] {message}', style='{'
    ))
    logger.addHandler(handler)
    logger.setLevel(logging.DEBUG if verbose is True else logging.INFO)
    return logger


def main(args):
    """Main entry point for the container tool"""
    logger = _get_logger(args.verbose)
    try:
        config = Config.from_toml(args.config_file, args.config_profile)
        runtime = args.runtime or config.runtime
        cls = Runtimes.get(runtime) if runtime else Runtimes.find()
    except ValueError as ex:
        logger.error(ex)
        return 1
    logger.debug("runtime: %s", cls.name)
    image = args.image or config.image
    if not image:
        logger.error("no image specified")
        return 1
    logger.debug("image: %s", image)
    return cls(args, config, logger).run(image, args.cmd)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        'container',
        description="See the documentation for more details: "
        "https://docs.kernel.org/dev-tools/container.html"
    )
    parser.add_argument(
        '-c', '--config-file',
        help="Path to the config file.  If not specified, the default is to "
        "look for .container.toml in the current working directory."
    )
    parser.add_argument(
        '-e', '--env-file',
        help="Path to an environment file to load in the container."
    )
    parser.add_argument(
        '-g', '--gid', type=int,
        help="Group ID to use inside the container."
    )
    parser.add_argument(
        '-i', '--image',
        help="Container image name."
    )
    parser.add_argument(
        '-p', '--config-profile',
        help="Profile section to use in the config file.  This will override "
        "any values defined in the DEFAULT section."
    )
    parser.add_argument(
        '-r', '--runtime', choices=Runtimes.get_names(),
        help="Container runtime name.  If not specified, the first one found "
        "on the system will be used i.e. Podman if present, otherwise Docker."
    )
    parser.add_argument(
        '-s', '--shell', action='store_true',
        help="Run the container in an interactive shell."
    )
    parser.add_argument(
        '-u', '--uid', type=int,
        help="User ID to use inside the container.  If the -g option is not "
        "specified and no group ID is defined in the configuration file, the "
        "user ID will also be set as the group ID."
    )
    parser.add_argument(
        '-v', '--verbose', action='store_true',
        help="Enable verbose output."
    )
    parser.add_argument(
        'cmd', nargs='+',
        help="Command to run in the container"
    )
    sys.exit(main(parser.parse_args(sys.argv[1:])))
