"""Convert a `docker run ...` command string into a Compose YAML fragment.

The result has to be a compose file that ``docker compose config`` accepts and
that means what the command meant. Three things got that wrong before:

* every flag the converter did not know was assumed to take a value, so a
  boolean like ``--read-only`` or ``-P`` swallowed the image name, and
  ``-dit ubuntu bash`` came out as a service running the image ``bash``;
* ``--network host`` became ``networks: [host]`` and ``-v data:/data`` a
  reference to a volume nobody declared — both of which compose rejects;
* flags with real effect (``--gpus``, ``--mount``, ``--entrypoint``,
  ``--cap-drop``, ``--shm-size`` …) were dropped without a word.

``docker run``'s boolean flags are a short, closed list; everything else takes a
value. That is what the parser goes by now.
"""
from __future__ import annotations

import re
import shlex

import yaml

#: docker run flags that take no value.
_BOOL_LONG = {
    "--detach", "--interactive", "--tty", "--rm", "--init", "--privileged",
    "--read-only", "--publish-all", "--oom-kill-disable", "--no-healthcheck",
    "--disable-content-trust", "--quiet", "--sig-proxy", "--use-api-socket",
}
#: Single-letter booleans, which may be bundled ("-dit").
_BOOL_SHORT = set("ditPq")
#: Single-letter flags that take a value ("-p 80:80", "-p80:80").
_VALUE_SHORT = {
    "p": "--publish", "v": "--volume", "e": "--env", "l": "--label",
    "u": "--user", "w": "--workdir", "h": "--hostname", "m": "--memory",
    "a": "--attach", "c": "--cpu-shares",
}
_SHORT_TO_LONG = {"d": "--detach", "i": "--interactive", "t": "--tty",
                  "P": "--publish-all", "q": "--quiet"}

#: network_mode values, as opposed to networks a service joins.
_NETWORK_MODES = {"host", "none", "bridge"}

_SERVICE_NAME_RE = re.compile(r"[^a-zA-Z0-9._-]+")


def _service_name_from_image(image: str) -> str:
    name = image.split("@", 1)[0].split("/")[-1].split(":")[0]
    return _SERVICE_NAME_RE.sub("-", name).strip("-.") or "app"


#: A line continuation whose newline is already gone. Pasted into a one-line
#: field, backslash-newline arrives as a backslash followed by the next line's
#: indentation, which shlex read as an escaped space: a lone " " token that
#: became the image. A backslash followed by an option, or by more than one
#: blank, is such a continuation; a single escaped space inside a word
#: (``My\ Files``) is left alone.
_FLATTENED_CONTINUATION = re.compile(r"\\[ \t]+(?=-)|\\[ \t]{2,}")


def _tokens(command: str) -> list[str]:
    # Shell line continuations, as pasted from a README.
    command = command.replace("\\\r\n", " ").replace("\\\n", " ")
    tokens = shlex.split(_FLATTENED_CONTINUATION.sub(" ", command))
    # Drop leading "sudo docker run" / "docker container run".
    while tokens and tokens[0] in ("docker", "run", "container", "sudo", "podman"):
        tokens.pop(0)
    return tokens


def _parse(tokens: list[str]) -> tuple[list[tuple[str, str | None]], str | None, list[str]]:
    """Split into ([(long_flag, value)], image, command)."""
    flags: list[tuple[str, str | None]] = []
    i, n = 0, len(tokens)
    while i < n:
        tok = tokens[i]
        if tok == "--":
            i += 1
            break
        if not tok.startswith("-") or tok == "-":
            break  # the image
        if tok.startswith("--"):
            name, eq, value = tok.partition("=")
            if eq:
                flags.append((name, value))
                i += 1
            elif name in _BOOL_LONG:
                flags.append((name, None))
                i += 1
            else:
                flags.append((name, tokens[i + 1] if i + 1 < n else ""))
                i += 2
            continue
        # Short form: bundled booleans ("-dit"), or one value flag with its
        # value attached ("-p8080:80") or following ("-p 8080:80").
        body = tok[1:]
        j = 0
        while j < len(body) and body[j] in _BOOL_SHORT:
            flags.append((_SHORT_TO_LONG[body[j]], None))
            j += 1
        if j == len(body):
            i += 1
            continue
        letter = body[j]
        long_name = _VALUE_SHORT.get(letter, f"-{letter}")
        attached = body[j + 1 :]
        if attached.startswith("="):
            attached = attached[1:]
        if attached:
            flags.append((long_name, attached))
            i += 1
        else:
            flags.append((long_name, tokens[i + 1] if i + 1 < n else ""))
            i += 2
    image = tokens[i] if i < n else None
    command = tokens[i + 1 :] if image is not None else []
    return flags, image, command


