{ config, lib, pkgs, ... }:

let
  inherit (builtins) elem attrNames;

  cfg = config.plugins.style.nvim-treesitter;

  treesitterWithGrammars = (cfg.treesitter.package.override { inherit (cfg) extraGrammars; }).withPlugins (lib.filterAtttrs (k: elem k (attrNames cfg.extraGrammars)));

  nvimTreeSitterWithBuiltinGrammars =
    if cfg.grammars == null then cfg.package.withAllGrammars
    else cfg.package.withPlugins cfg.grammars;

  finalNvimTreeSitter =
    if cfg.extraGrammars == { } then nvimTreeSitterWithBuiltinGrammars
    else
      nvimTreeSitterWithBuiltinGrammars.overrideAttrs (oldAttrs: {
        passthru.dependencies = oldAttrs.passthru.dependencies
          ++ (lib.flip lib.mapAttrsToList cfg.extraGrammars (k: { language, ... }:
          pkgs.runCommand "nvim-treesitter-${language}-grammar" { } ''
            mkdir -p $out/parser
            ln -s ${treesitterWithGrammars}/${language}.so $out/parser/${language}.so
          ''
        ));

        postPatch = oldAttrs.postPatch + (lib.concatLines
          (lib.flip lib.mapAttrsToList cfg.extraGrammars (k: { src, language, ... }:
            ''
              if [ -d ${src}/queries ]; then
                ln -s ${src}/queries queries/${language}
              fi
            ''
          ))
        );

      });
in
{
  options.plugins.style.nvim-treesitter = with lib; {
    enable = mkEnableOption "nvim-treesitter";

    package = mkPackageOption pkgs.vimPlugins "nvim-treesitter" { };

    finalPackage = mkOption {
      type = types.package;
      visible = false;
      readOnly = true;
      description = "Resulting customized nvim-treesitter package.";
    };

    treesitter.package = mkPackageOption pkgs "tree-sitter" { };

    grammars = mkOption {
      type = with types; nullOr (functionTo (listOf package));
      default = null;
    };

    extraGrammars = mkOption {
      type = with types; attrsOf attrs;
      default = { };
    };

    settings = mkOption {
      type = types.attrs;
      default = { };
      description = ''
        Settings to configure nvim-treesitter.

        See: https://github.com/nvim-treesitter/nvim-treesitter?tab=readme-ov-file#modules

        - `highlight` and `indent` are enabled by default.
        - `ensure_installed`, `sync_install`, `auto_install` will be ignored.
      '';
    };

    fold.enable = mkEnableOption "use nvim treesitter to fold expression" // { default = true; };
  };

  config = lib.mkIf cfg.enable {
    plugins.style.nvim-treesitter.finalPackage = finalNvimTreeSitter;

    plugin.nvim-treesitter = {
      name = "nvim-treesitter.configs";
      package = finalNvimTreeSitter;

      beforeSetup = with lib.nix2lua; lib.optionals (cfg.extraGrammars != { }) (lib.flatten [
        (local (set "parser_config")
          (pipe1 (require "nvim-treesitter.parsers") (call0 "get_parser_configs"))
        )

        (lib.mapAttrsToList (k: v: set "parser_config.${v.language}" { }) cfg.extraGrammars)
      ]);

      setupSettings = lib.mkMerge [
        # enable hihlight and indent by default
        {
          highlight.enable = true;
          indent.enable = true;
        }

        # apply user settings
        cfg.settings

        # force disable manual installation
        # TODO: warn about that behavior
        {
          ensure_installed = { };
          sync_install = false;
          auto_install = false;
        }
      ];
    };

    vim.opt = lib.mkIf cfg.fold.enable {
      foldmethod = "expr";
      foldexpr = "nvim_treesitter#foldexpr()";
      foldenable = lib.mkDefault false; # disable folding at startup
    };
  };

}