# Component resolution, aggregation and other utility functions.
# Provide the content of `rust-bin`.
{
  lib,
  pkgs,
  manifests,
  nightly,
}:
let
  inherit (builtins)
    compareVersions
    fromTOML
    match
    readFile
    tryEval
    ;

  inherit (lib)
    any
    attrNames
    attrValues
    concatStringsSep
    elem
    elemAt
    filter
    flatten
    foldl'
    hasPrefix
    head
    isString
    length
    listToAttrs
    makeOverridable
    mapAttrs
    mapAttrsToList
    optional
    optionalAttrs
    replaceStrings
    substring
    trace
    unique
    ;

  inherit (pkgs) stdenv callPackage fetchurl;

  # Remove keys from attrsets whose value is null.
  removeNulls = set: removeAttrs set (filter (name: set.${name} == null) (attrNames set));

  toRustTarget = platform: platform.rust.rustcTarget;

  # The platform where `rustc` is running.
  rustHostPlatform = toRustTarget stdenv.hostPlatform;
  # The platform of binary which `rustc` produces.
  rustTargetPlatform = toRustTarget stdenv.targetPlatform;

  mkComponentSet = callPackage ./mk-component-set.nix {
    inherit removeNulls toRustTarget;
  };

  mkAggregated = callPackage ./mk-aggregated.nix { };

  # Manifest selector.
  selectManifest =
    {
      channel,
      date ? null,
    }:
    let
      assertWith =
        cond: msg: body:
        if cond then body else throw msg;

      # https://rust-lang.github.io/rustup/concepts/toolchains.html#toolchain-specification
      # <channel> = stable|beta|nightly|<major.minor>|<major.minor.patch>

      asVersion = match "[0-9]+\\.[0-9]+(\\.[0-9]+)?" channel;
      asNightlyDate =
        let
          m = match "nightly-([0-9]+-[0-9]+-[0-9]+)" channel;
        in
        if m == null then null else elemAt m 0;
      asBetaDate =
        let
          m = match "beta-([0-9]+-[0-9]+-[0-9]+)" channel;
        in
        if m == null then null else elemAt m 0;

      maxWith = zero: f: foldl' (lhs: rhs: if lhs == zero || f lhs rhs < 0 then rhs else lhs) zero;

      latestStableWithMajorMinor = maxWith "" compareVersions (
        filter (hasPrefix (channel + ".")) (attrNames manifests.stable)
      );

    in
    # "stable"
    if channel == "stable" then
      assertWith (
        date == null
      ) "Stable version with specific date is not supported" manifests.stable.latest
    # "nightly"
    else if channel == "nightly" then
      manifests.nightly.${if date != null then date else "latest"}
        or (throw "Nightly ${date} is not available")
    # "beta"
    else if channel == "beta" then
      manifests.beta.${if date != null then date else "latest"} or (throw "Beta ${date} is not available")
    # "1.49.0" or "1.49"
    else if asVersion != null then
      assertWith (date == null) "Stable version with specific date is not supported" (
        # "1.49"
        if asVersion == [ null ] then
          manifests.stable.${latestStableWithMajorMinor} or (throw "No stable ${channel}.* is available")
        # "1.49.0"
        else
          manifests.stable.${channel} or (throw "Stable ${channel} is not available")
      )
    # "beta-2021-01-01"
    else if asBetaDate != null then
      assertWith (date == null) "Cannot specify date in both `channel` and `date`"
        manifests.beta.${asBetaDate} or (throw "Beta ${asBetaDate} is not available")
    # "nightly-2021-01-01"
    else if asNightlyDate != null then
      assertWith (date == null) "Cannot specify date in both `channel` and `date`"
        manifests.nightly.${asNightlyDate} or (throw "Nightly ${asNightlyDate} is not available")
    # Otherwise
    else
      throw "Unknown channel: ${channel}";

  # Select a toolchain and aggregate components by rustup's `rust-toolchain` file format.
  # See: https://rust-lang.github.io/rustup/concepts/profiles.html
  # Or see source: https://github.com/rust-lang/rustup/blob/84974df1387812269c7b29fa5f3bb1c6480a6500/doc/src/overrides.md#the-toolchain-file
  fromRustupToolchain =
    {
      path ? null,
      channel ? null,
      profile ? null,
      components ? [ ],
      targets ? [ ],
    }:
    if path != null then
      throw "`path` is not supported, please directly add it to your PATH instead"
    else if channel == null then
      throw "`channel` is required"
    else
      let
        toolchain = toolchainFromManifest (selectManifest {
          inherit channel;
        });
        profile' = if profile == null then "default" else profile;
        pkg =
          if toolchain._profiles != { } then
            toolchain._profiles.${profile'} or (throw ''
              Rust ${toolchain._version} doesn't have profile `${profile'}`.
              Available profiles are: ${concatStringsSep ", " (attrNames toolchain._profiles)}
            '')
          # Fallback to package `rust` when profiles are not supported and not specified.
          else if profile == null then
            toolchain.rust
          else
            throw "Cannot select profile `${profile'}` since rust ${toolchain._version} is too early to support profiles";
      in
      pkg.override {
        extensions = components;
        inherit targets;
      };

  # Same as `fromRustupToolchain` but read from a `rust-toolchain` file (legacy one-line string or in TOML).
  fromRustupToolchainFile =
    path:
    let
      content = readFile path;
      legacy = match "([^\r\n]+)\r?\n?" content;
    in
    if legacy != null then
      fromRustupToolchain { channel = head legacy; }
    else
      fromRustupToolchain (fromTOML content).toolchain;

  mkComponentSrc =
    { url, sha256 }:
    let
      url' = replaceStrings [ " " ] [ "%20" ] url; # This is required or download will fail.
      # Filter names like `llvm-tools-1.34.2 (6c2484dc3 2019-05-13)-aarch64-unknown-linux-gnu.tar.xz`
      matchParenPart = match ".*/([^ /]*) [(][^)]*[)](.*)" url;
      name = if matchParenPart == null then "" else (elemAt matchParenPart 0) + (elemAt matchParenPart 1);
    in
    fetchurl {
      inherit name sha256;
      url = url';
    };

  # Resolve final components to install from mozilla-overlay style `extensions`, `targets` and `targetExtensions`.
  #
  # `componentSet` has a layout of `componentSet.<name>.<rust-target> : Derivation`.
  # `targetComponentsList` is a list of all component names for target platforms.
  # `name` is only used for error message.
  #
  # Returns a list of component derivations, or throw if failed.
  resolveComponents =
    {
      name,
      componentSet,
      allComponentSet,
      allPlatformSet,
      targetComponentsList,
      profileComponents,
      extensions,
      targets,
      targetExtensions,
    }:
    let
      # Components for target platform like `rust-std`.
      collectTargetComponents =
        allowMissing: name:
        let
          targetSelected = flatten (map (tgt: componentSet.${tgt}.${name} or [ ]) targets);
        in
        if !allowMissing -> targetSelected != [ ] then
          targetSelected
        else
          "Component `${name}` doesn't support any of targets: ${concatStringsSep ", " targets}";

      collectComponents =
        allowMissing: name:
        if elem name targetComponentsList then
          collectTargetComponents allowMissing name
        else
          # Components for host platform like `rustc`.
          componentSet.${rustHostPlatform}.${name}
            or (if allowMissing then [ ] else "Host component `${name}` doesn't support `${rustHostPlatform}`");

      # Profile components can be skipped silently when missing.
      # Eg. `rust-mingw` on non-Windows platforms, or `rust-docs` on non-tier1 platforms.
      result =
        flatten (map (collectComponents true) profileComponents)
        ++ flatten (map (collectComponents false) extensions)
        ++ flatten (map (collectTargetComponents false) targetExtensions);

      isTargetUnused =
        target:
        !any (name: componentSet ? ${target}.${name})
          # FIXME: Get rid of the legacy component `rust`.
          (
            filter (name: name == "rust" || elem name targetComponentsList) (profileComponents ++ extensions)
            ++ targetExtensions
          );

      # Fail-fast for typo in `targets`, `extensions`, `targetExtensions`.
      fastErrors = flatten (
        map (
          tgt:
          optional (
            !(allPlatformSet ? ${tgt})
          ) "Unknown target `${tgt}`, typo or not supported by this version?"
        ) targets
        ++ map (
          name:
          optional (
            !(allComponentSet ? ${name})
          ) "Unknown component `${name}`, typo or not support by this version?"
        ) (profileComponents ++ extensions ++ targetExtensions)
      );

      errors =
        if fastErrors != [ ] then
          fastErrors
        else
          filter isString result
          ++ map (tgt: "Target `${tgt}` is not supported by any components or extensions") (
            filter isTargetUnused targets
          );

      notes = [
        "note: profile components: ${toString profileComponents}"
      ]
      ++ optional (targets != [ ]) "note: selected targets: ${toString targets}"
      ++ optional (extensions != [ ]) "note: selected extensions: ${toString extensions}"
      ++ optional (
        targetExtensions != [ ]
      ) "note: selected targetExtensions: ${toString targetExtensions}"
      ++ flatten (
        map (
          platform:
          optional (componentSet ? ${platform})
            "note: components available for ${platform}: ${toString (attrNames componentSet.${platform})}"
        ) (unique ([ rustHostPlatform ] ++ targets))
      )
      ++ [
        ''
          note: check here to see all targets and which components are available on each targets:
                https://rust-lang.github.io/rustup-components-history
        ''
      ];

    in
    if errors == [ ] then
      result
    else
      throw ''
        Component resolution failed for ${name}
        ${concatStringsSep "\n" (errors ++ notes)}
      '';

  # Generate the toolchain set from a parsed manifest.
  #
  # Manifest files are organized as follow:
  # { date = "2017-03-03";
  #   pkg.cargo.version= "0.18.0-nightly (5db6d64 2017-03-03)";
  #   pkg.cargo.target.x86_64-unknown-linux-gnu = {
  #     available = true;
  #     hash = "abce..."; # sha256
  #     url = "https://static.rust-lang.org/dist/....tar.gz";
  #     xz_hash = "abce..."; # sha256
  #     xz_url = "https://static.rust-lang.org/dist/....tar.xz";
  #   };
  # }
  #
  # The packages available usually are:
  #   cargo, rust-analysis, rust-docs, rust-src, rust-std, rustc, and
  #   rust, which aggregates them in one package.
  #
  # For each package the following options are available:
  #   extensions        - The extensions that should be installed for the package.
  #                       For example, install the package rust and add the extension rust-src.
  #   targets           - The package will always be installed for the host system, but with this option
  #                       extra targets can be specified, e.g. "mips-unknown-linux-musl". The target
  #                       will only apply to components of the package that support being installed for
  #                       a different architecture. For example, the rust package will install rust-std
  #                       for the host system and the targets.
  #   targetExtensions  - If you want to force extensions to be installed for the given targets, this is your option.
  #                       All extensions in this list will be installed for the target architectures.
  #                       *Attention* If you want to install an extension like rust-src, that has no fixed architecture (arch *),
  #                       you will need to specify this extension in the extensions options or it will not be installed!
  toolchainFromManifest =
    manifest:
    let
      # platform -> true
      # For fail-fast test.
      allPlatformSet = listToAttrs (
        flatten (
          mapAttrsToList (
            compName:
            { target, ... }:
            map (platform: {
              name = platform;
              value = true;
            }) (attrNames target)
          ) manifest.pkg
        )
      );

      # componentName -> true
      # May also contains unavailable components. Just for fail-fast test.
      allComponentSet = mapAttrs (compName: _: true) (manifest.pkg // manifest.renames);

      # componentSet.x86_64-unknown-linux-gnu.cargo = <derivation>;
      componentSet = mapAttrs (
        platform: _:
        mkComponentSet {
          inherit (manifest) version renames;
          inherit platform;
          srcs = removeNulls (
            mapAttrs (
              compName:
              { target, ... }:
              let
                content = target.${platform} or target."*" or null;
              in
              if content == null then
                null
              else
                mkComponentSrc {
                  url = content.xz_url;
                  sha256 = content.xz_hash;
                }
            ) manifest.pkg
          );
        }
      ) allPlatformSet;

      mkProfile =
        name: profileComponents:
        makeOverridable
          (
            {
              extensions,
              targets,
              targetExtensions,
            }:
            mkAggregated {
              pname = "rust-${name}";
              inherit (manifest) version date;
              availableComponents = componentSet.${rustHostPlatform};
              selectedComponents = resolveComponents {
                name = "rust-${name}-${manifest.version}";
                inherit
                  allPlatformSet
                  allComponentSet
                  componentSet
                  profileComponents
                  targetExtensions
                  ;
                inherit (manifest) targetComponentsList;
                extensions = extensions;
                targets = unique (
                  [
                    rustHostPlatform # Build script requires host std.
                    rustTargetPlatform
                  ]
                  ++ targets
                );
              };
            }
          )
          {
            extensions = [ ];
            targets = [ ];
            targetExtensions = [ ];
          };

      profiles = mapAttrs mkProfile manifest.profiles;

      result =
        # Individual components.
        componentSet.${rustHostPlatform}
        //
          # Profiles.
          profiles
        // {
          # Legacy support for special pre-aggregated package.
          # It has more components than `default` profile but less than `complete` profile.
          rust =
            let
              pkg = mkProfile "legacy" [ "rust" ];
            in
            if profiles != { } then
              trace ''
                Rust ${manifest.version}:
                Pre-aggregated package `rust` is not encouraged for stable channel since it contains almost all and uncertain components.
                Consider use `default` profile like `rust-bin.stable.latest.default` and override it with extensions you need.
                See README for more information.
              '' pkg
            else
              pkg;
        };

    in
    # If the platform is not supported for the current version, return nothing here,
    # so others can easily check it by `toolchain ? default`.
    optionalAttrs (componentSet ? ${rustHostPlatform}) result
    // {
      # Internal use.
      _components = componentSet;
      _profiles = profiles;
      _version = manifest.version;
      _manifest = manifest;
    };

  # From a git revision of rustc.
  # This does the same thing as crate `rustup-toolchain-install-master`.
  # But you need to manually provide component hashes.
  fromRustcRev =
    {
      # Package name of the derivation.
      pname ? "rust-custom",
      # Git revision of rustc.
      rev,
      # Version of the built package.
      version ? substring 0 7 rev,
      # Attrset with component name as key and its SRI hash as value.
      components,
      # Rust target to download.
      target ? rustTargetPlatform,
    }:
    let
      hashToSrc =
        compName: hash:
        fetchurl {
          url =
            if compName == "rust-src" then
              "https://ci-artifacts.rust-lang.org/rustc-builds/${rev}/${compName}-nightly.tar.xz"
            else
              "https://ci-artifacts.rust-lang.org/rustc-builds/${rev}/${compName}-nightly-${target}.tar.xz";
          inherit hash;
        };
      components' = mkComponentSet {
        inherit version;
        platform = target;
        srcs = mapAttrs hashToSrc components;
        # We cannot know aliases in this case.
        renames = { };
      };
    in
    mkAggregated {
      inherit pname version;
      date = null;
      selectedComponents = attrValues components';
    };

  # Select latest nightly toolchain which makes selected profile builds.
  # Some components are missing in some nightly releases.
  # Usage:
  # `selectLatestNightlyWith (toolchain: toolchain.default.override { extensions = ["llvm-tools-preview"]; })`
  selectLatestNightlyWith =
    selector:
    let
      nightlyDates = attrNames (removeAttrs nightly [ "latest" ]);
      dateLength = length nightlyDates;
      go =
        idx:
        let
          ret = selector (nightly.${elemAt nightlyDates idx});
        in
        if idx == 0 then
          ret
        else if dateLength - idx >= 256 then
          trace "Failed to select nightly version after 100 tries" ret
        else if ret != null && (tryEval ret.drvPath).success then
          ret
        else
          go (idx - 1);
    in
    go (length nightlyDates - 1);

in
# For each channel:
#   rust-bin.stable.latest.{minimal,default,complete} # Profiles.
#   rust-bin.stable.latest.rust   # Pre-aggregate from upstream.
#   rust-bin.stable.latest.cargo  # Components...
#   rust-bin.stable.latest.rustc
#   rust-bin.stable.latest.rust-docs
#   ...
#
# For a specific version of stable:
#   rust-bin.stable."1.47.0".default
#
# For a specific date of beta:
#   rust-bin.beta."2021-01-01".default
#
# For a specific date of nightly:
#   rust-bin.nightly."2020-01-01".default
mapAttrs (channel: mapAttrs (version: toolchainFromManifest)) manifests
// {
  inherit fromRustupToolchain fromRustupToolchainFile;
  inherit selectLatestNightlyWith;
  inherit fromRustcRev;

  _internal = {
    inherit toolchainFromManifest;
    inherit selectManifest;
  };
}