(require 'treesit)

(define-derived-mode powershell-ts-mode prog-mode "PS[ts]"
  (when (treesit-ready-p 'powershell)
    (treesit-parser-create 'powershell)
    
    (setq-local treesit-font-lock-settings
                (treesit-font-lock-rules
                 :language 'powershell
                 :feature 'comment
                 '((comment) @font-lock-comment-face)
                 :language 'powershell
                 :feature 'variable
                 '((unary_expression (variable) @font-lock-variable-name-face)
                   (variable) @font-lock-variable-name-face)
                 :language 'powershell
                 :feature 'string
                 '((string_literal (verbatim_string_characters) @font-lock-string-face)
                   (string_literal (expandable_string_literal) @font-lock-string-face)
                   (string_literal (expandable_here_string_literal) @font-lock-string-face))
                 :language 'powershell
                 :feature 'function
                 '((enum_statement "enum" @font-lock-operator-face (simple_name) @font-lock-function-name-face)
                   (class_statement "class" @font-lock-keyword-face (simple_name))
                   (class_method_definition (simple_name) @font-lock-function-name-face)
                   (function_statement "function" @font-lock-operator-face (function_name) @font-lock-function-name-face)
                   (function_statement "workflow" @font-lock-operator-face (function_name) @font-lock-function-name-face)
                   (command command_name: (command_name) @font-lock-function-call-face)
                   (command_parameter) @font-lock-constant-face
                   (param_block "param" @font-lock-keyword-face)
                   (comparison_operator) @font-lock-builtin-face
                   (if_statement "if" @font-lock-keyword-face)
                   (else_clause "else" @font-lock-keyword-face)
                   (elseif_clause "elseif" @font-lock-keyword-face)
                   (switch_statement "switch" @font-lock-keyword-face)
                   (foreach_statement "foreach" @font-lock-keyword-face)
                   (for_statement "for" @font-lock-keyword-face)
                   (while_statement "while" @font-lock-keyword-face)
                   (do_statement "do" @font-lock-keyword-face "while" @font-lock-keyword-face)
                   (do_statement "do" @font-lock-keyword-face "until" @font-lock-keyword-face)
                   (flow_control_statement "continue" @font-lock-keyword-face)
                   (flow_control_statement "break" @font-lock-keyword-face)
                   (flow_control_statement "return" @font-lock-keyword-face)
                   (flow_control_statement "throw" @font-lock-keyword-face)
                   (flow_control_statement "exit" @font-lock-keyword-face)
                   (type_literal) @font-lock-type-face
                   (try_statement "try" @font-lock-keyword-face)
                   (catch_clause "catch" @font-lock-keyword-face)
                   (finally_clause "finally" @font-lock-keyword-face)
                   (trap_statement "trap" @font-lock-keyword-face)
                   (parallel_statement "parallel" @font-lock-keyword-face))))

    (setq-local treesit-font-lock-feature-list '((comment) (variable) (string) (function)))
    (setq-local treesit-font-lock-level 5)
    
    (setq-local treesit-simple-indent-rules
                `((powershell
                   ((parent-is "program") parent-bol 0)
                   ((node-is "{") parent-bol 0)
                   ((node-is "}") parent-bol 0)
                   ((node-is "(") parent-bol 0)
                   ((node-is ")") parent-bol 0)
                   ((parent-is "statement_block") parent-bol 4)
                   ((node-is "statement_block") parent-bol 0)
                   ((parent-is "script_block_expression") parent-bol 4)
                   ((parent-is "script_block") parent-bol 0)
                   ((parent-is "script_parameter") parent-bol 0)
                   ((parent-is "pipeline") parent-bol 4)
                   ((parent-is "command_elements") parent-bol 4)
                   ((parent-is "invokation_expression") parent-bol 4)
                   ((parent-is "argument_list") parent-bol 4)
                   ((parent-is "argument_expression_list") parent 0)
                   ((parent-is "class_statement") parent-bol 4)
                   ((parent-is "class_property_definition") parent-bol 4)
                   ((parent-is "class_method_definition") parent-bol 4)
                   ((parent-is "function_statement") parent-bol 4)
                   ((parent-is "function_parameter_declaration") parent-bol 4)
                   ((node-is "parameter_list") parent-bol 4)
                   ((parent-is "parameter_list") parent 0)
                   ((parent-is "attribute_list") parent-bol 0)
                   ((node-is "param") parent-bol 0)
                   ((parent-is "param_block") parent-bol 4)
                   ((n-p-gp nil "statement_list" "script_block_body") great-grand-parent 0)
                   ((parent-is "statement_list") parent-bol 0)
                   ((node-is "while_condition") parent-bol 4)
                   ((match "else" "if_statement") parent-bol 0)
                   ((match "while" "do_statement") parent-bol 0)
                   ((parent-is ,(rx (or (seq (or "if" "for" "foreach" "while" "do") "_statement") "else_clause" "elseif_clause"))) parent-bol 4)
                   ((parent-is "switch_statement") parent-bol 0)
                   ((parent-is ,(rx (seq (or "switch" "switch_clause") "_condition"))) parent-bol 4)
                   ((parent-is "switch_body") parent-bol 4)
                   ((parent-is "switch_clauses") parent 0)
                   ((match "catch" "try_statement") parent 0)
                   ((match "final" "try_statement") parent 0)
                   ((parent-is "hash_literal_expression") parent-bol 4)
                   ((parent-is "hash_literal_body") parent 0)
                   ((parent-is "array_literal_expression") parent-bol 0)
                   ((parent-is "array_expression") parent-bol 4)
                   ((parent-is "logical_expression") parent 0)
                   ((parent-is "assignment_expression") parent-bol 4)
                   (no-node parent-bol 0))))

    (setq-local comment-start "#")
    (setq-local comment-start-skip "#+ *")
    (setq-local electric-indent-chars (append "{}():;," electric-indent-chars))
    
    (treesit-major-mode-setup)))

(add-to-list 'auto-mode-alist '("\\.ps[dm]?1\\'" . powershell-ts-mode))
(provide 'my-powershell)