module.nix

# The importApply argument. Use this to reference things defined locally,
# as opposed to the flake where this is imported.
{
  gepetto,
  flakoboros,
  jrl-cmakemodulesv2,
  make-shell,
  mc-rtc-lib,
  ...
}: # localFlake

# Regular module arguments; self, inputs, etc all reference the final user flake,
# where this module was imported.
{
  lib,
  config,
  ...
}:
let
  cfg = config.mc-rtc-nix;
  flakoboros-cfg = config.flakoboros;
  rosDistro =
    let
      distro = flakoboros-cfg.rosShellDistro;
    in
    builtins.trace "rosDistro evaluated to: ${distro}" distro;

  qtVersion =
    let
      version = flakoboros.lib.ros2qt rosDistro;
    in
    builtins.trace "qtVersion evaluated to: ${version}" version;

  # Define overlays as pure functions to avoid accessing `pkgs` at the root config level
  mcRtcPkgsOverlay =
    final: prev:
    let
      qt = if (qtVersion == "5") then final.qt5 else final.qt6;
    in
    (import ./overlay.nix {
      inherit lib;
      stdenv = prev.stdenv;
      with-ros = cfg.with-ros;
      with-python = cfg.with-python;
      inherit mc-rtc-lib;
      inherit qt;
      inherit rosDistro;
    })
      final
      prev;

  mcRtcPrivateOverlay =
    final: prev:
    let
      qt = if (qtVersion == "5") then final.qt5 else final.qt6;
    in
    (import ./overlay-private.nix {
      with-ros = cfg.with-ros;
      with-python = cfg.with-python;
      inherit qt;
    })
      final
      prev;

  jrlCmakeModulesOverlay = _final: prev: {
    jrl-cmakemodulesv2 = jrl-cmakemodulesv2.packages.${prev.system}.default;
  };

  mcRtcCcacheOverlay = import ./overlay-ccache.nix {
    inherit lib;
    with-private = cfg.overlays.private;
  };

  superbuildFlakeModule = import ./modules/superbuild/superbuild.nix {
    inherit mc-rtc-lib;
    inherit rosDistro;
  };
