#!/usr/bin/env python3
import json
import os
import pathlib
import shlex
import subprocess
import sys
import time

PACKAGE_REF = "modulus/latest"
CONAN_VERSION = "2.26.2"
FETCH_IMAGE = "modulus-fetch:dev"
DEFAULT_BUILD_BASE_IMAGE = "debian:trixie-slim"
PICO_SDK_VERSION = "2.3.0"
COMMANDS = (
  "configure",
  "build",
  "test",
  "package",
  "lsp",
  "run",
)
COMMON_DIR = pathlib.Path("build") / "build_common"
TOOLS_DIR = pathlib.Path("build") / "build_tools"
RUNTIME_DIR = COMMON_DIR / "modulus"
CONAN_HOME_DIR = COMMON_DIR / "conan_cache"
DEFAULT_REMOTE_NAME = "khodo"
DEFAULT_REMOTE_URL = "https://git.thyme-labs.com/api/packages/khodo/conan"

def repo_root():
  root = pathlib.Path(__file__).resolve().parent
  if not (root / "modulus.json").is_file():
    raise SystemExit("modulus.py must be run from a Modulus-managed repository root")
  return root

def read_json(path):
  try:
    data = json.loads(path.read_text(encoding="utf-8"))
  except (OSError, json.JSONDecodeError) as error:
    raise SystemExit(f"failed to read {path}: {error}") from error
  if not isinstance(data, dict):
    raise SystemExit(f"{path} must contain a JSON object")
  return data

def remote_config(root):
  config = read_json(root / "modulus.json")
  remote = config.get(
    "tooling",
    {},
  ).get("remote", {})
  if not isinstance(remote, dict):
    remote = {}
  return (
    str(remote.get("name") or DEFAULT_REMOTE_NAME),
    str(remote.get("url") or DEFAULT_REMOTE_URL),
  )

def runtime_ready(root):
  return (
    (root / TOOLS_DIR / "configure.py").is_file()
    and (root / TOOLS_DIR / "lsp.py").is_file()
    and (root / TOOLS_DIR / "run.py").is_file()
    and (root / RUNTIME_DIR / "modulus" / "cli.py").is_file()
    and (root / RUNTIME_DIR / "modulus" / "export.py").is_file()
  )

def optional_credentials():
  user = os.environ.get("CONAN_USER")
  token = os.environ.get("CONAN_TOKEN")
  if bool(user) != bool(token):
    raise SystemExit("CONAN_USER and CONAN_TOKEN must either both be set or both be unset")
  return user, token

def env_flag(name):
  return os.environ.get(name, "").lower() in {"1", "true", "yes", "on"}

def direct_mode(root):
  return env_flag("MODULUS_DIRECT") or project_docker_marker(root).exists()

def run(command, input_text=None, capture=False, display_command=None, env=None, check=True):
  visible = display_command if display_command is not None else command
  print("+", " ".join(shlex.quote(str(part)) for part in visible))
  completed = subprocess.run(
    [str(part) for part in command],
    input=input_text,
    stdout=subprocess.PIPE if capture else None,
    stderr=subprocess.PIPE if capture else None,
    text=True,
    check=False,
    env=env,
  )
  if check and completed.returncode != 0:
    if capture and completed.stdout:
      print(completed.stdout, end="")
    if capture and completed.stderr:
      print(completed.stderr, end="", file=sys.stderr)
    raise SystemExit(completed.returncode)
  return completed

def run_with_retries(command, attempts=3, input_text=None, display_command=None, env=None):
  completed = None
  for attempt in range(1, attempts + 1):
    completed = run(
      command,
      input_text=input_text,
      display_command=display_command,
      env=env,
      check=False,
    )
    if completed.returncode == 0:
      return completed
    if attempt != attempts:
      delay = attempt * 20
      print(f"command failed with exit code {completed.returncode}; retrying in {delay}s")
      time.sleep(delay)
  raise SystemExit(completed.returncode if completed is not None else 1)

def fetch_dockerfile():
  return f"""FROM {DEFAULT_BUILD_BASE_IMAGE}

ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update \\
  && apt-get install -y --no-install-recommends \\
    ca-certificates \\
    python3 \\
    python3-venv \\
  && rm -rf /var/lib/apt/lists/*

RUN python3 -m venv /opt/modulus-python

ENV PATH=/opt/modulus-python/bin:$PATH

RUN python -m pip install --no-cache-dir conan=={CONAN_VERSION}

WORKDIR /workspace
"""

