dotfiles

local session = require 'bugabinga.pivi.session'
local help = require 'bugabinga.pivi.help'
local transcript = require 'bugabinga.pivi.transcript'
local context = require 'bugabinga.pivi.context'
local follow = require 'bugabinga.pivi.follow'
local publish = require 'bugabinga.pivi.publish'
local runtime = require 'bugabinga.pivi.runtime'
local comparison = require 'bugabinga.pivi.diff'
local guide = require 'bugabinga.pivi.guide'
local layout = require 'bugabinga.pivi.layout'
local compose = require 'bugabinga.pivi.compose'

local INPUT_FILETYPE = 'pivi-input'

local notify = function( message, level )
  vim.notify( 'pivi: ' .. message, level or vim.log.levels.INFO )
end

local fetch_commands = function( current, callback )
  current.request( { type = 'get_commands', }, function( response )
    local commands = {}
    if response.success and response.data then
      for _, entry in ipairs( response.data.commands or {} ) do
        table.insert( commands, entry )
      end
    end
    current.commands = commands
    if callback then callback( commands ) end
  end )
end

--- Returns the active session, starting one when none exists.
--- A missing or unusable Pi is reported, never raised.
local require_session = function()
  local current = session.active()
  if current then return current end

  local available, failure = runtime.executable()
  if not available then
    notify( failure, vim.log.levels.ERROR )
    return nil
  end

  local ok, created = pcall( session.create, {} )
  if not ok then
    notify( 'cannot start pi: ' .. tostring( created ), vim.log.levels.ERROR )
    return nil
  end

  local started = created.client.failure()
  if started then
    notify( 'cannot start pi: ' .. started, vim.log.levels.ERROR )
    created.stop()
    return nil
  end

  -- Warm the command list so completion has something true to offer.
  fetch_commands( created )
  return created
end

--- Decides whether `/` opens slash-command entry.
--- Only a pivi input buffer with the cursor on the first character qualifies.
--- @param bufnr number
--- @param col number zero based cursor column
local slash_opens_commands = function( bufnr, col )
  if not vim.api.nvim_buf_is_valid( bufnr ) then return false end
  if vim.bo[bufnr].filetype ~= INPUT_FILETYPE then return false end
  return col == 0
end

local command_names = function( current )
  local found = {}
  for _, entry in ipairs( current and current.commands or {} ) do
    table.insert( found, entry.name )
  end
  table.sort( found )
  return found
end

local knows = function( current, name )
  for _, entry in ipairs( current and current.commands or {} ) do
    if entry.name == name then return true end
  end
  return false
end

--- Runs a Pi command.
--- Pi dispatches a command it owns and treats anything else as something the
--- user said, so an unknown name would quietly become a message to the model.
local run_slash_command = function( current, name, arguments )
  if not knows( current, name ) then
    notify(
      ( 'pi has no command "%s" here; terminal-only commands are not offered over this channel' ):format( name ),
      vim.log.levels.WARN
    )
    return false
  end

  local message = '/' .. name
  if arguments and arguments ~= '' then message = message .. ' ' .. arguments end
  local accepted, rejection = current.submit( message, { capture = false, } )
  if not accepted then notify( rejection, vim.log.levels.WARN ) end
  return accepted
end

local pick_slash_command = function( current, on_pick )
  fetch_commands( current, function( commands )
    if #commands == 0 then
      notify( 'no pi commands available', vim.log.levels.WARN )
      return
    end
    -- A picker filters the text it is shown, so showing a description here
    -- means typing a command name competes with every description that
    -- happens to contain the same letters.
    vim.ui.select( commands, {
      prompt = 'pi command',
      format_item = function( entry ) return '/' .. entry.name end,
    }, function( choice )
      if not choice then return end
      on_pick( choice )
    end )
  end )
end

--- Opens a modifiable prompt buffer bound to a session.
local prompts = {}

