#!/usr/bin/env python3
"""omniopt_docker - Build and run OmniOpt inside a Docker container.

Python rewrite of the former ``omniopt_docker`` bash script.

Usage:
    omniopt_docker [inner command]

If ``inner command`` is empty, the script only builds and starts the
container.  Otherwise it runs the inner command inside a fresh container
that has the project's ``runs/``, ``logs/`` and matplotlib config mounted.

The script is split into small, pure functions so the behaviour can be
exercised without a real docker daemon (see
``/.tests/test_omniopt_docker``).

Behavioural compatibility with the bash version:
  * ``--help`` exits 0 and prints help
  * Inner command must start with one of: ``./omniopt``, ``omniopt``,
    ``./.tests/``, ``.tests/``, ``python3 `` (anything else -> exit 1)
  * ``python3 foo`` is mapped to ``python3 /var/opt/omniopt/foo``;
    everything else is mapped to ``bash /var/opt/omniopt/<inner>``
  * Volumes: runs/, logs/, current pwd, matplotlib config
  * If ``$DISPLAY`` is set: adds X11 mounts and ``--user=$(id -u)``
  * Uses ``sudo docker`` when the user is not in the ``docker`` group
"""

from __future__ import annotations

import argparse
import os
import shutil
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, List, Optional, Sequence


DOCKER_IMAGE_NAME = "omniopt-omniopt2"

INNER_PREFIXES = ("./omniopt", "omniopt", "./.tests/", ".tests/", "python3 ")


# ---------------------------------------------------------------------------
# Pure helpers (used by tests)
# ---------------------------------------------------------------------------


def determine_docker_cmd(*, in_docker_group: bool) -> tuple[str, str]:
    """Return ``(compose_cmd, run_cmd)`` based on group membership."""
    if in_docker_group:
        return "docker compose", "docker"
    return "sudo docker compose", "sudo docker"


def validate_inner_command(inner: str) -> None:
    """Reject dangerous-looking inner commands.

    The original bash checked the WHOLE remaining argv string against
    patterns like ``python3*`` -- so ``python3 .tests/main --foo`` was
    accepted because it starts with ``python3``.  This function takes a
    single string for the same reason; :func:`main` joins ``argv``'s
    tail back together before calling here.
    """
    if not any(inner.startswith(prefix) for prefix in INNER_PREFIXES):
        raise ValueError(
            f"Invalid inner command {inner!r}: must start with one of {INNER_PREFIXES}"
        )


def _strip_one_dot_slash(path: str) -> str:
    """Strip at most one leading ``./`` (``./.tests/main`` -> ``.tests/main``)."""
    if path.startswith("./"):
        return path[2:]
    return path


def _looks_like_python_script(inner: str) -> bool:
    """True if `inner` points at a file that has a Python shebang.

    Used so callers can write ``./.tests/main`` instead of the longer
    ``python3 .tests/main`` — the latter still wins when given.
    """
    if not inner or " " in inner or inner.startswith("-"):
        return False
    # Strip at most one leading `./` so `./.tests/main` -> `.tests/main`
    # but `..foo` and `.foo` are left alone.
    candidate = _strip_one_dot_slash(inner)
    candidates = [candidate, "./" + candidate, inner]
    for path in candidates:
        try:
            with open(path, "rb") as fh:
                first_line = fh.readline(256).decode("utf-8", errors="ignore")
        except OSError:
            continue
        return first_line.startswith("#!") and "python" in first_line
    return False


def _parse_inner_command(
    tokens: Sequence[str],
) -> tuple[str, str, List[str]]:
    """Map the inner command tokens to ``(interpreter, script, args)``.

    ``tokens`` are the raw argv entries (e.g. ``["python3", ".tests/main",
    "--quick"]``), so argument boundaries -- including spaces inside a
    single argument -- are preserved.
    """
    if tokens and tokens[0] == "python3":
        interpreter = "python3"
        script = tokens[1] if len(tokens) > 1 else ""
        script_args = list(tokens[2:])
    elif tokens and _looks_like_python_script(tokens[0]):
        # Heuristic: anything that starts with `./` and points at a file
        # with a `#!/usr/bin/env python3` (or `python`) shebang is
        # dispatched as python3 instead of bash.  This avoids the
        # `bash: ./.tests/main: line 13: ...: command not found` errors
        # we used to get when callers passed the old bash-typed names.
        interpreter = "python3"
        script = _strip_one_dot_slash(tokens[0])
        script_args = list(tokens[1:])
    else:
        interpreter = "bash"
        script = tokens[0] if tokens else ""
        script_args = list(tokens[1:])

    return interpreter, script, script_args


