#!/usr/bin/env ruby

$LOAD_PATH.unshift File.expand_path("../lib", __dir__)

require "optparse"
require "open3"
require "rb_sys/version"
require "rb_sys/toolchain_info"
require "rb_sys/cargo/metadata"
require "rb_sys/util/logger"
require "fileutils"
require "tmpdir"

OPTIONS = {
  docker_platform: "linux/amd64",
  version: RbSys::VERSION,
  directory: Dir.pwd
}

def cargo_metadata
  @cargo_metadata ||= RbSys::Cargo::Metadata.new("rb-sys-dock", deps: true)
end

def logger
  return @logger if @logger

  io = ARGV.include?("--quiet") ? File.open(File::NULL, "w") : $stderr
  @logger ||= RbSys::Util::Logger.new(io: io)
end

def list_platforms
  RbSys::ToolchainInfo.supported.each do |p|
    old = logger.io
    logger.io = $stdout
    puts "- #{p.platform}"
  ensure
    logger.io = old
  end
end

OptionParser.new do |opts|
  opts.banner = <<~MSG
    Usage: rb-sys-dock [OPTIONS] [COMMAND]

    A CLI to facillitate building Ruby on Rust extensions using Docker.

    Examples:

        Build for Linux (x86_64)
            $ rb-sys-dock --platform x86_64-linux --build

        Build for macOS (Ruby 3.1 and 3.2)
            $ rb-sys-dock -p arm64-darwin --ruby-versions 3.1,3.2 --build

        Enter a shell
            $ rb-sys-dock --list-platforms

        Run a command in the container with a specific directory mounted
            $ rb-sys-dock -p x64-mingw-ucrt -d /tmp/mygem -- "env | grep RUBY"

        List all supported platforms
            $ rb-sys-dock --list-platforms

    Options:
  MSG

  opts.on("--quiet", "Prints no logging output") do
  end

  opts.on("-v", "--version", "Prints version") do
    require "rb_sys/version"
    puts RbSys::VERSION
    exit
  end

  opts.on("--build", "Run the default command to cross-compile a gem") do
    OPTIONS[:build] = true
  end

  opts.on("-p", "--platform PLATFORM", "Platform to build for (i.e. x86_64-linux)") do |p|
    toolchain_info = begin
      RbSys::ToolchainInfo.new(p)
    rescue
      supported_list = RbSys::ToolchainInfo.all
      supported_list.select!(&:supported?)
      list = supported_list.map { |p| "- #{p} (#{p.rust_target})" }.join("\n")
      logger.error("Platform #{p} is not supported, please use one of:\n\n#{list}")
      exit(1)
    end

    OPTIONS[:platform] = p
    OPTIONS[:toolchain_info] = toolchain_info
  end

  opts.on("-r", "--ruby-versions LIST", "List all supported Ruby versions") do |arg|
    vers = arg.split(/[^0-9.]+/).map do |v|
      parts = v.split(".")
      parts[2] = "0" if parts[2].nil?
      parts.join(".")
    end

    OPTIONS[:ruby_versions] = vers

    logger.info("Building for Ruby requested versions: #{vers}")

    ENV["RUBY_CC_VERSION"] = vers.join(":")
  end

  opts.on("--tag TAG", "Use a specific version of the Docker image") do |tag|
    logger.info("Using version #{tag} of the Docker image")
    OPTIONS[:version] = tag
    OPTIONS[:no_cache] = tag == "latest"
  end

  opts.on("--list-platforms", "--list", "List all supported platforms") do
    logger.info("Supported platforms listed below:")
    list_platforms
    exit(0)
  end

  opts.on("-h", "--help", "Prints this help") do
    puts opts
    exit
  end

  opts.on("-V", "--verbose", "Prints verbose output") do
    ENV["LOG_LEVEL"] = "trace"
    ENV["VERBOSE"] = "1"
    logger.level = :trace
    OPTIONS[:verbose] = true
  end

  opts.on("--mount-toolchains", "Mount local Rustup toolchain (instead of pre-installed from Docker container)") do
    OPTIONS[:mount_rustup_toolchains] = true
  end

  opts.on("-d", "--directory DIR", "Directory to run the command in") do |val|
    OPTIONS[:directory] = File.expand_path(val)
  end
end.parse!

def default_docker_command
  return @default_docker_command if defined?(@default_docker_command)

  @default_docker_command = ENV.fetch("DOCKER") do
    if !(docker = `which docker`).empty?
      docker.strip
    elsif !(podman = `which podman`).empty?
      podman.strip
    else
      logger.fatal("Could not find docker or podman command, please install one of them")
    end
  end
end

def run_command!(*cmd)
  logger.trace("Running command:\n\t$ #{cmd.join(" ")}")
  stdout, stderr, status = Open3.capture3(*cmd)

  if status.success?
    stdout
  else
    logger.error("Error running command: $ #{cmd}")
    warn(stderr)
    exit(status.exitstatus)
  end
end

def docker(cmd)
  run_command!("#{default_docker_command} #{cmd}")
