dotfiles

local rpc = require 'bugabinga.pivi.rpc'
local status = require 'bugabinga.pivi.status'
local publish = require 'bugabinga.pivi.publish'
local transcript = require 'bugabinga.pivi.transcript'
local assemble = require 'bugabinga.pivi.assemble'
local context = require 'bugabinga.pivi.context'
local dialog = require 'bugabinga.pivi.dialog'
local partial = require 'bugabinga.pivi.partial'
local activity = require 'bugabinga.pivi.activity'
local follow = require 'bugabinga.pivi.follow'
local edit = require 'bugabinga.pivi.edit'
local runtime = require 'bugabinga.pivi.runtime'
local compose = require 'bugabinga.pivi.compose'
local entry = require 'bugabinga.pivi.entry'
local render = require 'bugabinga.pivi.render'
local fold = require 'bugabinga.pivi.fold'
local tool = require 'bugabinga.pivi.tool'

local registry = {}
local counter = 0

local model = function( requested )
  return requested or vim.g.pivi_model
end

local build_command = function( options )
  options = options or {}
  local command = { options.executable or runtime.executable() or 'pi', '--mode', 'rpc', }

  local extension = options.extension ~= false and runtime.extension_path() or nil
  if extension then
    table.insert( command, '-e' )
    table.insert( command, extension )
  end

  local skill = options.skill ~= false and runtime.skill_path() or nil
  if skill then
    table.insert( command, '--skill' )
    table.insert( command, skill )
  end
  local chosen = model( options.model )
  if chosen then
    table.insert( command, '--model' )
    table.insert( command, chosen )
  end

  local thinking = options.thinking or vim.g.pivi_thinking
  if thinking then
    table.insert( command, '--thinking' )
    table.insert( command, thinking )
  end
  if options.session then
    table.insert( command, '--session' )
    table.insert( command, options.session )
  end
  if options.name then
    table.insert( command, '--name' )
    table.insert( command, options.name )
  end
  for _, argument in ipairs( options.arguments or {} ) do
    table.insert( command, argument )
  end
  return command
end

local unique_name = function( cwd )
  counter = counter + 1
  local base = vim.fs.basename( cwd or vim.uv.cwd() or 'pi' )
  local name = base
  while registry[name] do
    name = base .. '-' .. tostring( counter )
    counter = counter + 1
  end
  return name
end


--- Ends the entry currently being written, if any.
local seal = function( self )
  if not self.entry then return end
  entry.close( self.buffer, self.entry )
  render.decorate( self.buffer, self.entry )
  fold.apply( self.buffer )
  self.entry = nil
  self.entry_kind = nil
end

local begin = function( self, kind, meta )
  seal( self )
  self.entry = entry.open( self.buffer, kind, meta )
  self.entry_kind = kind
  render.decorate( self.buffer, self.entry )
  return self.entry
end

--- Records a submission as its own entry.
--- The submitted text is written exactly as typed; the quoting a reader sees
--- is decoration.
local render_submission = function( self, text )
  begin( self, 'prompt', { text = text, } )
  for _, piece in ipairs( vim.split( text, '\n', { plain = true, } ) ) do
    entry.line( self.buffer, self.entry, piece )
  end
  seal( self )
end

local render_event = function( self, event )
  local bufnr = self.buffer

  if event.type == 'message_start' and event.message and event.message.role == 'assistant' then
    seal( self )
  elseif event.type == 'message_update' then
    local delta = event.assistantMessageEvent or {}
    if delta.type == 'text_delta' and ( delta.delta or '' ) ~= '' then
      local text = delta.delta
      -- Models often open with blank lines; the entry already separates turns.
      if self.entry_kind ~= 'answer' then text = text:gsub( '^%s*\n', '' ) end
      if text ~= '' then
        if self.entry_kind ~= 'answer' then begin( self, 'answer', {} ) end
        entry.write( bufnr, self.entry, text )
      end
    end
  elseif event.type == 'tool_execution_start' then
    self.streamed = nil
    begin( self, 'tool', {
      tool = event.toolName,
      args = event.args,
      call = event.toolCallId,
      cwd = self.cwd,
    } )
    entry.line( bufnr, self.entry, tool.headline( entry.meta( bufnr, self.entry ) ) )
  elseif event.type == 'tool_execution_update' then
    -- Output arrives while the tool still runs, so the entry shows it now
    -- rather than only once the tool ends.
    if self.entry_kind == 'tool' and event.partialResult then
      entry.annotate( bufnr, self.entry, { result = event.partialResult, } )
      local shown = tool.body( entry.meta( bufnr, self.entry ) )
      local joined = table.concat( shown, '\n' )
      if joined ~= self.streamed then
        self.streamed = joined
        entry.rewrite( bufnr, self.entry, shown )
        render.decorate( bufnr, self.entry, { head_only = true, } )
      end
    end
  elseif event.type == 'tool_execution_end' then
    if self.entry_kind == 'tool' then
      entry.annotate( bufnr, self.entry, { result = event.result, failed = event.isError == true, } )
      entry.rewrite( bufnr, self.entry, tool.body( entry.meta( bufnr, self.entry ) ) )
      self.streamed = nil
      seal( self )
    end
  elseif event.type == 'agent_settled' then
    seal( self )
  elseif event.type == 'message_end' and event.message and event.message.errorMessage then
    begin( self, 'error', { error = event.message.errorMessage, } )
    entry.line( bufnr, self.entry, tostring( event.message.errorMessage ) )
    seal( self )
  end