--- The window showing a session's transcript, when there is one.
--- A session is not always a live one: a caller may hold a stand-in, so the
--- buffer is only trusted once it looks like a buffer.
local transcript_window = function( current )
  if type( current.buffer ) ~= 'number' then return nil end
  if not vim.api.nvim_buf_is_valid( current.buffer ) then return nil end
  local win = vim.fn.bufwinid( current.buffer )
  if win == -1 then return nil end
  return win
end

--- Shows a session's prompt, reusing the one it already has.
--- A draft survives, because reopening a prompt is asking to continue writing,
--- not to start over.
--- What `@` can refer to: what is open first, then what is nearby.
--- A buffer is what the user is actually looking at, so it outranks a path
--- that merely exists.
local references = function( cwd )
  local found, seen = {}, {}

  for _, bufnr in ipairs( vim.api.nvim_list_bufs() ) do
    if vim.api.nvim_buf_is_loaded( bufnr ) and vim.bo[bufnr].buftype == '' then
      local name = vim.api.nvim_buf_get_name( bufnr )
      if name ~= '' then
        local shown = vim.fn.fnamemodify( name, ':.' )
        if not seen[shown] then
          seen[shown] = true
          table.insert( found, { path = shown, open = true, modified = vim.bo[bufnr].modified, } )
        end
      end
    end
  end

  local ok, listed = pcall( vim.fn.systemlist, { 'git', '-C', cwd, 'ls-files', } )
  if ok and vim.v.shell_error == 0 then
    for _, path in ipairs( listed ) do
      if not seen[path] then
        seen[path] = true
        table.insert( found, { path = path, open = false, } )
      end
    end
  end

  return found
end

local pick_reference = function( current, on_pick )
  local choices = references( current.cwd or vim.uv.cwd() )
  if #choices == 0 then
    notify( 'nothing to refer to', vim.log.levels.WARN )
    return
  end

  vim.ui.select( choices, {
    prompt = 'refer to',
    format_item = function( choice ) return choice.path end,
  }, function( choice )
    if choice then on_pick( choice.path ) end
  end )
end

local open_prompt = function( current )
  local existing = prompts[current.name]
  if existing and vim.api.nvim_buf_is_valid( existing ) then
    local win = vim.fn.bufwinid( existing )
    if win == -1 then
      win = layout.open_prompt( existing, transcript_window( current ) ) or -1
    end
    if win ~= -1 then
      vim.api.nvim_set_current_win( win )
      vim.cmd.startinsert()
    end
    return existing
  end

  local bufnr = vim.api.nvim_create_buf( false, true )
  vim.bo[bufnr].buftype = 'nofile'
  vim.bo[bufnr].bufhidden = 'hide'
  vim.bo[bufnr].filetype = INPUT_FILETYPE
  vim.b[bufnr].pivi_session = current.name
  prompts[current.name] = bufnr
  pcall( vim.api.nvim_buf_set_name, bufnr, 'pivi://' .. current.name .. '/prompt' )

  vim.keymap.set( { 'n', 'i', }, '/', function()
    local col = vim.api.nvim_win_get_cursor( 0 )[2]
    if not slash_opens_commands( bufnr, col ) then return '/' end
    vim.schedule( function()
      pick_slash_command( current, function( choice )
        run_slash_command( current, choice.name, nil )
        transcript.show( current.buffer )
      end )
    end )
    return ''
  end, { buffer = bufnr, expr = true, desc = 'pivi slash command entry', } )

  -- `@` names a file, as it does in pi itself. It only takes over when it
  -- starts a word, so an address or a decorator can still be typed.
  vim.keymap.set( 'i', '@', function()
    local line = vim.api.nvim_get_current_line()
    local column = vim.api.nvim_win_get_cursor( 0 )[2]
    local before = column > 0 and line:sub( column, column ) or ''
    if before ~= '' and not before:match '%s' then return '@' end

    vim.schedule( function()
      pick_reference( current, function( path )
        vim.api.nvim_put( { '@' .. path .. ' ', }, 'c', true, true )
        vim.cmd.startinsert { bang = true, }
      end )
    end )
    return ''
  end, { buffer = bufnr, expr = true, desc = 'pivi refer to a file', } )

  guide.bind( bufnr )

  local send = function()
    local lines = vim.api.nvim_buf_get_lines( bufnr, 0, -1, false )
    local text = vim.trim( table.concat( lines, '\n' ) )
    if text == '' then return end
    local accepted, rejection = current.submit( text, { capture = false, } )
    if not accepted then
      notify( rejection, vim.log.levels.WARN )
      return
    end
    vim.api.nvim_buf_set_lines( bufnr, 0, -1, false, { '', } )
    -- A submission the user cannot read the answer to is a dead end.
    transcript.show( current.buffer )
  end

  vim.keymap.set( { 'n', 'i', }, '<Plug>(pivi-submit)', send,
    { buffer = bufnr, desc = 'pivi submit the prompt', } )

  -- Accepting a completion and accepting a prompt are the same intent, so a
  -- key bound to this yields to the popup whenever one is open.
  vim.keymap.set( 'i', '<Plug>(pivi-accept)', function()
    if vim.fn.pumvisible() == 1 then return '<C-y>' end
    vim.schedule( send )
    return ''
  end, { buffer = bufnr, expr = true, desc = 'pivi accept the completion, else the prompt', } )

  local win = layout.open_prompt( bufnr, transcript_window( current ) )
  if win then
    transcript.decorate( win, current.name )
    vim.wo[win].winbar = '%#Title#π ' .. current.name ..
      '%*  %#Comment#' .. ( vim.g.pivi_submit_hint or 'send' ) .. '  / commands%*'
    if vim.g.pivi_start_insert ~= false then vim.cmd.startinsert() end
  end
  return bufnr
