class Tasks::SwitchRoot < SingletonTask
  module Utils
    # Resolves symlinks under a given root to get the real name of the file.
    # The returned path is not prefixed with given root either.
    def readlink_rooted(root, filename)
      # Resolve the full pathname
      loop do
        prev_filename = filename

        if File.symlink?(File.join(root, prev_filename))
          filename = File.readlink(File.join(root, prev_filename))

          # Relative link? Make absolute.
          unless filename.match(%r{^/})
            filename = File.join(File.dirname(prev_filename), filename)
          end
        end
        break if prev_filename == filename
      end

      filename
    end
  end

  # Extracted bits for stage-0 support.
  # This should make deprecating and removing easier down the line.
  module StageZeroSupport
    # Returns whether the current boot will `kexec` into the generation.
    def will_kexec?()
      # Only stage-0 bootloader-flavoured init will kexec.
      return false unless STAGE == 0

      # The user wants to use the generation's kernel
      return false unless @use_generation_kernel

      # AND if we find the required files.
      [
        "initrd",
        "kernel",
        "kernel-params",
      ]
        .map { |file| generation_file(file) }
        .map { |file| File.exist?(file) }
        .all?
    end

    def kexec_in_generation(selected_generation)
      if Tasks.constants.include?(:SetupGadgetMode)
        Progress.exec_with_message("Tearing down USB Gadget mode") do
          begin
            Tasks::SetupGadgetMode.instance.teardown()
          rescue => e
            $logger.fatal("Caught an error during teardown for kexec...")
            $logger.fatal(e.inspect)
          end
        end
      end

      log("About to kexec...")

      System.run(
        "kexec", "--load",
        generation_file("kernel"),
        "--initrd=#{generation_file("initrd")}",
        *[maybe_dtb()].compact(),
        "--command-line",
        [
          ["init", selected_generation.path("init", readlink: true)].join("="),
          # Flag used to describe we're in a kexec situation.
          # For the time being, the flag is the whole string, not the value yes to that key.
          "mobile-nixos.kexec=yes",
          "mobile-nixos.generation=#{selected_generation.path}",
          File.read(generation_file("kernel-params")),
        ].join(" ")
      )
      System.exec("kexec", "-e")
      # Execution shouldn't reach here, as `exec` replaces this process with `kexec`.
    ensure
      raise "Failed to kexec into #{selected_generation.path}"
    end

    # If the system build exports a DTB file name to load, return the appropriate
    # command-line argument for `kexec`. Otherwise nil.
    def maybe_dtb()
      log("Looking for a DTB file...")

      mapping_path = generation_file("mobile-nixos/dtb-mapping.json", missing_allowed: true)
      unless mapping_path && File.exist?(mapping_path)
        log("  DTB: dtb-mapping.json not found... skipping DTB mapping...")
        return nil
      end

      log("  DTB: Mappings from: `#{mapping_path}`...")
      mapping = JSON.parse(File.read(mapping_path))

      # Work only off the first compatible name.
      # It is assumed that the loader uses the exact same scheme as the kernel
      # build does. If this assumption stops holding true, it will be a new
      # feature to implement.
      board_compatible = File.read("/proc/device-tree/compatible").split("\0").first
      log("  DTB: board_compatible: #{board_compatible}")

      desired_dtb = mapping[board_compatible]
      if desired_dtb
        log("  DTB: wants: #{desired_dtb}")
      else
        log("  DTB: no DTB mapping found for #{board_compatible}...")
        return nil
      end

      # The desired_dtb path is an absolute path in the mounted system.
      file = File.join(SYSTEM_MOUNT_POINT, desired_dtb)

      if File.exist?(file)
        log("  DTB: file `#{file}` found")
        "--dtb=#{forward_fdt_bootloader_info(file)}"
      else
        log("  DTB: file `#{file}` not found... skipping DTB mapping")
        nil
      end
    end

    # Given a path to a DTB file, merges required properties that the bootloader
    # has setup. It will additionally merge optional properties.
    def forward_fdt_bootloader_info(path)
      args = [
        "--print-header",
        "--copy-dtb", path,
        Configuration["stage-0"]["forward"]["nodes"].map {|path| [ "--forward-node", path] },
        Configuration["stage-0"]["forward"]["props"].map {|pair| [ "--forward-prop", *pair] },
      ].flatten
      log(" $ fdt-forward #{args.shelljoin}")
      dts = `fdt-forward #{args.shelljoin}`

      # Declare we booted using stage-0's kexec
      # And additional useful debugging data...
      dts << [
        "\n",
        "// Declare we booted using kexec",
        %Q{/ { mobile-nixos,stage-0; };},
        %Q{/ { mobile-nixos,stage-0,timestamp = #{Time.now.to_s.to_json}; };},
        %Q{/ { mobile-nixos,stage-0,uname = #{`uname -a`.to_json}; };},
        %Q{/ { mobile-nixos,stage-0,uptime = #{`uptime`.to_json}; };},
      ].join("\n")

      FileUtils.mkdir_p("/run/boot/")
      File.write("/run/boot/fdt.dts", dts)
      System.run("fdt-forward --to-dtb < /run/boot/fdt.dts > /run/boot/fdt.dtb")

      return "/run/boot/fdt.dtb"
    end
  end

  include StageZeroSupport

  # Where the system rootfs will be mounted.
  SYSTEM_MOUNT_POINT = "/mnt"

  #
  # Represents resources and facts from a NixOS generation.
  # 
  # A generation can be:
  #  - a symlink (e.g. from `/nix/var/nix/profiles`)
  #  - a bare store path.
  #
  class NixOSGeneration
    include Utils

    NIX_PROFILES_PATH = "/nix/var/nix/profiles"
    DEFAULT_PROFILE_LINK = File.join(NIX_PROFILES_PATH, "system")

    @@default_generation = nil

    # path: path to a generation (absolute according to the system rootfs)
    def initialize(path)
      @generation_path = path
    end

    # Whether the generation path exists or not.
    def exist?()
      # A symlink to a dangling path does not exist.
      File.exist?(mounted_path) || File.symlink?(mounted_path)
    end

    # A label for the generation
    def label()
      date = File.lstat(mounted_path).mtime.strftime("%F")

      version_file = mounted_path("nixos-version", readlink: true)
      version =
        if File.exist?(version_file)
          File.read(version_file)
        else
          nil
        end

      details = [
        date,
        version,
      ].compact.join(" - ")

      num = path.split("-")[-2]

      "NixOS ##{num} (#{details})"
    end

    # This generation's path, or a path within.
    def path(target = nil, readlink: false)
      resolved_path = @generation_path
      resolved_path = readlink_rooted(SYSTEM_MOUNT_POINT, resolved_path) if readlink

      return resolved_path unless target

      target = File.join(resolved_path, target)
      return target unless readlink

      readlink_rooted(SYSTEM_MOUNT_POINT, target)
    end

    # Path of the generation, or a path within, according to its location during stage-1.
    def mounted_path(target = nil, readlink: false)
      target = path(target, readlink: readlink)
      File.join(SYSTEM_MOUNT_POINT, target)
    end

    def self.default_generation()
      return @@default_generation if @@default_generation
      @@default_generation = NixOSGeneration.new(DEFAULT_PROFILE_LINK)
    end

    # Returns the list of generations, ordered from newest to oldest.
    def self.generations()
      Dir.glob(File.join(SYSTEM_MOUNT_POINT, NixOSGeneration::DEFAULT_PROFILE_LINK) + "-*")
        .sort do |a, b|
          # a and b are reversed here to sort descending.
          File.lstat(b).mtime <=> File.lstat(a).mtime
        end
        .map do |path|
          NixOSGeneration.new(path.delete_prefix(SYSTEM_MOUNT_POINT))
        end
    end
  end

  def initialize()
    add_dependency(:Task, Tasks::Splash.instance)
    add_dependency(:Target, :SwitchRoot)

    # By default, with stage-0, we prefer using the generation kernel
    # This may be overriden by the user recovery user interface
    @use_generation_kernel = STAGE == 0
  end

  # Creates the generation selection list.
  def generate_selection()
    selection = NixOSGeneration.generations.map do |generation|
      {
        id: generation.path,
        name: generation.label,
      }
    end

    # Prepend a special entry representing the default action.
    selection.unshift(
      {
        id: "$default",
        name: "NixOS - Default",
      }
    )

    FileUtils.mkdir_p("/run/boot/")
    # The recovery applet uses this file to allow generation selection.
    File.write("/run/boot/selection.json", selection.to_json)
  end

  #
  # Pick the default generation.
  #
  # This does either of:
  #
  #  * Booting the generation given in boot ("kernel") parameters.
  #     * From stage-0
  #     * From bootloader
  #  * Booting the default `system` profile link.
  #  * Guess and find the generation store path for first boot.
  #
  # This is *always* a sane default to fallback on.
  # This is what the `$default` generation button picks, and what is started by default.
  #
  # Returns the path to a generation (absolute according to the system rootfs)
  #
  def default_selection_path()
    # Given as a command-line option, most likely from stage-0.
    generation_parameter = System.cmdline().grep(/^mobile-nixos.generation=/).first
    unless generation_parameter.nil?
      $logger.info("Using '#{generation_parameter}' cmdline parameter to select generation...")
      return generation_parameter.split("=", 2).last
    end

    # Given as a command-line option, from the bootloader (replacement for NixOS's stage-1)
    init_parameter = System.cmdline().grep(/^init=/).first
    if init_parameter == "init=/init" then
      $logger.info("Skipping '#{init_parameter}' cmdline parameter from quirky device...")
    else
      unless init_parameter.nil?
        $logger.info("Using '#{init_parameter}' cmdline parameter to select generation...")
        init_parameter = init_parameter.split("=", 2).last
        return init_parameter.rpartition("/").first
      end
    end

    # The default generation
    if NixOSGeneration.default_generation().exist?()
      generation = NixOSGeneration.default_generation().path
      $logger.info("Using '#{generation}' default generation...")
      return generation
    end

    # Otherwise, we need to re-hydrate a system!
    registration = File.join(SYSTEM_MOUNT_POINT, "nix-path-registration")
    if File.exist?(registration)
      $logger.info("Getting NixOS generation from nix-path-registration...")
      path = File.read(registration)
        .split("\n")
        .grep(%r{^/nix/store/[a-z0-9]+-nixos-system-})
        .first
      if path
        $logger.info("... found '#{path}'")
        return path
      end
    end

    System.failure("INIT_NOT_FOUND", "Stage-2 init not found", "Could not find init path for stage-2", color: "FF00FF")
  end

  # Pauses the boot to allow the user to select a generation.
  def choose_generation()
    generate_selection()

    # Synchronuously (pause the init code) show the selection applet.
    System.run(LOADER, "/applets/boot-selection.mrb")

    # Read data from the user
    data = JSON.parse(File.read("/run/boot/choice"))
    generation = data["generation"]
    @use_generation_kernel = data["use_generation_kernel"]

    # Why "$default" rather than passing a path?
    # Because there may be no generations folder. It's easier to cheat and
    # use "$default" and rely on the existing default "maybe rehydrate"
    # codepath.
    generation = default_selection_path() if generation == "$default"
    NixOSGeneration.new(generation)
  end

  def selected_generation()
    return @selected_generation if @selected_generation

    if Hal::Recovery.wants_recovery?
      Tasks::Splash.instance.quit("Continuing to recovery menu")
      @selected_generation = choose_generation()
    else
      @selected_generation = NixOSGeneration.new(default_selection_path())
      if will_kexec?()
        Tasks::Splash.instance.quit("Rebooting in generation kernel", sticky: true)
      else
        Tasks::Splash.instance.quit("Continuing to stage-2")
      end
    end
    @selected_generation
  end

  def generation_file(name, missing_allowed: false)
    begin
      selected_generation.mounted_path(name, readlink: true)
    rescue => e
      log "While searching for generation_file(#{name.inspect}):"
      log e.inspect
      if missing_allowed
        return nil
      else
        raise e
      end
    end
  end

  def run()
    init = File.join(selected_generation.path, "init")

    # This is the traditional way we printed the init path.
    # This is still helpful to take vertical real estate when visually looking
    # through the log.
    log("")
    log("***")
    log("")
    if will_kexec?
      log("Kexecing into #{selected_generation.path}")
    else
      log("Switching root to #{init}")
    end
    log("")
    log("***")
    log("")

    kexec_in_generation(selected_generation) if will_kexec?

    Tasks::UDev.instance.teardown()

    [
      "/proc",
      "/sys",
      "/dev",
      "/run",
    ].each do |mount_point|
      new_location = File.join(SYSTEM_MOUNT_POINT, mount_point)
      FileUtils.mkdir_p(new_location)
      System.run("mount", "--move", mount_point, new_location)
    end

    switch_root = System.which("switch_root")
    System.exec({}, switch_root, SYSTEM_MOUNT_POINT, init)
  end
end