def product_build_dockerfile():
  return f"""FROM {DEFAULT_BUILD_BASE_IMAGE}

ENV DEBIAN_FRONTEND=noninteractive

RUN dpkg --add-architecture arm64

RUN apt-get update \\
  && apt-get install -y --no-install-recommends \\
    bash \\
    build-essential \\
    ca-certificates \\
    clangd-19 \\
    cmake \\
    curl \\
    docker-buildx \\
    docker-cli \\
    docker.io \\
    g++-12 \\
    g++-12-aarch64-linux-gnu \\
    gcc-12 \\
    gcc-12-aarch64-linux-gnu \\
    gcc-arm-none-eabi \\
    git \\
    libdrm-dev \\
    libdrm-dev:arm64 \\
    libegl-dev \\
    libegl-dev:arm64 \\
    libgbm-dev \\
    libgbm-dev:arm64 \\
    libgles-dev \\
    libgles-dev:arm64 \\
    libglvnd-dev \\
    libglvnd-dev:arm64 \\
    libnewlib-arm-none-eabi \\
    libstdc++-arm-none-eabi-newlib \\
    ninja-build \\
    pkg-config \\
    python3 \\
    python3-venv \\
  && ln -sf /usr/bin/clangd-19 /usr/local/bin/clangd \\
  && rm -rf /var/lib/apt/lists/*

ENV CC=/usr/bin/gcc-12
ENV CXX=/usr/bin/g++-12
ENV VAMP_PICO_SDK_VERSION={PICO_SDK_VERSION}
ENV PICO_SDK_PATH=/opt/pico-sdk

RUN git clone \\
    --depth 1 \\
    --branch "$VAMP_PICO_SDK_VERSION" \\
    --recurse-submodules \\
    --shallow-submodules \\
    https://github.com/raspberrypi/pico-sdk.git \\
    "$PICO_SDK_PATH"

RUN python3 -m venv /opt/modulus-python

ENV PATH=/opt/modulus-python/bin:$PATH

RUN python -m pip install --no-cache-dir \\
    conan=={CONAN_VERSION}

WORKDIR /workspace

RUN mkdir -p /opt && touch /opt/modulus-build-image
"""

def sync_product_build_dockerfile(root):
  dockerfile = root / TOOLS_DIR / "docker" / "Dockerfile"
  if not dockerfile.parent.is_dir():
    return
  desired = product_build_dockerfile()
  try:
    current = dockerfile.read_text(encoding="utf-8")
  except OSError:
    current = None
  if current != desired:
    dockerfile.write_text(desired, encoding="utf-8")

def ensure_fetch_image():
  if not docker_platform_args():
    inspected = subprocess.run(
      [
        "docker",
        "image",
        "inspect",
        FETCH_IMAGE,
      ],
      stdout=subprocess.DEVNULL,
      stderr=subprocess.DEVNULL,
      check=False,
    )
    if inspected.returncode == 0:
      return
  run_with_retries(
    [
      "docker",
      "build",
      *docker_platform_args(),
      "--tag",
      FETCH_IMAGE,
      "-",
    ],
    input_text=fetch_dockerfile(),
  )

def in_container():
  return pathlib.Path("/.dockerenv").exists()

def docker_workspace_args(root):
  if in_container():
    container = os.environ.get("HOSTNAME", "").strip()
    if not container:
      raise SystemExit("HOSTNAME is required to share the CI workspace with the Modulus fetch container")
    return [
      "--volumes-from",
      container,
    ]
  return [
    "--volume",
    f"{root}:{root}",
  ]

def docker_socket_args():
  socket = pathlib.Path("/var/run/docker.sock")
  if not socket.is_socket():
    return []
  args = [
    "--volume",
    "/var/run/docker.sock:/var/run/docker.sock",
  ]
  try:
    args.extend(["--group-add", str(socket.stat().st_gid)])
  except OSError:
    pass
  return args

def docker_platform_args():
  platform = os.environ.get("MODULUS_DOCKER_PLATFORM", "").strip()
  if not platform:
    return []
  return [
    "--platform",
    platform,
  ]

def fetch_env(root):
  env = os.environ.copy()
  env["CONAN_HOME"] = str(root / CONAN_HOME_DIR)
  return env

def run_fetch(root, command, capture=False, display_command=None, check=True):
  if direct_mode(root):
    return run(command, capture=capture, display_command=display_command, env=fetch_env(root), check=check)
  ensure_fetch_image()
  docker_command = [
    "docker",
    "run",
    "--rm",
    "--interactive",
    "--user",
    f"{os.getuid()}:{os.getgid()}",
    "--env",
    "HOME=/tmp",
    "--env",
    "CONAN_HOME",
    "--env",
    "CONAN_USER",
    "--env",
    "CONAN_TOKEN",
    "--env",
    "DOCKER_HOST",
    *docker_workspace_args(root),
    *docker_socket_args(),
    "--workdir",
    str(root),
    FETCH_IMAGE,
    *command,
  ]
  return run(docker_command, capture=capture, display_command=display_command, env=fetch_env(root), check=check)

def configure_remote(root, remote_name, remote_url):
  user, token = optional_credentials()
  run_fetch(
    root,
    [
      "conan",
      "remote",
      "add",
      "-f",
      remote_name,
      remote_url,
    ],
  )
  if not user and not token:
    return
  run_fetch(
    root,
    [
      "conan",
      "remote",
      "login",
      remote_name,
      user,
      "-p",
      token,
    ],
    display_command=[
      "conan",
      "remote",
      "login",
      remote_name,
      user,
      "-p",
      "<hidden>",
    ],
  )

