import {
  applyUIMessageChunk,
  createUIMessageStreamSnapshotChunk,
  createUIMessageStreamState,
} from '@cradleapp/ai-sdk'
import type { UIMessageChunk } from 'ai'
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'

import type {
  ProductAnalyticsCorrelation,
  ProductAnalyticsTask,
} from '~/features/product-analytics/event-model'
import { chatSelectors, useChatStore } from '~/store/chat'
import { useRendererChatStore } from '~/store/renderer-chat'

import type { ChatStreamChunk } from './chat-stream-types'
import { liveChatStreamChunk, replayChatStreamChunk } from './chat-stream-types'
import { ChatStreamingHandler } from './chat-streaming-handler'
import {
  buildRawUIMessageChunkStreamFromResponse,
  buildUIMessageChunkStreamFromResponse,
  disposeChatRunBroadcast,
  onChatRunSettled,
} from './sse-chat-transport'

const productAnalyticsMocks = vi.hoisted(() => ({
  trackProductTaskFinished: vi.fn(),
  trackProductTaskStarted: vi.fn((
    task: ProductAnalyticsTask,
    startedAtMs = 0,
    correlation?: ProductAnalyticsCorrelation,
  ) => ({ task, startedAtMs, settled: false, correlation })),
}))

vi.mock('~/features/product-analytics/client', () => productAnalyticsMocks)

function resetChatStore(store: typeof useChatStore): void {
  store.setState(state => ({
    ...state,
    messagesMap: new Map(),
    hydratedSessionIds: new Set(),
    runStateMap: new Map(),
    streamLeaseMap: new Map(),
    activeAbortControllers: new Map(),
    runDisplayMetaMap: new Map(),
    errorMap: new Map(),
    assistantDisplaySplitMap: new Map(),
  }))
}

function createChunkStream(chunks: UIMessageChunk[], replay = false): ReadableStream<ChatStreamChunk> {
  return new ReadableStream<ChatStreamChunk>({
    start(controller) {
      for (const chunk of chunks) {
        controller.enqueue(replay ? replayChatStreamChunk(chunk) : liveChatStreamChunk(chunk))
      }
      controller.close()
    },
  })
}

const assistantTextChunks: UIMessageChunk[] = [
  { type: 'start', messageId: 'assistant-stream' },
  { type: 'text-start', id: 'text-1' },
  { type: 'text-delta', id: 'text-1', delta: 'Hello' },
  { type: 'text-end', id: 'text-1' },
  { type: 'finish', finishReason: 'stop' },
]