end

local publish_status = function( self )
  publish.session_status( self.name, self.status )
end

local located = function( arguments )
  if type( arguments ) ~= 'table' then return nil end
  local path = arguments.path or arguments.file_path or arguments.filePath
  if type( path ) ~= 'string' or path == '' then return nil end
  return { path = path, line = arguments.line or arguments.lnum, }
end

--- Renders streaming edit intent as inert ghost text.
--- Nothing here mutates a buffer: complete arguments are required for that.
local ghost_intent = function( self, call )
  local intent = partial.intent( call.name, call.arguments_text )
  if not intent.path or not intent.text then return end
  local bufnr = edit.buffer_for( intent.path )
  if not bufnr then return end
  self.ghost_buffer = bufnr
  activity.ghost( bufnr, 0, intent.text )
end

local clear_ghost = function( self )
  if self.ghost_buffer then
    activity.ghost_clear( self.ghost_buffer )
    self.ghost_buffer = nil
  end
end

local handle_event = function( self, event )
  if event.type == 'extension_ui_request' then
    dialog.handle( event, function( response ) self.client.notify( response ) end )
    return
  end

  self.status = status.reduce( self.status, event )
  self.message = assemble.apply( self.message, event )

  if event.type == 'message_update' then
    local delta = event.assistantMessageEvent or {}
    if delta.type == 'toolcall_delta' then
      for _, call in ipairs( assemble.tool_calls( self.message ) ) do
        if not call.complete then ghost_intent( self, call ) end
      end
    elseif delta.type == 'toolcall_end' then
      clear_ghost( self )
    end
  elseif event.type == 'tool_execution_start' then
    clear_ghost( self )
    local location = located( event.args )
    if location then follow.to( location ) end
  elseif event.type == 'agent_settled' then
    clear_ghost( self )
  end

  render_event( self, event )
  for _, listener in ipairs( self.listeners ) do
    pcall( listener, event, self )
  end
  publish_status( self )
end

--- Creates a session and starts its Pi process.
--- @param options table? name, cwd, model, session, arguments, spawn
local new = function( options )
  options = options or {}
  local cwd = options.cwd or vim.uv.cwd()
  local name = options.name or unique_name( cwd )

  local self = {
    name = name,
    cwd = cwd,
    status = status.initial(),
    message = assemble.new(),
    listeners = {},
    buffer = transcript.ensure( name ),
  }

  self.client = rpc.start {
    command = build_command {
      model = options.model,
      session = options.session,
      name = options.name,
      arguments = options.arguments,
      extension = options.extension,
      skill = options.skill,
      thinking = options.thinking,
      executable = options.executable,
    },
    cwd = cwd,
    env = options.env or runtime.bridge_environment(),
    spawn = options.spawn,
    on_event = function( event ) handle_event( self, event ) end,
    on_stderr = function( data ) self.last_stderr = data end,
    on_failure = function( failure )
      vim.notify( 'pivi: ' .. failure, vim.log.levels.ERROR )
    end,
    on_exit = function( code )
      self.status = status.stopped( self.status, code )
      publish_status( self )
    end,
  }

  self.status.connection = 'running'
  publish_status( self )

  registry[name] = self
  return self
end

local get = function( name ) return registry[name] end

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

local active = function()
  local value = publish.read()
  if value.active and registry[value.active] then return registry[value.active] end
  local names = list()
  if #names == 0 then return nil end
  return registry[names[1]]
end

local set_active = function( name )
  if not registry[name] then return false end
  publish.set_active( name )
  return true
end

--- Runs a shell command through Pi.
--- Pi keeps the output in the conversation unless asked not to, which is what
--- doubling the bang means.
local run_bash = function( self, command, exclude, callback )
  vim.validate( 'command', command, 'string' )
  if command == '' then return false, 'nothing to run' end

  local id = begin( self, 'tool', {
    tool = exclude and 'bang!' or 'bang',
    args = { command = command, },
    cwd = self.cwd,
  } )
  entry.line( self.buffer, id, tool.headline( entry.meta( self.buffer, id ) ) )

  self.client.request(
    { type = 'bash', command = command, excludeFromContext = exclude == true, },
    function( response )
      local data = response.data or {}
      entry.annotate( self.buffer, id, {
        result = { content = { { type = 'text', text = data.output or '', }, }, },
        failed = response.success ~= true or ( data.exitCode ~= nil and data.exitCode ~= 0 ),
        exit = data.exitCode,
      } )
      entry.rewrite( self.buffer, id, tool.body( entry.meta( self.buffer, id ) ) )
      seal( self )
      if callback then callback( response ) end
    end
  )

  return true, nil