end

local submit_text = function( text, options )
  local current = require_session()
  if not current then return false end

  local accepted, rejection = current.submit( text, options )
  if not accepted then
    notify( rejection, vim.log.levels.WARN )
    return false
  end
  transcript.show( current.buffer )
  return true
end

local subcommands = {}

subcommands.ask = function( arguments, modifiers )
  if arguments == '' then
    notify( 'ask needs a question', vim.log.levels.WARN )
    return
  end
  local options = {}
  if modifiers and modifiers.range and modifiers.range > 0 then
    options.snapshot = context.snapshot { range = {
      kind = 'range',
      first = { modifiers.line1, 0, },
      last = { modifiers.line2, 0, },
    }, }
  end
  submit_text( arguments, options )
end

subcommands.steer = function( arguments )
  submit_text( arguments, { delivery = 'steer', } )
end

subcommands['follow-up'] = function( arguments )
  submit_text( arguments, { delivery = 'follow_up', } )
end

subcommands.help = function( arguments )
  if arguments == '' then
    notify( 'help needs a question', vim.log.levels.WARN )
    return
  end
  help.ask( arguments )
end

subcommands.command = function( arguments )
  local current = require_session()
  if not current then return end

  local name, rest = arguments:match '^(%S+)%s*(.*)$'
  if name then
    -- The list may not have arrived yet, so it is fetched before judging a name.
    fetch_commands( current, function()
      if run_slash_command( current, name, rest ) then transcript.show( current.buffer ) end
    end )
    return
  end
  pick_slash_command( current, function( choice )
    run_slash_command( current, choice.name, nil )
    transcript.show( current.buffer )
  end )
end

subcommands.start = function( arguments )
  local available, failure = runtime.executable()
  if not available then
    notify( failure, vim.log.levels.ERROR )
    return
  end

  local created = session.create { name = arguments ~= '' and arguments or nil, }
  fetch_commands( created )
  publish.set_active( created.name )
  transcript.show( created.buffer )
  notify( 'session ' .. created.name )
end

subcommands.stop = function( arguments )
  local target = arguments ~= '' and session.get( arguments ) or session.active()
  if not target then
    notify( 'no session', vim.log.levels.WARN )
    return
  end
  target.stop()
  notify( 'stopped ' .. target.name )