describe('chat streaming handler store boundary', () => {
  beforeEach(() => {
    resetChatStore(useChatStore)
    resetChatStore(useRendererChatStore)
    disposeChatRunBroadcast()
    vi.clearAllMocks()
  })

  afterEach(() => {
    disposeChatRunBroadcast()
  })

  it('writes chunk output to the injected renderer store without touching the main store', async () => {
    const settledEvents: unknown[] = []
    const unsubscribe = onChatRunSettled(event => settledEvents.push(event))
    const handler = new ChatStreamingHandler('side:side-conversation-1', 'assistant-local', 0, {
      store: useRendererChatStore,
      emitSettledEvents: false,
    })

    handler.start(new AbortController())
    await handler.consume(createChunkStream(assistantTextChunks))
    handler.finish()
    unsubscribe()

    const rendererMessages = chatSelectors.messages('side:side-conversation-1')(useRendererChatStore.getState())

    expect(rendererMessages).toHaveLength(1)
    expect(rendererMessages[0]).toMatchObject({
      id: 'assistant-stream',
      role: 'assistant',
      parts: [{ type: 'text', text: 'Hello', state: 'done' }],
    })
    expect(chatSelectors.isVisibleStreamingMessage('side:side-conversation-1', 'assistant-stream')(useRendererChatStore.getState())).toBe(false)
    expect(useChatStore.getState().messagesMap.has('side:side-conversation-1')).toBe(false)
    expect(settledEvents).toEqual([])
  })

  it('keeps streaming output when the server assigns the message id before a start chunk', async () => {
    const handler = new ChatStreamingHandler('session-1', 'assistant-local', 0)
    handler.start(new AbortController())
    handler.adoptServerMessageId('assistant-server')

    await handler.consume(createChunkStream([
      { type: 'text-start', id: 'text-1' },
      { type: 'text-delta', id: 'text-1', delta: 'Hello' },
    ]))

    const messages = chatSelectors.messages('session-1')(useChatStore.getState())
    expect(messages).toEqual([expect.objectContaining({
      id: 'assistant-server',
      role: 'assistant',
      parts: [expect.objectContaining({ type: 'text', text: 'Hello' })],
    })])
    expect(useChatStore.getState().streamLeaseMap.get('assistant-server')).toMatchObject({
      sessionId: 'session-1',
      source: 'local',
    })
  })

  it('restores a snapshot chunk and continues active deltas without failing the run', async () => {
    const target = {
      message: { id: 'assistant-recovered', role: 'assistant', parts: [] } satisfies import('ai').UIMessage,
      state: createUIMessageStreamState(),
    }
    applyUIMessageChunk(target, { type: 'text-start', id: 'text-recovered' })
    applyUIMessageChunk(target, { type: 'text-delta', id: 'text-recovered', delta: 'before ' })
    const snapshot = createUIMessageStreamSnapshotChunk({ runId: 'run-recovered', cursor: 40, target })

    const handler = new ChatStreamingHandler('session-1', 'assistant-recovered', 0, { mode: 'passive' })
    handler.start(new AbortController())
    await handler.consume(createChunkStream([
      snapshot,
      { type: 'text-delta', id: 'text-recovered', delta: 'after' },
      { type: 'text-end', id: 'text-recovered' },
    ], true))

    expect(chatSelectors.messages('session-1')(useChatStore.getState())).toMatchObject([
      { id: 'assistant-recovered', parts: [{ type: 'text', text: 'before after', state: 'done' }] },
    ])
    expect(useChatStore.getState().errorMap.has('assistant-recovered')).toBe(false)
  })

  it('publishes a recovery snapshot without waiting for a live chunk or stream close', async () => {
    const target = {
      message: { id: 'assistant-recovered', role: 'assistant', parts: [] } satisfies import('ai').UIMessage,
      state: createUIMessageStreamState(),
    }
    applyUIMessageChunk(target, { type: 'text-start', id: 'text-recovered' })
    applyUIMessageChunk(target, { type: 'text-delta', id: 'text-recovered', delta: 'Latest output' })
    const snapshot = createUIMessageStreamSnapshotChunk({ runId: 'run-recovered', cursor: 40, target })
    let streamController!: ReadableStreamDefaultController<ChatStreamChunk>
    const stream = new ReadableStream<ChatStreamChunk>({
      start(controller) {
        streamController = controller
      },
    })
    const handler = new ChatStreamingHandler('session-1', 'assistant-recovered', 0, { mode: 'passive' })

    handler.start(new AbortController())
    const consumePromise = handler.consume(stream)
    streamController.enqueue(replayChatStreamChunk(snapshot))

    await vi.waitFor(() => {
      expect(chatSelectors.messages('session-1')(useChatStore.getState())).toMatchObject([
        {
          id: 'assistant-recovered',
          parts: [{ type: 'text', text: 'Latest output', state: 'streaming' }],
        },
      ])
    })

    streamController.close()
    await consumePromise
    handler.finish()
  })

  it('emits settled events for the default main chat store path', () => {
    const settledEvents: unknown[] = []
    const unsubscribe = onChatRunSettled(event => settledEvents.push(event))
    const handler = new ChatStreamingHandler('session-1', 'assistant-1', 0)

    handler.start(new AbortController())
    handler.finish()
    unsubscribe()

    expect(settledEvents).toEqual([{
      chatSessionId: 'session-1',
      messageId: 'assistant-1',
      status: 'complete',
    }])
  })

  it('finishes local analytics tasks when UI settled broadcasts are disabled', () => {
    const handler = new ChatStreamingHandler('side:conversation-1', 'assistant-1', 0, {
      store: useRendererChatStore,
      emitSettledEvents: false,
    })

    handler.start(new AbortController())
    handler.finish()

    expect(productAnalyticsMocks.trackProductTaskStarted).toHaveBeenCalledWith({
      feature_domain: 'chat',
      task_kind: 'agent_run',
      task_variant: null,
    }, 0)
    expect(productAnalyticsMocks.trackProductTaskFinished).toHaveBeenCalledWith(
      expect.objectContaining({
        task: {
          feature_domain: 'chat',
          task_kind: 'agent_run',
          task_variant: null,
        },
      }),
      'success',
      null,
    )
  })

  it('carries server-issued telemetry correlation into the task lifecycle', () => {
    const handler = new ChatStreamingHandler('session-1', 'assistant-1', 0)

    handler.start(new AbortController())
    handler.setTelemetryCorrelation({
      session_id: 'opaque-session',
      run_id: 'opaque-run',
    })
    handler.finish()

    expect(productAnalyticsMocks.trackProductTaskStarted).toHaveBeenCalledWith(
      {
        feature_domain: 'chat',
        task_kind: 'agent_run',
        task_variant: null,
      },
      0,
      {
        session_id: 'opaque-session',
        run_id: 'opaque-run',
      },
    )
    expect(productAnalyticsMocks.trackProductTaskFinished).toHaveBeenCalledWith(
      expect.objectContaining({
        correlation: {
          session_id: 'opaque-session',
          run_id: 'opaque-run',
        },
      }),
      'success',
      null,
    )
  })

  it('clones into the store at flush so later chunks cannot mutate published snapshots', async () => {
    let streamController!: ReadableStreamDefaultController<ChatStreamChunk>
    const stream = new ReadableStream<ChatStreamChunk>({
      start(controller) {
        streamController = controller
      },
    })
    const handler = new ChatStreamingHandler('session-1', 'assistant-local', 0)

    handler.start(new AbortController())
    const consumePromise = handler.consume(stream)

    streamController.enqueue(liveChatStreamChunk({ type: 'start', messageId: 'assistant-stream' }))
    streamController.enqueue(liveChatStreamChunk({ type: 'text-start', id: 'text-1' }))
    streamController.enqueue(liveChatStreamChunk({ type: 'text-delta', id: 'text-1', delta: 'Hello' }))

    await vi.waitFor(() => {
      expect(chatSelectors.messages('session-1')(useChatStore.getState())[0]).toMatchObject({
        id: 'assistant-stream',
        parts: [{ type: 'text', text: 'Hello' }],
      })
    })

    const published = chatSelectors.messages('session-1')(useChatStore.getState())[0]!
    const publishedTextPart = published.parts[0]

    streamController.enqueue(liveChatStreamChunk({ type: 'text-delta', id: 'text-1', delta: ' world' }))
    streamController.close()
    await consumePromise
    handler.finish()

    // Live target mutations must not rewrite the previously published object graph.
    expect(publishedTextPart).toMatchObject({ type: 'text', text: 'Hello' })
    const latest = chatSelectors.messages('session-1')(useChatStore.getState())[0]!
    expect(latest).not.toBe(published)
    expect(latest.parts[0]).toMatchObject({ type: 'text', text: 'Hello world' })
  })

  it('applies replay chunks as a single store update when replay catches up', async () => {
    let resolveStreamController!: (controller: ReadableStreamDefaultController<ChatStreamChunk>) => void
    const streamControllerPromise = new Promise<ReadableStreamDefaultController<ChatStreamChunk>>((resolve) => {
      resolveStreamController = resolve
    })
    const stream = new ReadableStream<ChatStreamChunk>({
      start(nextController) {
        resolveStreamController(nextController)
      },
    })
    const handler = new ChatStreamingHandler('session-1', 'assistant-local', 0)

    handler.start(new AbortController())
    const consumePromise = handler.consume(stream)
    const streamController = await streamControllerPromise

    for (const chunk of assistantTextChunks) {
      streamController.enqueue(replayChatStreamChunk(chunk))
    }
    await Promise.resolve()

    expect(chatSelectors.messages('session-1')(useChatStore.getState())[0]?.parts).toEqual([])

    streamController.close()
    await consumePromise
    handler.finish()

    expect(chatSelectors.messages('session-1')(useChatStore.getState())[0]).toMatchObject({
      id: 'assistant-stream',
      parts: [{ type: 'text', text: 'Hello', state: 'done' }],
    })
  })

  it('marks HTTP stream chunks before the replay boundary as replay', async () => {
    const encoder = new TextEncoder()
    const response = new Response(encoder.encode([
      'data: {"type":"start","messageId":"assistant-stream"}\n\n',
      'data: {"type":"text-start","id":"text-1"}\n\n',
      ': cradle-replay-end\n\n',
      'data: {"type":"text-delta","id":"text-1","delta":"Live"}\n\n',
      'data: [DONE]\n\n',
    ].join('')))

    const reader = buildUIMessageChunkStreamFromResponse(response, 'session-1', { initialReplay: true }).getReader()
    const items: ChatStreamChunk[] = []
    while (true) {
      const result = await reader.read()
      if (result.done) {
        break
      }
      items.push(result.value)
    }

    expect(items.map(item => item.replay)).toEqual([true, true, false])
  })

  it('reuses the UIMessageChunk parser for stateless streams and consumes DONE as stream termination', async () => {
    const encoder = new TextEncoder()
    const response = new Response(encoder.encode([
      ': cradle-stream-open\n\n',
      'data: {"type":"text-start","id":"text-1"}\n\n',
      'data: {"type":"text-delta","id":"text-1","delta":"Answer"}\n\n',
      'data: {"type":"text-end","id":"text-1"}\n\n',
      'data: {"type":"finish","finishReason":"stop"}\n\n',
      'data: [DONE]\n\n',
    ].join('')))

    const reader = buildRawUIMessageChunkStreamFromResponse(response).getReader()
    const chunks: UIMessageChunk[] = []
    while (true) {
      const result = await reader.read()
      if (result.done) {
        break
      }
      chunks.push(result.value)
    }

    expect(chunks).toEqual([
      { type: 'text-start', id: 'text-1' },
      { type: 'text-delta', id: 'text-1', delta: 'Answer' },
      { type: 'text-end', id: 'text-1' },
      { type: 'finish', finishReason: 'stop' },
    ])
  })
})