def _is_named_volume(source: str) -> bool:
    return bool(source) and not source.startswith(("/", ".", "~", "$"))


def _mount(value: str) -> dict:
    """``--mount type=bind,source=/a,target=/b,readonly`` as compose long syntax."""
    spec: dict = {}
    for part in value.split(","):
        key, eq, val = part.partition("=")
        key = key.strip()
        if key in ("source", "src"):
            spec["source"] = val
        elif key in ("target", "destination", "dst"):
            spec["target"] = val
        elif key == "type":
            spec["type"] = val
        elif key in ("readonly", "ro"):
            spec["read_only"] = (not eq) or val.lower() in ("1", "true")
        elif key == "bind-propagation":
            spec.setdefault("bind", {})["propagation"] = val
        elif key == "volume-nocopy":
            spec.setdefault("volume", {})["nocopy"] = (not eq) or val.lower() in ("1", "true")
        elif key == "tmpfs-size":
            spec.setdefault("tmpfs", {})["size"] = val
    spec.setdefault("type", "volume")
    return spec


def _gpus(value: str) -> dict:
    """``--gpus all`` / ``--gpus 2`` / ``--gpus device=0,1`` as a device reservation."""
    reservation: dict = {"driver": "nvidia", "capabilities": ["gpu"]}
    value = value.strip().strip('"').strip("'")
    if value.startswith("device="):
        reservation["device_ids"] = [d for d in value[len("device="):].split(",") if d]
    elif value.isdigit():
        reservation["count"] = int(value)
    else:
        reservation["count"] = "all"
    return reservation


def _ulimit(value: str, ulimits: dict) -> None:
    name, _, limits = value.partition("=")
    soft, _, hard = limits.partition(":")
    try:
        ulimits[name] = (
            {"soft": int(soft), "hard": int(hard)} if hard else int(soft)
        )
    except ValueError:
        pass