in
{
  options.mc-rtc-nix = {
    with-ros = lib.mkOption {
      type = lib.types.bool;
      default = true;
      description = "Whether to build with ROS support.";
    };

    with-python = lib.mkOption {
      type = lib.types.bool;
      default = true;
      description = "Whether to build with Python (bindings) support.";
    };

    overlays = {
      private = lib.mkEnableOption "enables the private repository overlay";
      ccache = lib.mkEnableOption "enables the ccache overlay" // {
        default = false;
      };
    };

    gepetto = {
      packages = lib.mkEnableOption "adds gepetto packages";
      devShells = lib.mkEnableOption "adds gepetto devShells";
    };

    packages = lib.mkOption {
      type = lib.types.bool;
      default = false;
      description = "adds mc-rtc-nix packages";
    };
  };

  options.mc-rtc-superbuild = lib.mkOption {
    type = lib.types.deferredModuleWith {
      staticModules = [
        { config.enable = lib.mkDefault false; }
      ];
    };
    default = { };
    description = "mc-rtc superbuild configuration. Accepts a module attrset or a function of { pkgs, ... } returning a module attrset.";
  };

  imports = [
    gepetto.flakeModule
    make-shell.flakeModules.default
  ];

  config = {
    flake.overlays = {
      mc-rtc-pkgs = mcRtcPkgsOverlay;
      make-shell = make-shell.overlays.default;
    }
    // lib.optionalAttrs cfg.overlays.private {
      mc-rtc-private = mcRtcPrivateOverlay;
    }
    // {
      jrl-cmakemodulesv2 = jrlCmakeModulesOverlay;
    }
    // lib.optionalAttrs cfg.overlays.ccache {
      mc-rtc-ccache = mcRtcCcacheOverlay;
    };

    flake.flakeModules.superbuild = superbuildFlakeModule;

    flakoboros = {
      # FIXME: Flakoboros pulls in qtwayland that does not exist on macos
      enableQt = false;
      rosShellDistro = "kilted";
      extraPackages = [ "ninja" ];
      overlays = [
        (builtins.trace "using overlay mc-rtc-pkgs (public)" mcRtcPkgsOverlay)
        (builtins.trace "using overlay make-shell" make-shell.overlays.default)
      ]
      ++ lib.optional cfg.overlays.private (
        builtins.trace "using overlay mc-rtc-private" mcRtcPrivateOverlay
      )
      ++ [ jrlCmakeModulesOverlay ]
      ++ lib.optional cfg.overlays.ccache (builtins.trace "using overlay ccache" mcRtcCcacheOverlay);

      nixpkgsConfig = {
        permittedInsecurePackages = [ "openssl-1.1.1w" ];
      };
    };

    perSystem = (
      { pkgs, inputs', ... }:
      let
        superbuildCfg =
          (lib.evalModules {
            modules = [
              { options = import ./modules/superbuild/options.nix { inherit lib; }; }
              config.mc-rtc-superbuild
            ];
            specialArgs = { inherit pkgs; };
          }).config;

        maybe-mc-mujoco = mujoco-pkg: lib.optional (!pkgs.stdenv.hostPlatform.isDarwin) mujoco-pkg;

        builtInConfigurations =
          let
            genRobotSuperbuild = mcRobotPkgName: mainRobotName: controllerName: {
              extends = [ "default" ];
              enabled = controllerName;
              runtime = {
                robots = [
                  pkgs.${mcRobotPkgName}
                ];
              };
              mainRobot = mainRobotName;
            };
          in
          {
            minimal = {
              runtime = { };
            };

            default = {
              extends = [ "minimal" ];
              runtime = {
                observers = [ pkgs.mc-state-observation ];
                apps = [
                  pkgs.mc-rtc-magnum
                ]
                ++ lib.optionals (cfg.with-ros || superbuildCfg.withRos) [ pkgs.mc-rtc-rviz ];
              };
            };

            # TODO: automatically generate based on a list of robots
            # This list could be automatically constructed from all packages by adding a passthru.mc-rtc.isRobotModule flag
            # to each robot module
            default-all-robots = {
              extends = [ "default" ];
              runtime = with pkgs; {
                robots = [
                  mc-g1
                  mc-h1
                  mc-human
                  mc-robogami
                  mc-triorb
                  mc-pepper
                ]
                ++ lib.optionals (!pkgs.stdenv.hostPlatform.isDarwin) [
                  mc-panda
                  mc-panda-lirmm
                ]
                ++ lib.optionals (cfg.with-ros || superbuildCfg.withRos) [
                  mc-ur5e
                  mc-openarm
                ]
                ++ lib.optionals cfg.overlays.private [
                  mc-hrp2
                  mc-hrp4
                  mc-hrp5-p
                  mc-miroki
                ];
                controllers = [ robogami-controller ];
              };
            };

            mc-openarm = genRobotSuperbuild "mc-openarm" "OpenArmV2" "Posture";
            mc-triorb = genRobotSuperbuild "mc-triorb" "TriOrb" "Posture";

            full = {
              extends = [ "default-all-robots" ];
              runtime = {
                plugins = [
                  pkgs.mc-robot-model-update
                ]
                ++ lib.optionals (!pkgs.stdenv.hostPlatform.isDarwin) [
                  pkgs.mc-force-shoe-plugin
                ];
                apps = maybe-mc-mujoco pkgs.mc-mujoco-full;
              };
            };

            # auto-generated configs from controller's passthru.mc-rtc
            ismpc-walking = mc-rtc-lib.mkControllerSuperbuild pkgs pkgs.ismpc-walking-controller { };
            lipm-walking = mc-rtc-lib.mkControllerSuperbuild pkgs pkgs.lipm-walking-controller { };
            robogami = mc-rtc-lib.mkControllerSuperbuild pkgs pkgs.robogami-controller { };
            panda-prosthesis = mc-rtc-lib.mkControllerSuperbuild pkgs pkgs.panda-prosthesis { };
          };

        configurations = builtInConfigurations // superbuildCfg.configurations;

        mkSuperbuildShell =
          {
            mode,
            shellPname,
            configuration,
          }:
          pkgs.make-shell {
            imports = [ superbuildFlakeModule ];
            mc-rtc-superbuild = superbuildCfg // {
              inherit mode;
              withRos = cfg.with-ros;
              configurations = configurations;
              project = superbuildCfg.project // {
                pname = shellPname;
                inherit configuration;
              };
            };
          };

        shellBaseName = superbuildCfg.project.pname;

        mkShellsByPreset =
          mode: configs:
          builtins.listToAttrs (
            map (
              preset:
              let
                name =
                  (lib.optionalString (shellBaseName != "") "${shellBaseName}-")
                  + "${preset}"
                  + lib.optionalString (mode == "devel") "-devel";
              in
              {
                inherit name;
                value = mkSuperbuildShell {
                  mode = mode;
                  shellPname = name;
                  configuration = preset;
                };
              }
            ) (builtins.attrNames configs)
          );

        filterConfgurations =
          defaultShells_: autoShells_:
          if defaultShells_ then
            configurations
          else
            lib.optionalAttrs autoShells_ superbuildCfg.configurations;
        releaseShellsByPreset = mkShellsByPreset "release" (
          filterConfgurations superbuildCfg.shells.defaultShells.release superbuildCfg.shells.autoShells.release
        );
        develShellsByPreset = mkShellsByPreset "devel" (
          filterConfgurations superbuildCfg.shells.defaultShells.devel superbuildCfg.shells.autoShells.devel
        );

        explicitShells = lib.mapAttrs (
          name: shellCfg:
          mkSuperbuildShell {
            mode = shellCfg.mode;
            shellPname = name;
            configuration = shellCfg.configuration;
          }
        ) superbuildCfg.shells.additionalShells;

        hasExplicitShells = superbuildCfg.shells.additionalShells != { };
        generatedShells =
          if hasExplicitShells then explicitShells else releaseShellsByPreset // develShellsByPreset;
      in
      {
        packages = lib.mkMerge [
          (lib.mkIf cfg.gepetto.packages inputs'.gepetto.packages)
          (lib.mkIf cfg.packages (
            lib.mergeAttrsList [
              {
                # Main dependencies
                inherit (pkgs)
                  eigen3-to-python
                  spacevecalg
                  rbdyn
                  sch-core
                  sch-core-python
                  tasks
                  tasks-qld
                  tvm
                  eigen-quadprog
                  eigen-qld
                  state-observation
                  mesh-sampling
                  copra
                  ;

                # mc-rtc
                inherit (pkgs)
                  mc-rtc-data
                  mc-rtc
                  mc-rtc-ros-compat
                  mc-rtc-ros-compat-noros
                  mc-rtc-ros-compat-ros
                  ;

                # Main GUIs and applications
                inherit (pkgs)
                  mc-rtc-ticker
                  mc-rtc-magnum
                  mc-rtc-imgui
                  imguizmo
                  ;

                # Main robots
                inherit (pkgs)
                  mc-g1
                  mc-h1
                  mc-ur5e
                  mc-human
                  mc-robogami
                  mc-triorb
                  mc-pepper
                  ;

                # Robot description
                inherit (pkgs)
                  mc-int-obj-description
                  jvrc-description
                  g1-description
                  h1-description
                  ur5e-description
                  human-description
                  openarm-description
                  robogami-description
                  triorb-description
                  pepper-description
                  ;

                # MuJoCo Robots
                inherit (pkgs)
                  h1-mj-description
                  jvrc1-mj-description
                  g1-mj-description
                  ur5e-mj-description
                  human-mj-description
                  env-mj-description
                  ;

                # Controllers
                inherit (pkgs)
                  panda-prosthesis
                  robogami-controller
                  lipm-walking-controller
                  # ismpc
                  footsteps-planner-plugin
                  pendulum-feasibility-solver
                  ismpc-walking-controller
                  mc-joystick-plugin
                  ;

                # Plugins
                inherit (pkgs) mc-force-shoe-plugin;

                # Tools
                inherit (pkgs) mc-robot-tools sphinx-cmake;
              }

              (lib.optionalAttrs (cfg.with-ros || superbuildCfg.withRos) {
                inherit (pkgs)
                  mc-rtc-rviz
                  mc-openarm
                  ;
              })
              (lib.optionalAttrs (!pkgs.stdenv.hostPlatform.isDarwin) {
                inherit (pkgs) mc-udp mc-mujoco mc-mujoco-full;
                inherit (pkgs)
                  libfranka_0_9_2
                  mc-franka
                  mc-panda
                  mc-panda-lirmm
                  ;
              })
              (lib.optionalAttrs cfg.overlays.private {
                inherit (pkgs)
                  hrp2-description
                  hrp4-description
                  hrp5-p-description
                  rhps1-description
                  miroki-description
                  ;
                inherit (pkgs)
                  # hrp2-mj-description
                  hrp4-mj-description
                  hrp5p-mj-description
                  rhps1-mj-description
                  # miroki-mj-description
                  ;
                inherit (pkgs)
                  mc-hrp2
                  mc-hrp4
                  # mc-hrp4cr
                  mc-hrp5-p
                  mc-rhps1
                  mc-miroki
                  ;

                inherit (pkgs)
                  tasks-lssol
                  tvm-lssol
                  politopix
                  mc-dynamic-polytopes
                  dcm-vrptask
                  polytopeController
                  ;
              })

              (lib.mapAttrs' (n: lib.nameValuePair "py-${n}") {
                inherit (pkgs.python3Packages)
                  mc-rtc
                  ;
              })
            ]
          ))
        ];

        devShells = lib.mkMerge [
          (lib.mkIf cfg.gepetto.devShells inputs'.gepetto.devShells)
          (lib.mkIf superbuildCfg.enable generatedShells)
        ];
      }
    );
  };
}