#!/usr/bin/env python3
import argparse
import re
import sys
from datetime import timedelta
from pathlib import Path

import yaml


TIME_REGEX = re.compile(r'((?P<hours>\d+?)h)?((?P<minutes>\d+?)m)?((?P<seconds>\d+?)s)?')

def parse_time(time_str):
    parts = TIME_REGEX.match(time_str)
    if not parts or parts.span() == (0, 0):
        return
    parts = parts.groupdict()
    time_params = {}
    for name, param in parts.items():
        if param:
            time_params[name] = int(param)
    return timedelta(**time_params)

class IniFile:
    PREAMBLE = "### Generated by dcqc - do not edit ###"
    def __init__(self):
        self._data = {}

    def __getitem__(self, key: str):
        return self._data[key]

    def __setitem__(self, key: str, value):
        self._data[key] = value

    def write(self, f):
        f.write(self.PREAMBLE)
        f.write("\n\n")
        for section, section_data in self._data.items():
            f.write(f"[{section}]\n")
            for key, value in section_data.items():
                if isinstance(value, list):
                    for item in value:
                        f.write(f"{key}={item}\n")
                else:
                    f.write(f"{key}={value}\n")
            f.write("\n")

def abort(msg: str):
    print(msg, file=sys.stderr)
    sys.exit(1)

def enforce_list(key: str, dict_: dict):
    if not isinstance(dict_[key], list):
        abort(f"{key} is not a list. Offending dict: {dict_}")


def add_part_of(args, unit_file: IniFile):
    unit_file["Unit"]["PartOf"] = args.target

def write_file(ini_file: IniFile, path: Path):
    with open(path, "w", encoding="utf-8") as f:
        ini_file.write(f)

def write_build_unit(args, yaml_dict: dict, service_name: str):
    out_file = args.output_dir / f"{service_name}.build"
    print(f'Generating build "{service_name}" ({out_file})')
    build = yaml_dict["services"][service_name]["build"]

    unit_file = IniFile()
    unit_file["Unit"] = {}
    unit_file["Unit"]["Description"] = f"{service_name.capitalize()} build"

    unit_file["Build"] = {}
    unit_file["Build"]["ImageTag"] = f"localhost/{service_name}"
    if isinstance(build, str):
        unit_file["Build"]["SetWorkingDirectory"] = build
    elif isinstance(build, dict):
        if "context" in build:
            unit_file["Build"]["SetWorkingDirectory"] = build["context"]
        if "dockerfile" in build:
            unit_file["Build"]["File"] = build["dockerfile"]

    write_file(unit_file, out_file)


def write_network_units(args, yaml_dict):
    networks = yaml_dict.get("networks", {})
    if "default" not in networks:
        networks["default"] = None
    for network_name in networks:
        out_file = args.output_dir / f"{network_name}.network"
        print(f'Generating network "{network_name}" ({out_file})')
        network = networks.get(network_name)
        if network is None:
            network = {}

        unit_file = IniFile()
        unit_file["Unit"] = {}
        unit_file["Unit"]["Description"] = f"{network_name.capitalize()} network"
        add_part_of(args, unit_file)

        unit_file["Network"] = {}
        if "name" in network:
            unit_file["Network"]["NetworkName"] = network_name
        unit_file["Network"]["Internal"] = str(network.get("internal", True)).lower()
        if "enable_ipv6" in network:
            unit_file["Network"]["IPv6"] = str(network["enable_ipv6"]).lower()
        unit_file["Network"]["DisableDNS"] = "false"
        unit_file["Network"]["Options"] = "isolate=true"

        unit_file["Service"] = {}
        unit_file["Service"]["Restart"] = "on-failure"

        unit_file["Install"] = {}
        unit_file["Install"]["WantedBy"] = "default.target"

        write_file(unit_file, out_file)