end

def determine_cache_dir
  return ENV["RB_SYS_DOCK_CACHE_DIR"] if ENV["RB_SYS_DOCK_CACHE_DIR"]
  return File.join(ENV["XDG_CACHE_HOME"], "rb-sys-dock") if ENV["XDG_CACHE_HOME"]

  File.join(ENV["HOME"], ".cache", "rb-sys-dock")
end

def docker_tmp
  @docker_tmp ||= "/tmp/rb-sys-dock"
end

def docker_bundle_home
  @docker_bundle_home ||= "/tmp/rb-sys-dock/bundle-home"
end

def cache_dir
  return @cache_dir if defined?(@cache_dir)

  @cache_dir = determine_cache_dir
  FileUtils.mkdir_p(@cache_dir)
  @cache_dir
end

def mount_cargo_registry
  local_registry_dir = if ENV["CARGO_HOME"]
    ENV["CARGO_HOME"]
  elsif File.exist?(cargo_home = File.join(ENV["HOME"], ".cargo"))
    cargo_home
  else
    File.join(cache_dir, "cargo")
  end

  dir = File.join("registry")
  logger.trace("Mounting cargo registry dir: #{dir}")
  FileUtils.mkdir_p(dir)

  mount_shared_bind_dir(File.join(local_registry_dir, dir), File.join("/usr/local/cargo", dir))
end

def mount_rustup_toolchains
  return unless OPTIONS[:mount_rustup_toolchains]

  local_rustup_dir = if OPTIONS[:mount_rustup_toolchains].is_a?(String)
    OPTIONS[:mount_rustup_toolchains]
  elsif ENV["RUSTUP_HOME"]
    ENV["RUSTUP_HOME"]
  elsif File.exist?(rustup_home = File.join(ENV["HOME"], ".rustup"))
    rustup_home
  else
    logger.fatal("Could not find Rustup home directory, please set RUSTUP_HOME")
  end

  logger.info("Mounting rustup toolchains from #{local_rustup_dir}")

  target_triple = OPTIONS[:toolchain_info].rust_target
  dkr_triple = "x86_64-unknown-linux-gnu"
  dkr_toolchain = "stable-#{dkr_triple}"
  dkr_toolchain_dir = "/usr/local/rustup/toolchains/#{dkr_toolchain}"
  installed_toolchains = Dir.glob(File.join(local_rustup_dir, "toolchains", "*")).map { |f| File.basename(f) }
  has_host_toolchain = installed_toolchains.any? { |t| t.end_with?(dkr_toolchain) }

  if !has_host_toolchain
    logger.info("Installing default toolchain for docker image (#{dkr_toolchain})")
    run_command!("rustup", "toolchain", "add", dkr_toolchain, "--force-non-host")
  end

  has_target = run_command!("rustup target list --installed --toolchain #{dkr_toolchain}").include?(target_triple)

  if !has_target
    logger.info("Installing target for docker image (#{target_triple})")
    run_command!("rustup", "target", "add", target_triple, "--toolchain", dkr_toolchain)
  end

  volume("#{local_rustup_dir}/toolchains/#{dkr_toolchain}", dkr_toolchain_dir, mode: "z,ro")
end

def volume(src, dest, mode: "rw")
  "--volume #{src}:#{dest}:rw"
end

def mount_bundle_cache
  dir = File.join(cache_dir, ruby_platform, "bundle")
  bundle_path = File.join(docker_tmp, "bundle")
  FileUtils.mkdir_p(dir)
  logger.trace("Mounting bundle cache: #{dir}")

  "#{volume(dir, bundle_path)} -e BUNDLE_PATH=#{bundle_path.inspect}"
end

def mount_bundle_config
  bundle_config = File.join(ENV["HOME"] || "~", ".bundle", "config")

  return unless File.exist?(bundle_config)

  logger.trace("Mounting bundle config: #{bundle_config}")
  parts = []
  parts << volume(bundle_config, File.join(docker_bundle_home, ".bundle", "config"), mode: "ro")
  parts << "-e BUNDLE_HOME=#{docker_bundle_home.inspect}"
  parts.join(" ")
end

def tmp_target_dir
  return @tmp_target_dir if defined?(@tmp_target_dir)

  dir = File.join(working_directory, "tmp", "rb-sys-dock", ruby_platform, "target")
  FileUtils.mkdir_p(dir)
  @tmp_target_dir = dir
end

def working_directory
  OPTIONS.fetch(:directory)
end

def mount_target_dir
  "-v #{tmp_target_dir}:#{File.join(working_directory, "target")}"
end

def mount_command_history
  return unless $stdin.tty?

  history_dir = File.join(cache_dir, OPTIONS.fetch(:platform), "commandhistory")
  FileUtils.mkdir_p(history_dir)
  "-v #{history_dir}:#{File.join(docker_tmp, "commandhistory")}"
end