def build_run_command(
    *,
    inner: str,
    docker_name: str,
    docker_cmd: str,
    pwd: str,
    home: str,
    has_display: bool,
    uid: Optional[int] = None,
    inner_argv: Optional[Sequence[str]] = None,
) -> List[str]:
    """Build the ``docker run`` argv.

    ``inner`` is the joined inner command string (validated against the
    allowed prefixes).  ``inner_argv``, when given, carries the original
    argv entries so argument boundaries survive; otherwise ``inner`` is
    whitespace-split.

    The bash version translated ``./foo`` -> ``/var/opt/omniopt/foo`` and
    ``python3 foo`` -> ``python3 /var/opt/omniopt/foo``.  We replicate
    that here.
    """
    validate_inner_command(inner)

    interpreter, script, script_args = _parse_inner_command(
        inner_argv if inner_argv is not None else inner.split()
    )
    target = "/var/opt/omniopt/" + script

    if uid is None:
        uid = os.getuid() if hasattr(os, "getuid") else 0

    cmd: List[str] = [
        *docker_cmd.split(),
        "run",
        "-v", f"{pwd}/logs:/var/opt/omniopt/logs:rw",
        "-v", f"{pwd}/runs:/var/opt/omniopt/runs:rw",
        "-v", f"{pwd}/:/var/opt/omniopt/docker_user_dir:rw",
        "-v", f"{home}/.config/matplotlib_docker_omniopt:{home}/.config/matplotlib:rw",
        "--mount", "type=tmpfs,destination=/tmp",
        "-t", "--rm", docker_name,
        interpreter, target, *script_args,
    ]

    if has_display:
        extra = [
            f"--user={uid}",
            "--env=DISPLAY",
            "--volume=/etc/group:/etc/group:ro",
            "--volume=/etc/passwd:/etc/passwd:ro",
            "--volume=/etc/shadow:/etc/shadow:ro",
            "--volume=/etc/sudoers.d:/etc/sudoers.d:ro",
            "--volume=/tmp/.X11-unix:/tmp/.X11-unix:rw",
        ]
        idx = cmd.index("-t")
        cmd = cmd[:idx] + extra + cmd[idx:]

    return cmd


# ---------------------------------------------------------------------------
# I/O helpers (mockable via the injected dependencies in main())
# ---------------------------------------------------------------------------


def _default_check_cmd(cmd: str) -> bool:
    return shutil.which(cmd) is not None


def _default_mkdir(paths: Sequence[str]) -> None:
    for p in paths:
        Path(p).mkdir(parents=True, exist_ok=True)


def _default_docker_build(compose_cmd: str) -> int:
    args = compose_cmd.split() + [
        "build",
        "--build-arg",
        f"GetMyUsername={os.environ.get('USER', 'root')}",
    ]
    return subprocess.run(args).returncode


def _default_docker_up(compose_cmd: str) -> int:
    args = compose_cmd.split() + ["up", "-d"]
    return subprocess.run(args).returncode


def _default_docker_run(cmd: List[str]) -> int:
    return subprocess.run(cmd).returncode


def _default_in_docker_group() -> bool:
    try:
        groups = subprocess.run(
            ["groups"], check=False, capture_output=True, text=True
        ).stdout
    except OSError:
        return False
    return "docker" in groups.split()


def _default_has_display() -> bool:
    return bool(os.environ.get("DISPLAY"))


HELP_TEXT = """\
Usage: omniopt_docker [OPTIONS] [INNER COMMAND]

Options:
  --help     Show this help and exit.

If INNER COMMAND is omitted, the docker image is built and the container
started.  Otherwise the inner command is executed inside a fresh
container that mounts runs/, logs/ and the matplotlib config directory.

Allowed INNER COMMAND prefixes:
  ./omniopt ...
  omniopt ...
  ./.tests/...
  .tests/...
  python3 ...
"""


def _build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="omniopt_docker",
        description="Build and run OmniOpt inside a Docker container.",
        add_help=False,
    )
    p.add_argument("--help", action="store_true", help="Show this help and exit.")
    return p


def main(
    argv: Optional[Sequence[str]] = None,
    *,
    check_cmd: Callable[[str], bool] = _default_check_cmd,
    mkdir: Callable[[Sequence[str]], None] = _default_mkdir,
    docker_build: Callable[[str], int] = _default_docker_build,
    docker_up: Callable[[str], int] = _default_docker_up,
    docker_run: Callable[[List[str]], int] = _default_docker_run,
    docker_cmd: Optional[str] = None,
    in_docker_group: Optional[bool] = None,
    has_display: Optional[bool] = None,
    pwd: Optional[str] = None,
    home: Optional[str] = None,
) -> int:
    """Entry point with all side-effecting helpers injectable for tests."""
    parser = _build_parser()
    args, inner = parser.parse_known_args(list(argv) if argv is not None else sys.argv[1:])
    if args.help:
        print(HELP_TEXT)
        return 0

    if has_display is None:
        has_display = _default_has_display()
    if in_docker_group is None:
        in_docker_group = _default_in_docker_group()
    if docker_cmd is None:
        compose_cmd, run_cmd = determine_docker_cmd(in_docker_group=in_docker_group)
    else:
        compose_cmd, run_cmd = ("docker compose", docker_cmd)
    if pwd is None:
        pwd = os.getcwd()
    if home is None:
        home = os.path.expanduser("~")

    if inner:
        try:
            # Join back into a single string so ``python3 .tests/main``
            # matches the ``python3 `` prefix the same way it did in
            # the original bash script.
            validate_inner_command(" ".join(inner))
        except ValueError as e:
            print(f"Error: {e}", file=sys.stderr)
            return 1

    mkdir([f"{pwd}/runs", f"{pwd}/logs"])

    if docker_build(compose_cmd) != 0:
        print("Failed composing container", file=sys.stderr)
        return 1

    if docker_up(compose_cmd) != 0:
        print("Failed to build container", file=sys.stderr)
        return 1

    if not inner:
        return 0

    mkdir([f"{home}/.config/matplotlib_docker_omniopt"])
    cmd = build_run_command(
        inner=" ".join(inner),
        inner_argv=list(inner),
        docker_name=DOCKER_IMAGE_NAME,
        docker_cmd=run_cmd,
        pwd=pwd,
        home=home,
        has_display=has_display,
    )
    rc = docker_run(cmd)
    if rc != 0:
        print("Command 2 failed. Docker images:", file=sys.stderr)
        subprocess.run(run_cmd.split() + ["images"], check=False)
        return rc
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        sys.exit(0)