end

subcommands.sessions = function()
  local names = session.list()
  if #names == 0 then
    notify( 'no sessions', vim.log.levels.WARN )
    return
  end
  vim.ui.select( names, { prompt = 'pi session', }, function( choice )
    if not choice then return end
    session.set_active( choice )
    local chosen = session.get( choice )
    if chosen then transcript.show( chosen.buffer ) end
  end )
end

subcommands.show = function()
  local current = session.active()
  if not current then
    notify( 'no session', vim.log.levels.WARN )
    return
  end
  transcript.show( current.buffer )
end

subcommands.prompt = function()
  local current = require_session()
  if not current then return end
  open_prompt( current )
end

subcommands.abort = function()
  local current = session.active()
  if not current then
    notify( 'no session', vim.log.levels.WARN )
    return
  end
  current.abort()
end

subcommands.queue = function()
  local current = session.active()
  if not current then
    notify( 'no session', vim.log.levels.WARN )
    return
  end
  local queue = current.status.queue or {}
  local entries = {}
  for _, item in ipairs( queue.steering or {} ) do table.insert( entries, 'steer: ' .. item ) end
  for _, item in ipairs( queue.follow_up or {} ) do table.insert( entries, 'follow-up: ' .. item ) end
  if #entries == 0 then
    notify 'queue is empty'
    return
  end
  notify( table.concat( entries, '\n' ) )
end

subcommands.clear = function()
  local current = session.active()
  if not current then
    notify( 'no session', vim.log.levels.WARN )
    return
  end
  current.clear_queue( function() notify 'queue cleared' end )
end

subcommands.diff = function()
  local bufnr = vim.api.nvim_get_current_buf()
  if vim.bo[bufnr].filetype ~= transcript.filetype then
    notify( 'move the cursor into a transcript first', vim.log.levels.WARN )
    return
  end

  local changed, _, faithful = comparison.open( bufnr )
  if not changed then
    notify( 'this entry reports no change to compare', vim.log.levels.WARN )
    return
  end
  if faithful == false then
    notify( 'the file has moved on since this change', vim.log.levels.WARN )
  end
end

--- Switches model.
--- `/model` is a terminal-only command, but the capability itself is on this
--- channel, so pivi offers it natively instead of pretending the command works.
subcommands.model = function( arguments )
  local current = require_session()
  if not current then return end

  current.request( { type = 'get_available_models', }, function( response )
    if not response.success or not response.data then
      notify( 'pi did not offer a model list', vim.log.levels.WARN )
      return
    end

    local models = response.data.models or {}
    if #models == 0 then
      notify( 'pi has no models configured', vim.log.levels.WARN )
      return
    end

    local choose = function( model )
      current.request(
        { type = 'set_model', provider = model.provider, modelId = model.id, },
        function( answer )
          if answer.success then
            notify( model.provider .. '/' .. model.id )
          else
            notify( answer.error or 'pi refused the model', vim.log.levels.ERROR )
          end
        end
      )
    end

    if arguments ~= '' then
      for _, model in ipairs( models ) do
        if ( model.provider .. '/' .. model.id ) == arguments or model.id == arguments then
          choose( model )
          return
        end
      end
      notify( 'no model matches ' .. arguments, vim.log.levels.WARN )
      return
    end

    vim.ui.select( models, {
      prompt = 'pi model',
      format_item = function( model ) return model.provider .. '/' .. model.id end,
    }, function( model )
      if model then choose( model ) end
    end )
  end )
end

--- Runs a shell command through pi, so its output joins the conversation.
subcommands.bang = function( arguments )
  local current = require_session()
  if not current then return end
  if arguments == '' then
    notify( 'bang needs a command', vim.log.levels.WARN )
    return
  end

  local accepted, rejection = current.bash( arguments, false )
  if not accepted then
    notify( rejection, vim.log.levels.WARN )
    return
  end
  transcript.show( current.buffer )
end