def package_refs(data):
  refs = []
  for location_data in data.values():
    if not isinstance(location_data, dict):
      continue
    for ref, ref_data in location_data.items():
      if not ref.startswith(PACKAGE_REF) or not isinstance(ref_data, dict):
        continue
      revisions = ref_data.get(
        "revisions",
        {},
      )
      if not isinstance(revisions, dict):
        continue
      for recipe_revision, revision_data in revisions.items():
        if not isinstance(revision_data, dict):
          continue
        packages = revision_data.get(
          "packages",
          {},
        )
        if not isinstance(packages, dict):
          continue
        for package_id, package_data in packages.items():
          package_revisions = package_data.get(
            "revisions",
            {},
          )if isinstance(package_data, dict) else {}
          if isinstance(package_revisions, dict) and package_revisions:
            for package_revision in package_revisions:
              refs.append(f"{PACKAGE_REF}#{recipe_revision}:{package_id}#{package_revision}")
          else:
            refs.append(f"{PACKAGE_REF}#{recipe_revision}:{package_id}")
  return sorted(refs)

def remove_cached_modulus(root):
  run_fetch(
    root,
    [
      "conan",
      "remove",
      PACKAGE_REF + ":*",
      "-c",
    ],
    check=False,
  )

def download_modulus(root, remote_name):
  remove_cached_modulus(root)
  run_fetch(
    root,
    [
      "conan",
      "download",
      PACKAGE_REF + ":*",
      "-r",
      remote_name,
    ],
  )
  completed = run_fetch(
    root,
    [
      "conan",
      "list",
      PACKAGE_REF + ":*",
      "-c",
      "-f",
      "json",
    ],
    capture=True,
  )
  refs = package_refs(json.loads(completed.stdout))
  if not refs:
    raise SystemExit(f"no downloaded {PACKAGE_REF} package found in the Conan cache")
  package_ref = refs[-1]
  completed = run_fetch(
    root,
    [
      "conan",
      "cache",
      "path",
      package_ref,
    ],
    capture=True,
  )
  package_root = pathlib.Path(completed.stdout.strip())
  if not (package_root / "modulus" / "export.py").is_file() or not (package_root / "tools").is_dir():
    raise SystemExit(f"downloaded {PACKAGE_REF} package does not contain Modulus export files: {package_root}")
  return package_root

def materialize_runtime(root):
  remote_name, remote_url = remote_config(root)
  configure_remote(root, remote_name, remote_url)
  package_root = download_modulus(root, remote_name)
  sys.path.insert(0, str(package_root))
  import modulus.export
  modulus.export.sync_project(root, package_root)

def ensure_runtime(root):
  if not runtime_ready(root):
    materialize_runtime(root)
  sync_product_build_dockerfile(root)

def project_docker_marker(root):
  config = read_json(root / "modulus.json")
  tooling = config.get("tooling", {})
  if not isinstance(tooling, dict):
    tooling = {}
  docker = tooling.get("docker", {})
  if not isinstance(docker, dict):
    docker = {}
  return pathlib.Path(str(docker.get("marker") or "/opt/modulus-build-image"))

def run_configure(argv):
  return delegate("configure", argv)

def run_package(argv):
  return delegate("package", argv)

def delegate(command, argv):
  root = repo_root()
  ensure_runtime(root)
  sys.path.insert(0, str(root / RUNTIME_DIR))
  import modulus.cli
  original_docker_run = modulus.cli.docker_run

  def docker_run_with_host_socket(context, command, docker_args=None, env_overrides=None):
    merged_docker_args = [
      *docker_socket_args(),
      *(docker_args or []),
    ]
    return original_docker_run(context, command, docker_args=merged_docker_args, env_overrides=env_overrides)

  modulus.cli.docker_run = docker_run_with_host_socket
  previous_argv = sys.argv
  previous_cwd = pathlib.Path.cwd()
  sys.argv = [
    str(root / "modulus.py"),
    *argv,
  ]
  os.chdir(root)
  try:
    return modulus.cli.main(command)
  finally:
    os.chdir(previous_cwd)
    sys.argv = previous_argv

def main(argv):
  if not argv or argv[0] in ("-h", "--help"):
    print("usage: python3 modulus.py " + "|".join(COMMANDS) + " ...")
    return 0 if argv else 2
  command = argv[0]
  if command not in COMMANDS:
    raise SystemExit(f"unknown Modulus command: {command}")
  if command == "configure":
    return run_configure(argv[1:])
  if command == "package":
    return run_package(argv[1:])
  return delegate(command, argv[1:])

if __name__ == "__main__":
  raise SystemExit(main(sys.argv[1:]))