(require 'consult)
(require 'eglot)
(require 'jsonrpc)
(require 'subr-x)

(defvar eglot-jdt--symbolkind-map
  '((1 . "File") (2 . "Module") (3 . "Namespace") (4 . "Package")
    (5 . "Class") (6 . "Method") (7 . "Property") (8 . "Field")
    (9 . "Constructor") (10 . "Enum") (11 . "Interface") (12 . "Function")
    (13 . "Variable") (14 . "Constant") (15 . "String") (16 . "Number")
    (17 . "Boolean") (18 . "Array") (19 . "Object") (20 . "Key")
    (21 . "Null") (22 . "EnumMember") (23 . "Struct") (24 . "Event")
    (25 . "Operator") (26 . "TypeParameter"))
  "Mapping of LSP SymbolKind integers to human-readable names.")

(defun eglot-jdt--get (obj key)
  "Robustly retrieve KEY from LSP OBJ (plist/alist)."
  (let* ((kname (substring (symbol-name key) 1))
         (sym-key (intern kname)))
    (or (plist-get obj key)
        (alist-get key obj)
        (alist-get sym-key obj)
        (alist-get kname obj))))

(defun eglot-jdt--fetch-symbols (&optional query)
  "Synchronously fetch workspace symbols for QUERY from Eglot/JDTLS."
  (unless (eglot-managed-p)
    (user-error "Eglot is not managing this buffer"))
  (let* ((server (eglot-current-server))
         (symbols (jsonrpc-request server :workspace/symbol `(:query ,(or query "")))))
    (when (vectorp symbols)
      (setq symbols (append symbols nil)))
    symbols))

;; Helper to pad strings
(defun eglot-jdt--pad-right (str width)
  "Pad STR on the right with spaces to WIDTH."
  (let ((len (length str)))
    (if (< len width)
        (concat str (make-string (- width len) ?\s))
      str)))

;; Main formatting function with faces
(defun eglot-jdt--format-symbols-table (symbols)
  "Format SYMBOLS into a table with faces for Consult, showing Kind first."
  (let* ((rows (mapcar (lambda (sym)
                         (let ((kind-num (eglot-jdt--get sym :kind)))
                           (list (or (alist-get kind-num eglot-jdt--symbolkind-map)
                                     (format "Kind %s" kind-num))
                                 (eglot-jdt--get sym :name)
                                 (or (eglot-jdt--get sym :containerName) "<no-package>")
                                 (let ((location (eglot-jdt--get sym :location)))
                                   (eglot-jdt--get location :uri)))))
                       symbols))
         ;; compute max width for each column
         (max-kind (apply #'max (mapcar (lambda (r) (length (nth 0 r))) rows)))
         (max-name (apply #'max (mapcar (lambda (r) (length (nth 1 r))) rows)))
         (max-package (apply #'max (mapcar (lambda (r) (length (nth 2 r))) rows)))
         ;; faces
         (divider-face 'shadow)       ;; low opacity / dim for dividers
         (kind-face 'font-lock-function-name-face)
         (name-face 'font-lock-variable-name-face)
         (package-face 'font-lock-keyword-face))
    ;; create table rows with faces
    (mapcar (lambda (r)
              (let* ((kind (propertize (eglot-jdt--pad-right (nth 0 r) max-kind) 'face kind-face))
                     (name (propertize (eglot-jdt--pad-right (nth 1 r) max-name) 'face name-face))
                     (package (propertize (eglot-jdt--pad-right (nth 2 r) max-package) 'face package-face))
                     (divider (propertize "|" 'face divider-face)))
                (cons (format "%s %s %s %s %s"
                              kind divider name divider package)
                      (nth 3 r))))
            rows)))

(defun eglot-jdt--open-uri (uri)
  "Open a URI returned by the language server, handling jdt:// URIs."
  (if (and uri (string-prefix-p "jdt://" uri))
      (find-file (expand-file-name uri))  ; your jdt-file-name-handler takes care of this
    (find-file (eglot-uri-to-path uri))))

(defun eglot-jdt--sort-by-query (symbols query)
  "Sort SYMBOLS by how closely their names match QUERY."
  (let ((query (or query "")))
    (sort symbols
          (lambda (a b)
            (< (string-distance (eglot-jdt--get a :name) query)
               (string-distance (eglot-jdt--get b :name) query))))))

(defun consult-eglot-jdt-symbols (&optional query)
  "Consult interface for Java workspace symbols via Eglot/JDTLS.
Sorts results based on closeness to QUERY."
  (interactive "sSymbol query (blank for all): ")
  (let* ((symbols (eglot-jdt--fetch-symbols query))
         (sorted-symbols (eglot-jdt--sort-by-query symbols query))
         (candidates (eglot-jdt--format-symbols-table sorted-symbols))
         (selection (consult--read
                     (mapcar #'car candidates)
                     :prompt "Symbol: "
                     :sort nil
                     :require-match t
                     :category 'symbol)))
    (when selection
      (let ((uri (cdr (assoc selection candidates))))
        (when uri
          (eglot-jdt--open-uri uri))))))