import assert from 'node:assert/strict'
import { ChatStreamCommitter, type ChatStreamKey, type ChatStreamScheduler } from '../src/renderer/src/stores/chat-stream-committer'
import { applyAssistantEvent } from '../src/shared/message-parts'
import type { ChatEvent, ChatMessage, ToolCallRecord } from '../src/shared/types'

class FakeScheduler implements ChatStreamScheduler {
  private currentTime = 0
  private nextId = 1
  private readonly frames = new Map<number, () => void>()
  private readonly timers = new Map<number, { at: number; callback: () => void }>()

  now(): number { return this.currentTime }
  requestFrame(callback: () => void): number {
    const id = this.nextId++
    this.frames.set(id, callback)
    return id
  }
  cancelFrame(handle: unknown): void { this.frames.delete(handle as number) }
  setTimeout(callback: () => void, delayMs: number): number {
    const id = this.nextId++
    this.timers.set(id, { at: this.currentTime + delayMs, callback })
    return id
  }
  clearTimeout(handle: unknown): void { this.timers.delete(handle as number) }

  runFrame(at: number): void {
    this.currentTime = at
    const callbacks = [...this.frames.values()]
    this.frames.clear()
    callbacks.forEach((callback) => callback())
  }

  advanceTo(at: number): void {
    while (true) {
      const next = [...this.timers.entries()].sort((a, b) => a[1].at - b[1].at)[0]
      if (!next || next[1].at > at) break
      this.currentTime = next[1].at
      this.timers.delete(next[0])
      next[1].callback()
    }
    this.currentTime = at
  }
}

const key: ChatStreamKey = { sessionId: 'session-1', assistantId: 'assistant-1' }
const tool: ToolCallRecord = { id: 'tool-1', name: 'read_file', args: {}, argsRaw: '{}', status: 'running' }

function eventTypes(events: ChatEvent[]): string[] { return events.map((event) => event.type) }

function testSameFrameTextDeltasAreOneStoreCommit(): void {
  const scheduler = new FakeScheduler()
  const commits: ChatEvent[][] = []
  const committer = new ChatStreamCommitter((_key, events) => commits.push(events), scheduler)
  let message: ChatMessage = { id: key.assistantId, role: 'assistant', content: '', createdAt: 1, parts: [], status: 'streaming' }
  const apply = (events: ChatEvent[]) => { message = events.reduce((current, event) => applyAssistantEvent(current, event), message) }

  committer.enqueue(key, { runId: 'run-1', type: 'text', delta: 'start' })
  assert.equal(commits.length, 1, 'the first text block must render immediately')
  for (const delta of [' ', 'batched', ' ', 'text']) committer.enqueue(key, { runId: 'run-1', type: 'text', delta })
  assert.equal(commits.length, 1, 'same-frame deltas must not submit the store repeatedly')

  scheduler.runFrame(16)
  scheduler.advanceTo(149)
  assert.equal(commits.length, 1, 'the first batch must respect the 150ms minimum interval')
  scheduler.advanceTo(150)
  assert.equal(commits.length, 2, 'all same-frame deltas must share one store submission')
  assert.deepEqual(commits[1].map((event) => event.type), ['text'])
  apply(commits.flat())
  assert.equal(message.content, 'start batched text', 'coalescing must preserve exact text concatenation')
}

function testStructuralEventFlushesTextInOrder(): void {
  const scheduler = new FakeScheduler()
  const commits: ChatEvent[][] = []
  const committer = new ChatStreamCommitter((_key, events) => commits.push(events), scheduler)
  committer.enqueue(key, { runId: 'run-2', type: 'text', delta: 'before ' })
  committer.enqueue(key, { runId: 'run-2', type: 'text', delta: 'tool' })
  committer.enqueue(key, { runId: 'run-2', type: 'tool_call', call: tool })

  assert.equal(commits.length, 2, 'the tool event must force the queued text out immediately')
  assert.deepEqual(eventTypes(commits[1]), ['text', 'tool_call'], 'the pending text must precede its structural event')
  let message: ChatMessage = { id: key.assistantId, role: 'assistant', content: '', createdAt: 1, parts: [] }
  message = commits.flat().reduce((current, event) => applyAssistantEvent(current, event), message)
  assert.equal(message.content, 'before tool')
  assert.deepEqual(message.parts, [{ type: 'text', text: 'before tool' }, { type: 'tool', toolId: 'tool-1' }])
}

function testNewReasoningBlockIsImmediateAndUpdatesAreThrottled(): void {
  const scheduler = new FakeScheduler()
  const commits: ChatEvent[][] = []
  const committer = new ChatStreamCommitter((_key, events) => commits.push(events), scheduler)
  const reasoning = (text: string): ChatEvent => ({ runId: 'run-3', type: 'part', part: { type: 'reasoning', id: 'reason-1', text } })

  committer.enqueue(key, reasoning('first'))
  assert.equal(commits.length, 1, 'a new reasoning block must appear immediately')
  committer.enqueue(key, reasoning('second'))
  scheduler.runFrame(16)
  scheduler.advanceTo(149)
  assert.equal(commits.length, 1)
  scheduler.advanceTo(150)
  assert.equal(commits.length, 2)
  assert.deepEqual(commits[1], [reasoning('second')], 'repeated updates to the same reasoning block can be compacted')
}

function testTerminalFlushCancelsPendingSchedule(): void {
  const scheduler = new FakeScheduler()
  const commits: ChatEvent[][] = []
  const committer = new ChatStreamCommitter((_key, events) => commits.push(events), scheduler)
  committer.enqueue(key, { runId: 'run-4', type: 'text', delta: 'kept' })
  committer.enqueue(key, { runId: 'run-4', type: 'text', delta: ' before cancel' })
  committer.enqueue(key, { runId: 'run-4', type: 'cancelled' })

  assert.equal(commits.length, 2, 'cancellation must flush queued text before applying its terminal state')
  assert.deepEqual(eventTypes(commits[1]), ['text', 'cancelled'])
  scheduler.runFrame(16)
  scheduler.advanceTo(1000)
  assert.equal(commits.length, 2, 'terminal cleanup must cancel the scheduled frame and timer')
}

testSameFrameTextDeltasAreOneStoreCommit()
testStructuralEventFlushesTextInOrder()
testNewReasoningBlockIsImmediateAndUpdatesAreThrottled()
testTerminalFlushCancelsPendingSchedule()
console.log('chat-stream-committer-regression: all checks passed')