--- The same, kept out of the conversation.
subcommands['bang!'] = function( arguments )
  local current = require_session()
  if not current then return end
  if arguments == '' then
    notify( 'bang needs a command', vim.log.levels.WARN )
    return
  end

  local accepted, rejection = current.bash( arguments, true )
  if not accepted then
    notify( rejection, vim.log.levels.WARN )
    return
  end
  transcript.show( current.buffer )
end

--- Switches thinking level, the other capability `/thinking` cannot reach here.
subcommands.thinking = function( arguments )
  local current = require_session()
  if not current then return end

  current.request( { type = 'get_available_thinking_levels', }, function( response )
    local levels = response.success and response.data and response.data.levels or {}
    if #levels == 0 then
      notify( 'this model offers no thinking levels', vim.log.levels.WARN )
      return
    end

    local choose = function( level )
      current.request( { type = 'set_thinking_level', level = level, }, function( answer )
        if answer.success then
          notify( 'thinking ' .. level )
        else
          notify( answer.error or 'pi refused the level', vim.log.levels.ERROR )
        end
      end )
    end

    if arguments ~= '' then
      if vim.list_contains( levels, arguments ) then
        choose( arguments )
      else
        notify( 'no such level: ' .. arguments, vim.log.levels.WARN )
      end
      return
    end

    vim.ui.select( levels, { prompt = 'thinking', }, function( level )
      if level then choose( level ) end
    end )
  end )
end

--- Compacts the conversation, optionally saying what to keep.
subcommands.compact = function( arguments )
  local current = require_session()
  if not current then return end

  local request = { type = 'compact', }
  if arguments ~= '' then request.customInstructions = arguments end

  notify 'compacting'
  current.request( request, function( response )
    if not response.success then
      notify( response.error or 'compaction failed', vim.log.levels.ERROR )
      return
    end
    local data = response.data or {}
    notify( ( 'compacted %s to about %s tokens' ):format(
      tostring( data.tokensBefore or '?' ), tostring( data.estimatedTokensAfter or '?' )
    ) )
  end )
end

--- Turns automatic compaction on or off.
subcommands.autocompact = function( arguments )
  local current = require_session()
  if not current then return end

  local wanted = arguments ~= 'off'
  current.request( { type = 'set_auto_compaction', enabled = wanted, }, function( response )
    if response.success then
      notify( 'automatic compaction ' .. ( wanted and 'on' or 'off' ) )
    else
      notify( response.error or 'pi refused', vim.log.levels.ERROR )
    end
  end )
end

--- What this session has cost so far.
subcommands.session = function()
  local current = require_session()
  if not current then return end

  current.request( { type = 'get_session_stats', }, function( response )
    if not response.success or not response.data then
      notify( 'pi did not report session statistics', vim.log.levels.WARN )
      return
    end

    local data = response.data
    local tokens = data.tokens or {}
    local usage = data.contextUsage or {}
    local told = {
      'session ' .. tostring( data.sessionId or current.name ),
      ( 'messages %s, tool calls %s' ):format(
        tostring( data.totalMessages or '?' ), tostring( data.toolCalls or '?' )
      ),
      ( 'tokens %s in, %s out, %s total' ):format(
        tostring( tokens.input or 0 ), tostring( tokens.output or 0 ), tostring( tokens.total or 0 )
      ),
      ( 'cost %s' ):format( tostring( data.cost or 0 ) ),
    }
    if usage.percent then
      table.insert( told, ( 'context %s%% of %s' ):format(
        tostring( usage.percent ), tostring( usage.contextWindow )
      ) )
    end
    notify( table.concat( told, '\n' ) )
  end )
end

--- Names the session, so it can be recognised when resumed.
subcommands.name = function( arguments )
  local current = require_session()
  if not current then return end
  if arguments == '' then
    notify( 'name needs a name', vim.log.levels.WARN )
    return
  end

  current.request( { type = 'set_session_name', name = arguments, }, function( response )
    if response.success then
      notify( 'named ' .. arguments )
    else
      notify( response.error or 'pi refused the name', vim.log.levels.ERROR )
    end
  end )
