dotfiles

---@diagnostic disable: duplicate-set-field, need-check-nil, param-type-mismatch, assign-type-mismatch
local assemble = require 'bugabinga.pivi.assemble'
local fixture = dofile 'tests/helpers/pivi_fixture.lua'

local delta = function( value )
  return { type = 'message_update', assistantMessageEvent = value, }
end

local feed = function( events )
  local state = assemble.new()
  for _, event in ipairs( events ) do
    state = assemble.apply( state, event )
  end
  return state
end

describe( 'pivi assemble', function()
  it( 'concatenates text deltas by content index', function()
    local state = feed {
      { type = 'message_start', message = { role = 'assistant', }, },
      delta { type = 'text_start', contentIndex = 0, },
      delta { type = 'text_delta', contentIndex = 0, delta = 'Hello', },
      delta { type = 'text_delta', contentIndex = 0, delta = ' world', },
    }
    assert.equal( 'Hello world', assemble.text( state ) )
  end )

  it( 'keeps interleaved content blocks in index order', function()
    local state = feed {
      delta { type = 'text_delta', contentIndex = 2, delta = 'second', },
      delta { type = 'text_delta', contentIndex = 0, delta = 'first ', },
    }
    assert.equal( 'first second', assemble.text( state ) )
  end )

  it( 'separates thinking from visible text', function()
    local state = feed {
      delta { type = 'thinking_delta', contentIndex = 0, delta = 'pondering', },
      delta { type = 'text_delta', contentIndex = 1, delta = 'answer', },
    }
    assert.equal( 'answer', assemble.text( state ) )
    assert.equal( 'pondering', assemble.thinking( state ) )
  end )

  it( 'buffers tool call arguments and marks completion', function()
    local state = feed {
      delta { type = 'toolcall_start', contentIndex = 0, id = 'call_1', toolName = 'edit', },
      delta { type = 'toolcall_delta', contentIndex = 0, delta = '{"path":"a.lua"', },
    }

    local streaming = assemble.tool_calls( state )[1]
    assert.equal( 'call_1', streaming.id )
    assert.equal( 'edit', streaming.name )
    assert.equal( '{"path":"a.lua"', streaming.arguments_text )
    assert.is_false( streaming.complete )

    state = assemble.apply( state, delta {
      type = 'toolcall_end',
      contentIndex = 0,
      toolCall = { id = 'call_1', name = 'edit', arguments = { path = 'a.lua', }, },
    } )

    local finished = assemble.tool_calls( state )[1]
    assert.is_true( finished.complete )
    assert.equal( 'a.lua', finished.arguments.path )
  end )

  it( 'treats message_end as authoritative over accumulated deltas', function()
    local state = feed {
      delta { type = 'text_delta', contentIndex = 0, delta = 'partial junk', },
      {
        type = 'message_end',
        message = {
          role = 'assistant',
          content = { { type = 'text', text = 'final answer', }, },
        },
      },
    }
    assert.equal( 'final answer', assemble.text( state ) )
    assert.is_true( state.complete )
  end )

  it( 'resets on a new assistant message and ignores user messages', function()
    local state = feed {
      delta { type = 'text_delta', contentIndex = 0, delta = 'old', },
      { type = 'message_start', message = { role = 'user', }, },
    }
    assert.equal( 'old', assemble.text( state ) )

    state = assemble.apply( state, { type = 'message_start', message = { role = 'assistant', }, } )
    assert.equal( '', assemble.text( state ) )
  end )

  it( 'reproduces a recorded streaming answer', function()
    local state = assemble.new()
    for _, frame in ipairs( fixture.frames 'rpc-text-stream.jsonl' ) do
      state = assemble.apply( state, frame )
    end

    local final
    for _, frame in ipairs( fixture.of_type( 'rpc-text-stream.jsonl', 'message_end' ) ) do
      if frame.message and frame.message.role == 'assistant' then final = frame.message end
    end

    local expected = ''
    for _, block in ipairs( final.content ) do
      if block.type == 'text' then expected = expected .. block.text end
    end

    assert.equal( expected, assemble.text( state ) )
    assert.is_true( #assemble.thinking( state ) > 0 )
  end )

  it( 'assembles a recorded tool call from its argument deltas', function()
    local state = assemble.new()
    local streaming, completed = {}, {}

    -- A recorded run holds several assistant messages, and each one resets the
    -- assembled state, so a completed call is captured while it is current.
    for _, frame in ipairs( fixture.frames 'rpc-tool-stream.jsonl' ) do
      state = assemble.apply( state, frame )
      for _, call in ipairs( assemble.tool_calls( state ) ) do
        if call.complete then
          -- The first completion carries the streamed argument text; the later
          -- authoritative message_end rebuilds the call without it.
          completed[call.id] = completed[call.id] or call
        elseif call.arguments_text ~= '' then
          table.insert( streaming, call.arguments_text )
        end
      end
    end

    assert.is_true( #streaming > 1 )
    for index = 2, #streaming do
      assert.equal( streaming[index - 1], streaming[index]:sub( 1, #streaming[index - 1] ) )
    end

    local recorded
    for _, event in ipairs( fixture.deltas 'rpc-tool-stream.jsonl' ) do
      if event.type == 'toolcall_end' then recorded = event.toolCall end
    end

    assert.is_table( recorded )

    local finished = completed[recorded.id]
    assert.is_table( finished )
    assert.equal( recorded.name, finished.name )
    assert.equal( recorded.arguments.path, finished.arguments.path )
    assert.equal( recorded.arguments.path, vim.json.decode( finished.arguments_text ).path )
  end )

  it( 'survives a recorded delta that carries an empty payload', function()
    local empty = 0
    for _, event in ipairs( fixture.deltas 'rpc-tool-stream.jsonl' ) do
      if event.type == 'toolcall_delta' and ( event.delta == nil or event.delta == '' ) then
        empty = empty + 1
      end
    end
    assert.is_true( empty > 0 )

    local state = assemble.new()
    local seen = 0
    for _, frame in ipairs( fixture.frames 'rpc-tool-stream.jsonl' ) do
      state = assemble.apply( state, frame )
      seen = math.max( seen, #assemble.tool_calls( state ) )
    end
    assert.is_true( seen > 0 )
  end )

  it( 'ignores unknown events and malformed deltas', function()
    local state = feed {
      { type = 'totally_unknown_event', },
      { type = 'message_update', },
      delta { type = 'nonsense', contentIndex = 0, },
      'not a table',
    }
    assert.equal( '', assemble.text( state ) )
  end )
end )