end

--- Submits text to the session.
--- A submission during an active run must choose its delivery explicitly;
--- rejection is returned, never silently swallowed.
--- @return boolean accepted
--- @return string? rejection
local submit = function( self, text, options )
  options = options or {}
  vim.validate( 'text', text, 'string' )

  if not self.client.is_running() then
    return false, 'pi process is not running'
  end

  local busy = status.is_busy( self.status )
  local delivery = options.delivery

  if busy and not delivery then
    return false, 'pi is running; use :Pi steer <text>, :Pi follow-up <text>, or :Pi abort'
  end

  if delivery and delivery ~= 'steer' and delivery ~= 'follow_up' then
    return false, 'unknown delivery "' .. tostring( delivery ) .. '"; use steer or follow_up'
  end

  -- A written prompt may not be a prompt at all.
  if options.compose ~= false then
    local intent = compose.parse( text )
    if intent.kind == 'bash' then
      return run_bash( self, intent.command, intent.exclude )
    end
  end

  local snapshot = options.snapshot
  if snapshot == nil and options.capture ~= false then
    snapshot = context.snapshot( options.source )
  end

  local expanded, images, missing = text, nil, {}
  if options.compose ~= false then
    expanded, images, missing = compose.expand( text, { cwd = self.cwd, } )
    for _, reference in ipairs( missing ) do
      vim.notify( 'pivi: cannot read @' .. reference, vim.log.levels.WARN )
    end
  end

  local message = expanded
  if snapshot then
    message = context.render( snapshot ) .. expanded
  end

  render_submission( self, text )

  local command
  if delivery == 'steer' then
    command = { type = 'steer', message = message, }
  elseif delivery == 'follow_up' then
    command = { type = 'follow_up', message = message, }
  else
    command = { type = 'prompt', message = message, }
  end
  if images and #images > 0 then command.images = images end

  table.insert( self.submissions, { text = text, delivery = delivery, snapshot = snapshot, } )

  self.client.request( command, function( response )
    if response.success then return end
    local rejected = begin( self, 'error', { error = response.error, } )
    entry.line( self.buffer, rejected, tostring( response.error or 'submission rejected' ) )
    seal( self )
    self.status = vim.tbl_extend( 'force', self.status, { error = response.error, } )
    publish_status( self )
  end )

  return true, nil
end

local request = function( self, command, callback )
  return self.client.request( command, callback )
end

local abort = function( self, callback )
  return self.client.request( { type = 'abort', }, callback )
end

local clear_queue = function( self, callback )
  return self.client.request( { type = 'clear_queue', }, function( response )
    if response.success and response.data then
      self.status = status.reduce( self.status, {
        type = 'queue_update',
        steering = {},
        followUp = {},
      } )
      publish_status( self )
    end
    if callback then callback( response ) end
  end )
end

local refresh = function( self, callback )
  return self.client.request( { type = 'get_state', }, function( response )
    if response.success then
      self.status = status.reduce_state( self.status, response.data )
      publish_status( self )
    end
    if callback then callback( response ) end
  end )
end

local stop = function( self )
  self.client.stop()
  registry[self.name] = nil
  publish.remove( self.name )
end

local on = function( self, listener )
  table.insert( self.listeners, listener )
  return function()
    for index, candidate in ipairs( self.listeners ) do
      if candidate == listener then
        table.remove( self.listeners, index )
        return
      end
    end
  end
end

local stop_all = function()
  for _, name in ipairs( list() ) do
    local live = registry[name]
    if live then stop( live ) end
  end
end

local wrap = function( self )
  self.submissions = self.submissions or {}
  self.submit = function( text, options ) return submit( self, text, options ) end
  self.request = function( command, callback ) return request( self, command, callback ) end
  self.bash = function( command, exclude, callback ) return run_bash( self, command, exclude, callback ) end
  self.abort = function( callback ) return abort( self, callback ) end
  self.clear_queue = function( callback ) return clear_queue( self, callback ) end
  self.refresh = function( callback ) return refresh( self, callback ) end
  self.stop = function() return stop( self ) end
  self.on = function( listener ) return on( self, listener ) end
  return self
end

local create = function( options ) return wrap( new( options ) ) end

return {
  create = create,
  get = get,
  list = list,
  active = active,
  set_active = set_active,
  stop_all = stop_all,
  build_command = build_command,
}