dotfiles

local edit = require 'bugabinga.pivi.edit'

local IMAGES = {
  png = 'image/png',
  jpg = 'image/jpeg',
  jpeg = 'image/jpeg',
  gif = 'image/gif',
  webp = 'image/webp',
}

--- What a written prompt turns out to mean.
--- Pi reads a leading bang as a shell command rather than something said, and
--- doubles it to keep the output out of the conversation.
--- @param text string
--- @return table action kind, and what that kind needs
local parse = function( text )
  vim.validate( 'text', text, 'string' )
  local written = vim.trim( text )

  local quiet = written:match '^!!%s*(.*)$'
  if quiet then
    return { kind = 'bash', command = vim.trim( quiet ), exclude = true, }
  end

  local loud = written:match '^!%s*(.*)$'
  if loud then
    return { kind = 'bash', command = vim.trim( loud ), exclude = false, }
  end

  return { kind = 'prompt', message = written, }
end

local image_type = function( path )
  local suffix = path:match '%.([%a%d]+)$'
  if not suffix then return nil end
  return IMAGES[suffix:lower()]
end

local absolute = function( path, cwd )
  local expanded = vim.fn.expand( path )
  if expanded:match '^/' or expanded:match '^%a:' then return vim.fs.normalize( expanded ) end
  return vim.fs.normalize( ( cwd or vim.uv.cwd() ) .. '/' .. expanded )
end

--- Reads a file for the model.
--- A loaded buffer wins over what is on disk, because that is what the user is
--- looking at and what every other pivi read already answers with.
local contents = function( path )
  local text = edit.read( path )
  return text
end

local encode_image = function( path )
  local read = io.open( path, 'rb' )
  if not read then return nil end
  local bytes = read:read '*a'
  read:close()
  if not bytes or bytes == '' then return nil end

  local ok, encoded = pcall( vim.base64.encode, bytes )
  if not ok then return nil end
  return encoded
end

--- Replaces every file reference with the file itself.
--- The shape matches what Pi builds for the same reference elsewhere, so a
--- model sees one convention rather than two.
--- @param text string
--- @param options table? cwd
--- @return string message
--- @return table images
--- @return string[] missing references that could not be read
local expand = function( text, options )
  vim.validate( 'text', text, 'string' )
  options = options or {}

  local images = {}
  local missing = {}
  local attached = {}

  local message = text:gsub( '@([%w%._%-~/\\:]+)', function( reference )
    local path = absolute( reference, options.cwd )

    local mime = image_type( path )
    if mime then
      local data = encode_image( path )
      if not data then
        table.insert( missing, reference )
        return '@' .. reference
      end
      table.insert( images, { type = 'image', data = data, mimeType = mime, } )
      return ( '<file name="%s"></file>' ):format( path )
    end

    local body = contents( path )
    if not body then
      table.insert( missing, reference )
      return '@' .. reference
    end

    if attached[path] then return ( '<file name="%s"></file>' ):format( path ) end
    attached[path] = true
    return ( '<file name="%s">\n%s\n</file>' ):format( path, body )
  end )

  return message, images, missing
end

--- Every reference a written prompt names.
local references = function( text )
  local found = {}
  for reference in ( text or '' ):gmatch '@([%w%._%-~/\\:]+)' do
    table.insert( found, reference )
  end
  return found
end

return {
  image_types = IMAGES,
  parse = parse,
  expand = expand,
  references = references,
  absolute = absolute,
}