def convert_docker_run(command: str) -> str:
    tokens = _tokens(command)
    flags, image, command_args = _parse(tokens)
    if not image:
        raise ValueError("Could not find an image in the docker run command")

    service: dict = {"image": image}
    top_volumes: dict = {}
    top_networks: dict = {}
    name = None
    lists: dict[str, list] = {}
    healthcheck: dict = {}
    logging: dict = {}
    ulimits: dict = {}
    gpus: list[dict] = []

    def add(key: str, value) -> None:
        lists.setdefault(key, []).append(value)

    for flag, value in flags:
        value = value if value is not None else ""
        if flag == "--name":
            name = value
        elif flag == "--publish":
            add("ports", value)
        elif flag == "--volume":
            add("volumes", value)
            source = value.split(":", 1)[0] if ":" in value else ""
            if _is_named_volume(source):
                top_volumes[source] = None
        elif flag == "--mount":
            spec = _mount(value)
            add("volumes", spec)
            if spec.get("type") == "volume" and _is_named_volume(spec.get("source", "")):
                top_volumes[spec["source"]] = None
        elif flag == "--tmpfs":
            add("tmpfs", value)
        elif flag == "--env":
            add("environment", value)
        elif flag == "--env-file":
            add("env_file", value)
        elif flag == "--device":
            add("devices", value)
        elif flag == "--cap-add":
            add("cap_add", value)
        elif flag == "--cap-drop":
            add("cap_drop", value)
        elif flag == "--label":
            add("labels", value)
        elif flag in ("--network", "--net"):
            if value in _NETWORK_MODES or value.startswith(("container:", "service:")):
                service["network_mode"] = value
            elif value:
                add("networks", value)
                # `docker run --network x` needs x to exist already.
                top_networks[value] = {"external": True}
        elif flag == "--network-alias":
            add("_aliases", value)
        elif flag == "--add-host":
            add("extra_hosts", value)
        elif flag == "--dns":
            add("dns", value)
        elif flag == "--dns-search":
            add("dns_search", value)
        elif flag == "--security-opt":
            add("security_opt", value)
        elif flag == "--sysctl":
            add("sysctls", value)
        elif flag == "--group-add":
            add("group_add", value)
        elif flag == "--expose":
            add("expose", value)
        elif flag == "--gpus":
            gpus.append(_gpus(value))
        elif flag == "--ulimit":
            _ulimit(value, ulimits)
        elif flag == "--restart":
            service["restart"] = value
        elif flag == "--privileged":
            service["privileged"] = True
        elif flag == "--read-only":
            service["read_only"] = True
        elif flag == "--init":
            service["init"] = True
        elif flag == "--interactive":
            service["stdin_open"] = True
        elif flag == "--tty":
            service["tty"] = True
        elif flag == "--hostname":
            service["hostname"] = value
        elif flag == "--user":
            service["user"] = value
        elif flag == "--workdir":
            service["working_dir"] = value
        elif flag == "--entrypoint":
            service["entrypoint"] = value
        elif flag == "--shm-size":
            service["shm_size"] = value
        elif flag == "--memory":
            service["mem_limit"] = value
        elif flag == "--memory-reservation":
            service["mem_reservation"] = value
        elif flag == "--cpus":
            try:
                service["cpus"] = float(value)
            except ValueError:
                pass
        elif flag == "--cpu-shares":
            try:
                service["cpu_shares"] = int(value)
            except ValueError:
                pass
        elif flag == "--pids-limit":
            try:
                service["pids_limit"] = int(value)
            except ValueError:
                pass
        elif flag == "--pid":
            service["pid"] = value
        elif flag == "--ipc":
            service["ipc"] = value
        elif flag == "--stop-signal":
            service["stop_signal"] = value
        elif flag in ("--stop-timeout", "--time"):
            service["stop_grace_period"] = f"{value}s" if value.isdigit() else value
        elif flag == "--platform":
            service["platform"] = value
        elif flag == "--pull":
            service["pull_policy"] = value
        elif flag == "--log-driver":
            logging["driver"] = value
        elif flag == "--log-opt":
            key, _, val = value.partition("=")
            logging.setdefault("options", {})[key] = val
        elif flag == "--health-cmd":
            healthcheck["test"] = ["CMD-SHELL", value]
        elif flag == "--health-interval":
            healthcheck["interval"] = value
        elif flag == "--health-timeout":
            healthcheck["timeout"] = value
        elif flag == "--health-start-period":
            healthcheck["start_period"] = value
        elif flag == "--health-retries":
            try:
                healthcheck["retries"] = int(value)
            except ValueError:
                pass
        elif flag == "--no-healthcheck":
            healthcheck = {"disable": True}
        # Everything else (--detach, --rm, --publish-all, --cidfile, …) has no
        # compose equivalent worth writing down.

    aliases = lists.pop("_aliases", [])
    if aliases:
        nets = lists.get("networks") or []
        if nets:
            # Aliases apply to the networks given on the command line.
            lists["networks"] = {net: {"aliases": aliases} for net in nets}
    service.setdefault("restart", "unless-stopped")
    for key in (
        "ports", "expose", "volumes", "tmpfs", "environment", "env_file", "devices",
        "cap_add", "cap_drop", "security_opt", "sysctls", "group_add", "extra_hosts",
        "dns", "dns_search", "labels", "networks",
    ):
        if key in lists:
            service[key] = lists[key]
    if ulimits:
        service["ulimits"] = ulimits
    if logging:
        service["logging"] = logging
    if healthcheck:
        service["healthcheck"] = healthcheck
    if gpus:
        service["deploy"] = {"resources": {"reservations": {"devices": gpus}}}
    if command_args:
        service["command"] = command_args
    if name:
        service["container_name"] = name

    svc_name = _SERVICE_NAME_RE.sub("-", name).strip("-.") if name else ""
    doc: dict = {"services": {svc_name or _service_name_from_image(image): service}}
    if top_volumes:
        doc["volumes"] = top_volumes
    if top_networks:
        doc["networks"] = top_networks
    return yaml.safe_dump(doc, sort_keys=False, default_flow_style=False)