end

--- Starts a fresh conversation in the process that is already running.
--- `start` spawns another pi; this one does not.
subcommands.new = function()
  local current = require_session()
  if not current then return end

  current.request( { type = 'new_session', }, function( response )
    if not response.success then
      notify( response.error or 'pi refused a new session', vim.log.levels.ERROR )
      return
    end
    if response.data and response.data.cancelled then
      notify( 'an extension cancelled the new session', vim.log.levels.WARN )
      return
    end
    transcript.clear( current.buffer )
    notify 'new conversation'
  end )
end

--- Writes the conversation out as it is.
subcommands.export = function( arguments )
  local current = require_session()
  if not current then return end

  local request = { type = 'export_html', }
  if arguments ~= '' then request.outputPath = vim.fn.expand( arguments ) end

  current.request( request, function( response )
    if response.success and response.data then
      notify( 'exported to ' .. tostring( response.data.path ) )
    else
      notify( response.error or 'export failed', vim.log.levels.ERROR )
    end
  end )
end

subcommands.guide = function()
  guide.show( vim.api.nvim_get_current_buf() )
end

subcommands.follow = function( arguments )
  follow.toggle( arguments ~= '' and arguments or nil )
  notify( 'follow ' .. ( follow.mode() or 'off' ) )
end

local names = function()
  local list = {}
  for key in pairs( subcommands ) do table.insert( list, key ) end
  table.sort( list )
  return list
end

local dispatch = function( options )
  local input = options.args or ''
  local name, rest = input:match '^(%S+)%s*(.*)$'
  -- Bare `:Pi` is the way in: it starts a session when there is none and puts
  -- the cursor where the user can type.
  if not name then
    subcommands.prompt()
    return
  end
  local handler = subcommands[name]
  if not handler then
    notify( 'unknown subcommand: ' .. name, vim.log.levels.ERROR )
    return
  end
  handler( rest or '', options )
end

local setup = function()
  vim.api.nvim_create_user_command( 'Pi', dispatch, {
    nargs = '*',
    range = true,
    desc = 'pivi native Pi integration',
    complete = function( lead, line )
      local matching = function( candidates )
        return vim.tbl_filter( function( candidate )
          return candidate:sub( 1, #lead ) == lead
        end, candidates )
      end

      -- Completing an argument needs the session's own knowledge, so the list
      -- is kept warm and offered as it is, rather than blocking on a request.
      local subcommand = line:match '^%s*Pi%s+(%S+)%s+%S*$'
      if subcommand == 'command' then return matching( command_names( session.active() ) ) end
      if subcommand == 'stop' or subcommand == 'sessions' then return matching( session.list() ) end
      if subcommand == 'follow' then return matching { 'window', 'tab', 'current', } end
      if subcommand == 'autocompact' then return matching { 'on', 'off', } end

      if line:match '^%s*Pi%s+%S*$' == nil then return {} end
      return matching( names() )
    end,
  } )

  vim.keymap.set( 'n', '<Plug>(pivi-show)', function() subcommands.show() end, { desc = 'pivi show transcript', } )
  vim.keymap.set( 'n', '<Plug>(pivi-prompt)', function() subcommands.prompt() end, { desc = 'pivi open prompt', } )
  vim.keymap.set( 'n', '<Plug>(pivi-abort)', function() subcommands.abort() end, { desc = 'pivi abort run', } )
  vim.keymap.set( 'n', '<Plug>(pivi-follow)', function() subcommands.follow '' end, { desc = 'pivi toggle follow', } )
  vim.keymap.set( 'n', '<Plug>(pivi-guide)', function() subcommands.guide() end, { desc = 'pivi show the guide', } )
end

return {
  input_filetype = INPUT_FILETYPE,
  setup = setup,
  dispatch = dispatch,
  subcommands = subcommands,
  names = names,
  slash_opens_commands = slash_opens_commands,
  open_prompt = open_prompt,
}