;; -*- lexical-binding: t -*-
(require 's)
(require 'cl-lib)
(require 'dash)

(progn ;; idiosyncratic code-style
  (put 'treesit-parent-until 'lisp-indent-function 1) 
  (put 'treesit-parent-while 'lisp-indent-function 1) 
  t)

(defvar-local evil-extensible-jump-item/*table*
  ;; (test * do) list
  `(; other
    ))

;; org-mode header-to-header jumping
(defun z03-extjmp/org-extensible-jump-h ()
  (add-to-list
   'evil-extensible-jump-item/*table*
   (let (org-heading-matches (direction t))
     (cons
      (lambda ()
        (when (eq major-mode 'org-mode)
          (let ((m (s-matched-positions-all "^\\*.*$"
                                            (buffer-string))))
            (unless org-heading-matches (setf org-heading-matches m))
            ;; if (point) is within a given header
            (cl-some (lambda (r) (<= (car r) (point) (cdr r))) m))))
      (lambda (&optional count)
        (cond
         (direction
          (let ((r (seq-elt org-heading-matches
                            (1- (length org-heading-matches)))))
            (if (<= (car r) (point) (cdr r))
                (setf direction nil))))
         (t
          (let ((r (seq-elt org-heading-matches 0)))
            (if (<= (car r) (point) (cdr r))
                (setf direction t)))))
        (unless count (setf count 1))
        (if direction
            (org-next-visible-heading count)
          (org-previous-visible-heading count)))))))
(add-hook 'org-mode-hook #'z03-extjmp/org-extensible-jump-h)

;;; HTML
(defun z03-extjmp/html-ts-test ()
  (cl-block nil
    (unless (featurep 'treesit)
      (cl-return nil))
    (when (or (eq major-mode 'html-ts-mode)
              (eq major-mode 'xml-ts-mode))
      (let ((ty (treesit-node-type (treesit-node-at (point)))))
        (cl-return
         (or (string= ty "tag_name") (string= ty "<")
             (string= ty "</") (string= ty ">")
             (comment
              ;; for if or when i support this
              (string= ty "=") (string= ty (string ?\"))
              (string= ty "attribute_name")
              (string= ty "attribute_value")
              )
             ;; i think these are neccessarily parent types
             ;; and not leaf-node types
             (string= ty "start_tag")
             (string= ty "end_tag")))))))

(defun z03-extjmp/html-ts-go (&optional count)
  (let* ((n (treesit-node-at (point)))
         (ty (treesit-node-type n)))
    ;; only works if point is on tag or brackets
    (let* ((tag (treesit-node-parent n))
           (elm (treesit-node-parent tag))

           (children (treesit-node-children elm))
           (other-tag
            (if (string= (treesit-node-type tag) "start_tag")
                (seq-elt children (1- (seq-length children)))
              (seq-elt children 0))))
      (goto-char (treesit-node-start other-tag)))))
(defun z03-extjmp/html-ts-extensible-jump-h ()
  (add-to-list
   'evil-extensible-jump-item/*table*
   ;; HTML start-tag to end-tag jumping
   (cons
    'z03-extjmp/html-ts-test
    'z03-extjmp/html-ts-go)))

(add-hook 'html-ts-mode-hook #'z03-extjmp/html-ts-extensible-jump-h)

;;; BASH
(defun z03-extjmp/bash-ts-if-test ()
  (cl-block nil
    (unless (featurep 'treesit)
      (cl-return 1))
    (when (or (eq major-mode 'bash-ts-mode)
              (eq major-mode 'zsh-ts-mode))
      (let ((ty (treesit-node-type (treesit-node-at (point)))))
        (cl-return (or (string= ty "if")
                       (string= ty "fi")
                       (string= ty "else")
                       (string= ty "elif")))))))

(defun z03-extjmp/bash-ts-if-go (&optional count)
  (let* ((n (treesit-node-at (point)))
         (ty (treesit-node-type n))
         (p (treesit-node-parent n))
         (gp (treesit-node-parent p))
         (c (treesit-node-children p))
         (other
          (cond ((string= ty "if")
                 'if-next)
                ((string= ty "fi")
                 'fi)
                (t 'else-next)))
         (else-parent-p 
          (or (string= (treesit-node-type p) "else_clause")
              (string= (treesit-node-type p) "elif_clause")))
         (elses
          (cl-remove-if-not
           (lambda (x)
             (let* ((ty (treesit-node-type x)))
               (and (or (string= ty "else_clause")
                        (string= ty "elif_clause"))
                    ;; because final ELSE has ELSES otherwise and shit
                    ;; breaks
                    (<= (treesit-node-end n) (treesit-node-start x)))))
           (if else-parent-p
               (treesit-node-children gp)
             c))))
    (cl-block nil
      ;; go to fi if there are no elses
      (when (and (eq other 'if-next) (length= elses 0))
        (setf other (seq-elt c (1- (seq-length c))))
        (cl-return))
      ;; go to beginning
      (when (eq other 'fi)
        (setf other (seq-elt c 0))
        (cl-return))
      ;; if its the final else
      (when (and (length= elses 0) (eq other 'else-next) else-parent-p)
        (let ((c (treesit-node-children gp)))
          (setf other (seq-elt c (1- (seq-length c)))))
        (cl-return))
      ;; i.e. else-next
      (setf other
            (car (cl-remove-if-not
                  (lambda (x)
                    (<= (treesit-node-end n) (treesit-node-start x)))
                  elses)))
      t)
    (goto-char (treesit-node-start other))))

(let (case-stmt case-item)
  (defun z03-extjmp/bash-ts-case-test ()
    (setf case-stmt nil case-item nil)
    (let* ((curn (treesit-node-at (point)))
           (case-item-n
            (treesit-parent-until curn
              (lambda (n)
                (treesit-node-type= n "case_item"))))
           (case-stmt-n
            (treesit-parent-until case-item-n
              (lambda (n)
                (treesit-node-type= n "case_statement")))))
      (and
       ;; but this is to prevent code inside a `case_item' also
       ;; applying to this, so we test here
       (treesit-node-type= curn "extglob_pattern")
       ;; also, this also applies for the | version, since most likely
       ;; you are on a pattern, and can move if needed
       case-item-n
       case-stmt-n
       (setf case-item case-item-n)
       (setf case-stmt case-stmt-n))))
  (defun z03-extjmp/bash-ts-case-go (&optional count)
    ;;; XXX:
    ;; maybe jump from esac to case/in and back?
    ;; that would require a separate function and a dispatch
    (let* ((sibs
            (treesit-node-children case-stmt))
           ;; case_item siblings,, 
           (cisibs
            (-filter (lambda (n)  (treesit-node-type= n "case_item"))
                     sibs))
           (next-stmt
            (cadr
             (cl-member case-item cisibs
                        :test #'treesit-node-eq))))
      (if next-stmt
          (goto-char (treesit-node-start next-stmt))
        (let ((first-node
               (car (-filter
                     (lambda (n)
                       (treesit-node-type= n "case_item"))
                     cisibs))))
          (message "first-node:=%s" first-node)
          (goto-char (treesit-node-start first-node))))))
  t)

(defun z03-extjmp/bash-ts-generic-block-test ())
(defun z03-extjmp/bash-ts-generic-block-go ())

(let (pipeline)
  (defun z03-extjmp/bash-ts-pipeline-test ()
    (when (or (eq major-mode 'bash-ts-mode)
              (eq major-mode 'zsh-ts-mode))
      (let* ((node (treesit-node-at (point)))
             (forward-pipe
              (condition-case err
                  (save-excursion
                    (re-search-forward "|" (point-max))
                    (treesit-node-at (point)))
                (error nil)))
             (backward-pipe
              (condition-case err
                  (save-excursion
                    (re-search-backward "|" (point-min))
                    (treesit-node-at (point)))
                (error nil)))
             (forward-pipeline
              (-some-> forward-pipe (treesit-node-parent)))
             (backward-pipeline
              (-some-> backward-pipe (treesit-node-parent)))
             (forward-enclosedp
              (-some->> forward-pipeline (treesit-node-enclosed-p node)))
             (backward-enclosedp
              (-some->> backward-pipeline (treesit-node-enclosed-p node)))
             ;; we use `*-pipelinep's here because otherwise we get
             ;; PROGRAM nodes in here, which obviously enclose
             ;; every other node.

             ;; the reason it gives us the PROGRAM node as a parent to
             ;; the node with | in it in `open-cuis-image-zenity' is
             ;; because there is a pipe symbol within the comment
             ;; block at the beginning
             (forward-pipeline-pipelinep
              (and forward-pipeline
                   (treesit-node-p forward-pipeline)
                   (string= "pipeline" (treesit-node-type forward-pipeline))))
             (backward-pipeline-pipelinep
              (and backward-pipeline
                   (treesit-node-p backward-pipeline)
                   (string= "pipeline" (treesit-node-type backward-pipeline)))))
        (cond ((and forward-pipeline forward-pipeline-pipelinep
                    forward-enclosedp)
               (setf pipeline forward-pipeline))
              ((and backward-pipeline backward-pipeline-pipelinep
                    backward-enclosedp)
               (setf pipeline backward-pipeline))
              (t nil))
        (or (and forward-pipeline-pipelinep forward-enclosedp)
            (and backward-pipeline-pipelinep backward-enclosedp)))))
  (defun z03-extjmp/bash-ts-pipeline-go (&optional count)
    (let* ((n (treesit-node-at (point)))
           (pipe-children (treesit-node-children pipeline))
           (pipe-end (seq-elt pipe-children (1- (length pipe-children))))
           (endp
            (treesit-node-enclosed-p n pipe-end 'partial)))
      (if endp
          (goto-char (treesit-node-start (seq-elt pipe-children 0)))
        (let* ((next
                (car 
                 (cl-remove-if-not 
                  (lambda (x)
                    (and
                     (not (string= (treesit-node-type x) "|"))
                     (<= (treesit-node-end n) (treesit-node-start x))))
                  pipe-children))))
          (goto-char (treesit-node-start next))))))
  'defined-pipeline-jumps)

(defun z03-extjmp/bash-ts-extensible-jump-h ()
  (add-to-list
   'evil-extensible-jump-item/*table*
   (cons
    'z03-extjmp/bash-ts-if-test
    'z03-extjmp/bash-ts-if-go))

  ;; TODO:
  ;;  - support for jumping between block start and end 
  (add-to-list
   'evil-extensible-jump-item/*table*
   (cons
    'z03-extjmp/bash-ts-generic-block-test
    'z03-extjmp/bash-ts-generic-block-go))
   
  (add-to-list
   'evil-extensible-jump-item/*table*
   (cons
    'z03-extjmp/bash-ts-pipeline-test
    'z03-extjmp/bash-ts-pipeline-go))

  ;; make higher priority so generic block isnt pranged
  ;; TODO:
  ;;  - support for jumping between CASE clauses
  (add-to-list
   'evil-extensible-jump-item/*table*
   (cons
    'z03-extjmp/bash-ts-case-test
    'z03-extjmp/bash-ts-case-go))
  t)

(add-hook 'bash-ts-mode-hook #'z03-extjmp/bash-ts-extensible-jump-h)
(add-hook 'zsh-ts-mode-hook #'z03-extjmp/bash-ts-extensible-jump-h)


;;; SML

(defsubst treesit-node-type= (node typespec)
  (string= (treesit-node-type node) typespec))

(let (fmrule)
  (defun z03-extjmp/sml-patfun-test ()
    (let* ((cur (treesit-node-at (point)))
           (patnode 
            (treesit-parent-until cur
              (lambda (n) (treesit-node-type= n "vid_pat"))))
           ;; this is so that fmrules from a case exp do not affect
           ;; this, and return true
           (caseexp-patnode
            (treesit-parent-until patnode 
              (lambda (n) (treesit-node-type= n "case_exp"))))
           (caseexp-curn
            (treesit-parent-until cur 
              (lambda (n) (treesit-node-type= n "case_exp"))))
           ;; --
           (fmruleprt-of-cur
            (treesit-parent-until cur 
              (lambda (n) (treesit-node-type= n "fmrule"))))
           (fmruleprt-def
            (treesit-parent-until (treesit-node-at (point))
              (lambda (n)
                (string= "def" (treesit-node-field-name n)))))
           (fmruleprt-of-patnode
            (treesit-parent-until patnode
              (lambda (n) (treesit-node-type= n "fmrule")))))
      ;;(print (alist! caseexp-curn caseexp-patnode
      ;;               fmruleprt-of-patnode fmruleprt-of-cur
      ;;               patnode))
      (and (or (and patnode fmruleprt-of-patnode)
               (and fmruleprt-of-cur (not fmruleprt-def)))
           (not (or caseexp-patnode caseexp-curn))
           (setf fmrule
                 (or fmruleprt-of-patnode fmruleprt-of-cur)))))

  (defun z03-extjmp/sml-patfun-go (&optional count)
    (let* ((prt (treesit-node-parent fmrule))
           (sibs (treesit-node-children prt))
           (fmrulenxt
            (caddr
             (cl-member fmrule sibs
                        :test #'treesit-node-eq))))
      (if fmrulenxt
          (goto-char (treesit-node-start fmrulenxt))
        (let ((first-node
               (cl-find-if (lambda (n) (treesit-node-type= n "fmrule"))
                           sibs)))
          (goto-char (treesit-node-start first-node))))))
  t)

(let (casenode)
  (defun z03-extjmp/sml-case-test ()
    (let* ((cur (treesit-node-at (point)))
           (patnode 
            (treesit-parent-until cur
              (lambda (n) (or (treesit-node-type= n "vid_pat")
                         (treesit-node-type= n "scon_pat")))))
           (parent-casenode
            (treesit-parent-until patnode
              (lambda (n) (treesit-node-type= n "case_exp")))))
      (and patnode
           parent-casenode
           (setf casenode parent-casenode))))

  (defun z03-extjmp/sml-case-go (&optional count)
    (let* ((curn (treesit-node-at (point)))
           (mrulen (treesit-parent-until curn
                     (lambda (n) (treesit-node-type= n "mrule"))))
           (mrulenxt 
            (caddr 
             (cl-member mrulen (treesit-node-children casenode)
                        :test #'treesit-node-eq))))
      (if mrulenxt
          (goto-char (treesit-node-start mrulenxt))
        (let ((first-node
               (cl-find-if (lambda (n) (treesit-node-type= n "mrule"))
                           (treesit-node-children casenode))))
          (goto-char (treesit-node-start first-node))))))
  t)

(let (what)
  (defun z03-extjmp/sml-let-test ()
    (let* ((curn (treesit-node-at (point))) 
           (test-on-let
            (string= (treesit-node-type curn)
                     "let"))
           (test-on-in
            (string= (treesit-node-type curn)
                     "in"))
           (test-on-end
            (string= (treesit-node-type curn)
                     "end")))
      (cl-macrolet ((or-setf-what (&rest clauses)
                      (let ((c (cl-map 'list
                                 (lambda (x)
                                   (list 'and x `(setf what (quote ,x)) x))
                                 clauses)))
                        `(or ,@c))))
        (or-setf-what
         ;; test-in-val
         test-on-end
         test-on-in
         test-on-let))))
  (comment
   ;; jumping from vals is kinda weird
   (test-in-val
    (-some-> (treesit-parent-until
              curn 
              (lambda (n)
                (string= "val_dec"
                         (treesit-node-type n))))
      (treesit-node-parent)
      ;; not changing this, shrimply prevent the bytecompiler
      ;; warnings
      ((lambda (n)
         (string= (treesit-node-type n)
                  "let_exp")))))
   ;; may be neccessary to get to the body from inside app_exp 
   (treesit-parent-until
    curn
    (lambda (n)
      (string= (treesit-node-field-name n)
               "body")))
   ;; ---
   t)
  (defun z03-extjmp/sml-let-go (&optional count)
    (let* ((curn (treesit-node-at (point)))
           (sibs (treesit-node-children (treesit-node-parent curn))))
      (cl-case what
        (test-on-end
         (let* ((let (car sibs))
                (npt (treesit-node-start let)))
           (goto-char npt)))
        (test-on-let
         (let* ((let (car
                      (-filter (lambda (n) (string= (treesit-node-type n) "in"))
                              sibs)))
                (npt (treesit-node-start let)))
           (goto-char npt))
         'test-on-let)
        (test-on-in
         (let* ((let (car (last sibs)))
                (npt (treesit-node-start let)))
           (goto-char npt))))))
       
  t)
(defun z03-extjmp/sml-ts-extensible-jump-item-h ()
  ;; dubious 
  ;;(add-to-list
  ;; 'evil-extensible-jump-item/*table*
  ;; (cons 'z03-extjmp/sml-if-test
  ;;       'z03-extjmp/sml-if-go))
  
  ;; has to be last since it clobbers other movers.
  ;; in particular, it will bind to a function's FMRULE even if its on
  ;; a `let` node, so this will try to go to the next FMRULE even if
  ;; its a stupid movement
  (add-to-list
   'evil-extensible-jump-item/*table*
   (cons 'z03-extjmp/sml-patfun-test
         'z03-extjmp/sml-patfun-go))

  (add-to-list
   'evil-extensible-jump-item/*table*
   (cons 'z03-extjmp/sml-let-test
         'z03-extjmp/sml-let-go))
               
  (add-to-list
   'evil-extensible-jump-item/*table*
   (cons 'z03-extjmp/sml-case-test
         'z03-extjmp/sml-case-go))
  t)

(add-hook 'sml-ts-mode-hook
          #'z03-extjmp/sml-ts-extensible-jump-item-h)

(comment
 "For those using normal emacs."
 ;; NOTE: this is very similar to `sgml-skip-tag-backward' et al, but
 ;;       it functions closer to how vim does it with `%', so i am
 ;;       keeping it
 (defun treesit-html-goto-opposite-tag ()
   (interactive)
   (when (z03-extjmp/html-ts-test)
     (z03-extjmp/html-ts-go)))

 ;; C-c C-j is html-line
 (define-key html-ts-mode-map (kbd "C-c j") #'treesit-html-goto-opposite-tag)
 t)

(defun evil-extensible-jump-item/advice (fn &optional count)
  "Advice for enabling the extensible jumping behaviour."
  (let (do)
    (if (cl-some (lambda (x)
                   (let ((tst (funcall (car x))))
                     (when tst
                       (setf do (cdr x))
                       tst)))
                 evil-extensible-jump-item/*table*)
        (funcall do count)
      ;; fallthrough to original function
      (apply fn count))))

(advice-add 'evil-jump-item :around
  #'evil-extensible-jump-item/advice
  '((name . "evil-extensible-jump-item")))

(comment
 (advice-remove 'evil-jump-item "evil-extensible-jump-item")
 t)

(comment 
 (with-current-buffer "lib.zsh"
   (mapcar (lambda (x)
             (->> (symbol-name (car x))
                  (string-replace
                   "z03-extjmp/bash-ts-" "")
                  (string-replace "-test" "")))
           evil-extensible-jump-item/*table*))
 =>
 ("case-test" "pipeline-test" "generic-block-test" "if-test")
 t)

(defun z03-org/within-any-block-p ()
  (org-in-block-p
   (append '("quote" "center" "verse" "comment")
           org-protecting-blocks)))