def write_service_units(args, yaml_dict):
    for service_name in yaml_dict["services"]:
        out_file = args.output_dir / f"{service_name}.container"
        print(f'Generating container "{service_name}" ({out_file})')
        service = yaml_dict["services"][service_name]

        unit_file = IniFile()
        unit_file["Unit"] = {}
        unit_file["Unit"]["Description"] = f"{service_name.capitalize()} container"
        add_part_of(args, unit_file)
        # currently decreases reliability instead of increasing it
        # if "depends_on" in service:
        #     enforce_list("depends_on", service)
        #     for dependency in service["depends_on"]:
        #         unit_file["Unit"].setdefault("Requires", []).append(f"{dependency}.service")
        #         unit_file["Unit"].setdefault("After", []).append(f"{dependency}.service")

        unit_file["Container"] = {}
        if "container_name" in service:
            unit_file["Container"]["ContainerName"] = service["container_name"]
        if "image" in service:
            unit_file["Container"]["Image"] = service["image"]
            unit_file["Container"]["AutoUpdate"] = "registry"
        elif "build" in service:
            write_build_unit(args, yaml_dict, service_name)
            unit_file["Container"]["Image"] = f"{service_name}.build"
            unit_file["Container"]["AutoUpdate"] = "local"
        if "ports" in service:
            enforce_list("ports", service)
            unit_file["Container"]["PublishPort"] = service["ports"]
        if "command" in service:
            enforce_list("command", service)
            command = service["command"]
            cmd = f'"{command[0]}"'
            if len(command) > 1:
                cmd += " \\\n"
            for i, cmd_part in enumerate(command[1:]):
                cmd += f'  "{cmd_part}"'
                if i < len(command) - 2:
                    cmd += " \\\n"
            unit_file["Container"]["Exec"] = cmd
        if "entrypoint" in service:
            if not isinstance(service["entrypoint"], str):
                abort(f"Only supporting str for entrypoint for now. Got {type(service['entrypoint'])}")
            unit_file["Container"]["Entrypoint"] = service["entrypoint"]
        if "labels" in service:
            enforce_list("labels", service)
            unit_file["Container"]["Label"] = service["labels"]
        if "environment" in service:
            if isinstance(service["environment"], dict):
                env_vars = []
                for key, value in service["environment"].items():
                    env_vars.append(f"{key}={value}")
                service["environment"] = env_vars
            enforce_list("environment", service)
            unit_file["Container"]["Environment"] = service["environment"]
        if "env_file" in service:
            unit_file["Container"]["EnvironmentFile"] = service["env_file"]
        if "volumes" in service:
            enforce_list("volumes", service)
            unit_file["Container"]["Volume"] = service["volumes"]
        if "networks" in service:
            enforce_list("networks", service)
            unit_file["Container"]["Network"] = list(map(lambda x: f"{x}.network", service["networks"]))
        else:
            unit_file["Container"]["Network"] = ["default.network"]
        if "devices" in service:
            unit_file["Container"]["AddDevice"] = service["devices"]
        if "userns_mode" in service:
            unit_file["Container"]["UserNS"] = service["userns_mode"]
        if "group_add" in service:
            enforce_list("group_add", service)
            unit_file["Container"]["GroupAdd"] = service["group_add"]
        if "security_opt" in service:
            enforce_list("security_opt", service)
            unmask = []
            mask = []
            for line in service["security_opt"]:
                parts = line.split(":")
                if len(parts) != 2:
                    abort(f'Can only parts security_opts of form "type:path". Got: {line}')
                type_ = parts[0]
                match type_:
                    case "mask":
                        mask.append(parts[1])
                    case "unmask":
                        unmask.append(parts[1])
                    case _:
                        abort(f'security_opt type "{type_}" is not supported"')
            if mask:
                unit_file["Container"]["Mask"] = mask
            if unmask:
                unit_file["Container"]["Unmask"] = unmask

        if "stop_grace_period" in service:
            unit_file["Container"]["StopTimeout"] = int(parse_time(service["stop_grace_period"]).total_seconds())


        unit_file["Service"] = {}
        unit_file["Service"]["Restart"] = "on-failure"

        unit_file["Install"] = {}
        unit_file["Install"]["WantedBy"] = "default.target"

        write_file(unit_file, out_file)


def parse_args():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("compose_file")
    parser.add_argument("-o", "--output-dir", help="output directory for unit files",
                        default="/etc/containers/systemd", type=Path)
    parser.add_argument("-t", "--target",
                        help="Add PartOf dependency to all unit files", default="containers.target")
    return parser.parse_args()


def main():
    args = parse_args()
    with open(args.compose_file, encoding="utf-8") as f:
        yaml_dict = yaml.safe_load(f)

    write_network_units(args, yaml_dict)
    write_service_units(args, yaml_dict)

    print('\n==> Remember to run "systemctl daemon-reload"')



if __name__ == '__main__':
    main()