From 755c15c0bdda1a91413dccf697b0161a8b8b9458 Mon Sep 17 00:00:00 2001
From: Aaron Pham <contact@aarnphm.xyz>
Date: Wed, 28 Aug 2024 11:52:12 -0400
Subject: [PATCH] chore(clipboard): separate logic to save as base64 and files
 (#329)

Signed-off-by: Aaron Pham <contact@aarnphm.xyz>
---
 lua/avante/clipboard/darwin.lua  | 28 +++++++++++++++++++++++++++-
 lua/avante/clipboard/init.lua    | 14 +++++++++++++-
 lua/avante/clipboard/linux.lua   | 17 ++++++++++++++++-
 lua/avante/clipboard/windows.lua | 18 +++++++++++++++++-
 lua/avante/config.lua            |  4 ++++
 lua/avante/init.lua              |  1 +
 lua/avante/providers/claude.lua  |  2 +-
 lua/avante/providers/openai.lua  | 11 +++++++----
 lua/avante/types.lua             |  3 +++
 9 files changed, 89 insertions(+), 9 deletions(-)

diff --git a/lua/avante/clipboard/darwin.lua b/lua/avante/clipboard/darwin.lua
index c9741d6..0d81e51 100644
--- a/lua/avante/clipboard/darwin.lua
+++ b/lua/avante/clipboard/darwin.lua
@@ -3,8 +3,10 @@ local Utils = require("avante.utils")
 ---@class AvanteClipboard
 local M = {}
 
+---@alias DarwinClipboardCommand "pngpaste" | "osascript"
 M.clip_cmd = nil
 
+---@return DarwinClipboardCommand
 M.get_clip_cmd = function()
   if M.clip_cmd then
     return M.clip_cmd
@@ -34,7 +36,31 @@ M.has_content = function()
   return false
 end
 
-M.get_content = function()
+M.save_content = function(filepath)
+  local cmd = M.get_clip_cmd()
+  ---@type vim.SystemCompleted
+  local output
+
+  if cmd == "pngpaste" then
+    output = Utils.shell_run(('pngpaste - > "%s"'):format(filepath))
+    return output.code == 0
+  elseif cmd == "osascript" then
+    output = Utils.shell_run(
+      string.format(
+        [[osascript -e 'set theFile to (open for access POSIX file "%s" with write permission)' ]]
+          .. [[-e 'try' -e 'write (the clipboard as «class PNGf») to theFile' -e 'end try' ]]
+          .. [[-e 'close access theFile' -e 'do shell script "cat %s > %s"']],
+        filepath,
+        filepath,
+        filepath
+      )
+    )
+    return output.code == 0
+  end
+  return false
+end
+
+M.get_base64_content = function()
   local cmd = M.get_clip_cmd()
   ---@type vim.SystemCompleted
   local output
diff --git a/lua/avante/clipboard/init.lua b/lua/avante/clipboard/init.lua
index bd88881..dae7344 100644
--- a/lua/avante/clipboard/init.lua
+++ b/lua/avante/clipboard/init.lua
@@ -1,16 +1,28 @@
 ---NOTE: this module is inspired by https://github.com/HakonHarnes/img-clip.nvim/tree/main
+---@see https://github.com/ekickx/clipboard-image.nvim/blob/main/lua/clipboard-image/paste.lua
 
+local Path = require("plenary.path")
 local Utils = require("avante.utils")
+local Config = require("avante.config")
 
 ---@class AvanteClipboard
 ---@field clip_cmd string
 ---@field get_clip_cmd fun(): string
 ---@field has_content fun(): boolean
----@field get_content fun(): string
+---@field get_base64_content fun(): string
+---@field save_content fun(filename: string): boolean
 ---
 ---@class avante.Clipboard: AvanteClipboard
 local M = {}
 
+M.paste_directory = Path:new(Config.history.storage_path):joinpath("pasted_images")
+
+M.setup = function()
+  if not M.paste_directory:exists() then
+    M.paste_directory:mkdir({ parent = true })
+  end
+end
+
 return setmetatable(M, {
   __index = function(t, k)
     local os_mapping = Utils.get_os_name()
diff --git a/lua/avante/clipboard/linux.lua b/lua/avante/clipboard/linux.lua
index d65519e..e0cab2a 100644
--- a/lua/avante/clipboard/linux.lua
+++ b/lua/avante/clipboard/linux.lua
@@ -37,7 +37,22 @@ M.has_content = function()
   return false
 end
 
-M.get_content = function()
+M.save_content = function(filepath)
+  local cmd = M.get_clip_cmd()
+  ---@type vim.SystemCompleted
+  local output
+
+  if cmd == "xclip" then
+    output = Utils.shell_run(('xclip -selection clipboard -o -t image/png > "%s"'):format(filepath))
+    return output.code == 0
+  elseif cmd == "wl-paste" then
+    output = Utils.shell_run(('wl-paste --type image/png > "%s"'):format(filepath))
+    return output.code == 0
+  end
+  return false
+end
+
+M.get_base64_content = function()
   local cmd = M.get_clip_cmd()
   ---@type vim.SystemCompleted
   local output
diff --git a/lua/avante/clipboard/windows.lua b/lua/avante/clipboard/windows.lua
index 8ee3b2f..8a233a5 100644
--- a/lua/avante/clipboard/windows.lua
+++ b/lua/avante/clipboard/windows.lua
@@ -30,7 +30,23 @@ M.has_content = function()
   return false
 end
 
-M.get_content = function()
+M.save_content = function(filepath)
+  local cmd = M.get_clip_cmd()
+  ---@type vim.SystemCompleted
+  local output
+
+  if cmd == "powershell.exe" then
+    output = Utils.shell_run(
+      ("Add-Type -AssemblyName System.Windows.Forms; [System.Windows.Forms.Clipboard]::GetImage().Save('%s')"):format(
+        filepath
+      )
+    )
+    return output.code == 0
+  end
+  return false
+end
+
+M.get_base64_content = function()
   local cmd = M.get_clip_cmd()
   ---@type vim.SystemCompleted
   local output
diff --git a/lua/avante/config.lua b/lua/avante/config.lua
index de65b7a..07eeade 100644
--- a/lua/avante/config.lua
+++ b/lua/avante/config.lua
@@ -73,6 +73,10 @@ M.defaults = {
   },
   history = {
     storage_path = vim.fn.stdpath("state") .. "/avante",
+    paste = {
+      extension = "png",
+      filename = "pasted-%Y-%m-%d-%H-%M-%S",
+    },
   },
   highlights = {
     ---@type AvanteConflictHighlights
diff --git a/lua/avante/init.lua b/lua/avante/init.lua
index 3bc6408..c065ca3 100644
--- a/lua/avante/init.lua
+++ b/lua/avante/init.lua
@@ -235,6 +235,7 @@ function M.setup(opts)
   require("avante.highlights").setup()
   require("avante.diff").setup()
   require("avante.providers").setup()
+  require("avante.clipboard").setup()
 
   -- setup helpers
   H.autocmds()
diff --git a/lua/avante/providers/claude.lua b/lua/avante/providers/claude.lua
index be9714a..00a8a6c 100644
--- a/lua/avante/providers/claude.lua
+++ b/lua/avante/providers/claude.lua
@@ -50,7 +50,7 @@ M.parse_message = function(opts)
       source = {
         type = "base64",
         media_type = "image/png",
-        data = Clipboard.get_content(),
+        data = Clipboard.get_base64_content(),
       },
     })
   end
diff --git a/lua/avante/providers/openai.lua b/lua/avante/providers/openai.lua
index 00ac2be..d363b95 100644
--- a/lua/avante/providers/openai.lua
+++ b/lua/avante/providers/openai.lua
@@ -61,18 +61,21 @@ M.get_user_message = function(opts)
 end
 
 M.parse_message = function(opts)
-  local user_content = {}
+  ---@type string | OpenAIMessage[]
+  local user_content
   if Config.behaviour.support_paste_from_clipboard and Clipboard.has_content() then
+    user_content = {}
     table.insert(user_content, {
       type = "image_url",
       image_url = {
-        url = "data:image/png;base64," .. Clipboard.get_content(),
+        url = "data:image/png;base64," .. Clipboard.get_base64_content(),
       },
     })
+    table.insert(user_content, { type = "text", text = M.get_user_message(opts) })
+  else
+    user_content = M.get_user_message(opts)
   end
 
-  table.insert(user_content, { type = "text", text = M.get_user_message(opts) })
-
   return {
     { role = "system", content = opts.system_prompt },
     { role = "user", content = user_content },
diff --git a/lua/avante/types.lua b/lua/avante/types.lua
index bec0604..9e40506 100644
--- a/lua/avante/types.lua
+++ b/lua/avante/types.lua
@@ -16,3 +16,6 @@
 --- @param opts vim.api.keyset.create_autocmd.opts
 --- @return integer
 function vim.api.nvim_create_autocmd(event, opts) end
+
+---@type boolean
+vim.g.avante_login = vim.g.avante_login