def default_command_to_run(input_args)
  input_cmd = input_args.empty? ? "true" : input_args.join(" ")

  if OPTIONS[:build]
    with_bundle = +"test -f Gemfile && bundle install && #{input_cmd} && bundle exec rake native:$RUBY_TARGET gem"
    without_bundle = "#{input_cmd} && rake native:$RUBY_TARGET gem"
    logger.info("Running default build command (rake native:#{ruby_platform} gem)")
    "bash -c '(#{with_bundle}) || (#{without_bundle})'"
  else
    input_args.empty? ? "bash" : "bash -c '#{input_args.join(" ")}'"
  end
end

def uid_gid
  explicit_uid = ENV["RB_SYS_DOCK_UID"]
  explicit_gid = ENV["RB_SYS_DOCK_GID"]

  if /darwin/.match?(RbConfig::CONFIG["host_os"]) && !default_docker_command.include?("podman")
    [explicit_uid || "1000", explicit_gid || "1000"]
  else
    [explicit_uid || Process.uid, explicit_gid || Process.gid]
  end
end

def user_mapping
  uid, gid = uid_gid
  "-e UID=#{uid} -e GID=#{gid} -e GROUP=_staff -e USER=rb-sys-dock"
end

def interactive?(input_args)
  $stdin.tty?
end

def mount_shared_bind_dir(src, dest)
  "--mount type=bind,source=#{src},destination=#{dest},readonly=false"
end

def mount_tmp_dir
  "--mount type=bind,source=#{Dir.mktmpdir},destination=#{working_directory}/tmp/#{ruby_platform},readonly=false"
end

def toolchain_info
  @toolchain_info ||= OPTIONS.fetch(:toolchain_info) do
    logger.error("Could not determine ruby platform, please set ruby platform with --platform to one of:")
    list_platforms
    logger.fatal("Exiting...")
  end
end

def ruby_platform
  @ruby_platform ||= toolchain_info.platform
end

def rcd(input_args)
  wrapper_command = []
  wrapper_command << "sigfw" unless interactive?(input_args)
  wrapper_command << "runas"

  docker_options = []
  docker_options << "--tty" if interactive?(input_args)

  cmd = <<~SH
    #{default_docker_command} run \
      --platform #{OPTIONS.fetch(:docker_platform)} \
      -v #{working_directory}:#{working_directory} \
      #{mount_tmp_dir} \
      #{mount_target_dir} \
      #{mount_cargo_registry} \
      #{mount_rustup_toolchains} \
      #{mount_bundle_cache} \
      #{mount_bundle_config} \
      #{mount_command_history} \
      #{user_mapping} \
      -e GEM_PRIVATE_KEY_PASSPHRASE \
      -e ftp_proxy \
      -e http_proxy \
      -e https_proxy \
      -e RCD_HOST_RUBY_PLATFORM=#{RbConfig::CONFIG["arch"]} \
      -e RCD_HOST_RUBY_VERSION=#{RUBY_VERSION} \
      -e RCD_IMAGE \
      -e RB_SYS_DOCK_TMPDIR="/tmp/rb-sys-dock" \
      -e RB_SYS_CARGO_TARGET_DIR=#{tmp_target_dir.inspect} \
      #{ENV["RUBY_CC_VERSION"] ? "-e RUBY_CC_VERSION=#{ENV["RUBY_CC_VERSION"]}" : ""} \
      -e RAKEOPT \
      -e TERM \
      -w #{working_directory} \
      --rm \
      --interactive \
      #{docker_options.join(" ")} \
      #{ENV.fetch("RCD_IMAGE")} \
      #{wrapper_command.join(" ")} \
      #{default_command_to_run(input_args)}
  SH

  cmd.gsub!(/\s+/, " ")

  logger.trace("Running command:\n\t$ #{cmd}")

  exec(cmd)
end

def download_image
  image = ENV.fetch("RCD_IMAGE")

  if docker("images -q #{image}").strip.empty? || OPTIONS[:no_cache]
    # Nicely formatted message that we are downloading the image which might take awhile
    logger.info("Downloading container #{image.inspect}, this might take awhile...")
    docker("pull #{image} --platform #{OPTIONS[:docker_platform]} --quiet > /dev/null")
  end
end

def log_some_useful_info
  return if OPTIONS[:build]

  if ARGV.empty?
    logger.info("Entering shell in Docker container #{ENV["RCD_IMAGE"].inspect}")
  else
    logger.info("Running command #{ARGV.inspect} in Docker container #{ENV["RCD_IMAGE"].inspect}")
  end
end

def set_env
  ENV["RCD_IMAGE"] ||= "rbsys/#{ruby_platform}:#{OPTIONS[:version]}"
end

def lint_rb_sys
  cargo_version = cargo_metadata.rb_sys_version
  return if cargo_version == RbSys::VERSION
  logger.warn("Cargo rb-sys version (#{cargo_version}) does not match Ruby gem version (#{RbSys::VERSION})")
rescue => e
  logger.warn("Could not determine Cargo rb-sys version")
  logger.trace("Error was: #{e.inspect}")
end

set_env
lint_rb_sys
download_image
log_some_useful_info
rcd(ARGV)