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

let
  inherit (builtins) attrNames;
  inherit (lib.nix2lua) pipe1 require call0 set;

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

  treesitterWithGrammars = (cfg.treesitter.package.override {
    inherit (cfg) extraGrammars;
  }).withPlugins (p: map (k: p.${k}) (attrNames cfg.extraGrammars));

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


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

  postPatchExtraQueries = lib.concatLines (lib.flatten
    (lib.flip lib.mapAttrsToList cfg.extraQueries (
      lang:
      lib.mapAttrsToList (queryKind: queries:
        let file = pkgs.writeText "${lang}-${queryKind}" queries; in
        ''
          mkdir -p queries/${lang}
          cat ${file} >> queries/${lang}/${queryKind}.scm
        ''
      )
    ))
  );

  finalNvimTreeSitter =
    if cfg.extraGrammars == { } && cfg.extraQueries == { } 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
          + postPatchExtraGrammars
          + postPatchExtraQueries;
      });
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 = { };
    };

    extraQueries = mkOption {
      type = with types; attrsOf (submodule {
        options = {
          injections = mkOption {
            type = str;
            default = "";
          };
        };
      });
      default = { };
      example = {
        javascript.injections = ''(comment) @comment'';
      };
    };

    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;

      extraImports = {
        parser_config = pipe1 (require "nvim-treesitter.parsers") (call0 "get_parser_configs");
      };

      beforeSetup = with lib; optionals (cfg.extraGrammars != { }) (
        mapAttrsToList (k: v: set (pipe1 "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
    };
  };

}