From b11abea019175e09f8e004107cfba260a1b98984 Mon Sep 17 00:00:00 2001
From: Will Hopkins <willothyh@gmail.com>
Date: Wed, 20 Dec 2023 16:14:42 -0800
Subject: [PATCH] feat: highlight search matches and search result score

---
 lua/hawtkeys/ui.lua | 80 +++++++++++++++++++++++++++------------------
 1 file changed, 49 insertions(+), 31 deletions(-)

diff --git a/lua/hawtkeys/ui.lua b/lua/hawtkeys/ui.lua
index 4429ee6..6cd9269 100644
--- a/lua/hawtkeys/ui.lua
+++ b/lua/hawtkeys/ui.lua
@@ -45,29 +45,6 @@ local function create_win(enter, opts)
     return win, buf
 end
 
----@param str string
----@param combo string
----@return string
-local function highlight_desc(str, combo)
-    -- returns str with the first unmarked occurrence of each letter of combo surrounded by []
-    local newStr = str:lower()
-    local marked = {} -- Keep track of characters already marked
-    for i = 1, #combo do
-        local char = combo:sub(i, i)
-        local pos = marked[char] or 1 -- Start searching from the last marked position or from the beginning
-        pos = newStr:find(char, pos, true) or 0
-        if pos then
-            newStr = newStr:sub(1, pos - 1)
-                .. "["
-                .. char
-                .. "]"
-                .. newStr:sub(pos + 1)
-            marked[char] = pos + 2 -- Mark this character's position
-        end
-    end
-    return newStr
-end
-
 local prompt_extmark
 
 M.search = function(text)
@@ -80,15 +57,14 @@ M.search = function(text)
     for i = 1, #results do
         local data = results[i]
         local lines = {}
-        table.insert(
-            lines,
-            "Key: "
-                .. highlight_desc(text, data.combo)
-                .. "<leader>"
-                .. data.combo
-                .. " - Hawt Score: "
-                .. data.score
+        local line = string.format(
+            "Key: %s <leader>%s - Hawt Score: %d",
+            text,
+            data.combo,
+            data.score
         )
+        table.insert(lines, line)
+
         line_count = line_count + 1
         local already_mapped = false
         if
@@ -138,6 +114,48 @@ M.search = function(text)
                 0,
                 -1
             )
+        else
+            local newStr = text:lower()
+            local marked = {} -- Keep track of characters already marked
+            for idx = 1, #data.combo do
+                local char = data.combo:sub(idx, idx)
+                local pos = marked[char] or 1 -- Start searching from the last marked position or from the beginning
+                pos = newStr:find(char, pos, true) or 0
+                if marked[char] and marked[char] == pos then
+                    pos = newStr:find(char, pos + 1, true) or 0
+                end
+                if pos then
+                    newStr = newStr:sub(1, pos - 1)
+                        .. char
+                        .. newStr:sub(pos + 1)
+                    local hl_col = pos + 5
+                    vim.api.nvim_buf_add_highlight(
+                        ResultBuf,
+                        -1,
+                        "Function",
+                        line_count - (already_mapped and 3 or 1),
+                        hl_col - 1,
+                        hl_col
+                    )
+                    marked[char] = pos -- Mark this character's position
+                end
+            end
+
+            local score_offset = #line - #tostring(data.score)
+            local score_hl = "MoreMsg"
+            if data.score <= 0 then
+                score_hl = "ErrorMsg"
+            elseif data.score <= 2 then
+                score_hl = "WarningMsg"
+            end
+            vim.api.nvim_buf_add_highlight(
+                ResultBuf,
+                -1,
+                score_hl,
+                line_count - (already_mapped and 3 or 1),
+                score_offset,
+                score_offset + #tostring(data.score)
+            )
         end
     end
 end