import { describe, it, expect } from "vitest"
import { ChatService, type ChatEvent, type MessageResponse } from "./chat"

describe("ChatService SSE Parsing", () => {
  it("should correctly parse a text stream chunk", () => {
    const rawChunk = `data: {"type": "on_chat_model_stream", "name": "synthesizer", "data": {"chunk": "Hello"}}\n\n`
    const events = ChatService.parseSSEChunk(rawChunk)
    
    expect(events).toHaveLength(1)
    expect(events[0]).toEqual({
      type: "on_chat_model_stream",
      name: "synthesizer",
      data: { chunk: "Hello" }
    })
  })

  it("should handle multiple events in one chunk", () => {
    const rawChunk = 
      `data: {"type": "on_chat_model_stream", "name": "synthesizer", "data": {"chunk": "Hello"}}\n\n` +
      `data: {"type": "on_chat_model_stream", "name": "synthesizer", "data": {"chunk": " World"}}\n\n`
    
    const events = ChatService.parseSSEChunk(rawChunk)
    
    expect(events).toHaveLength(2)
    expect(events[1].data!.chunk).toBe(" World")
  })

  it("should parse encoded plots from data_analyst_worker node", () => {
    const rawChunk = `data: {"type": "on_chain_end", "name": "data_analyst_worker", "data": {"encoded_plots": ["base64data"]}}\n\n`
    const events = ChatService.parseSSEChunk(rawChunk)
    
    expect(events[0].data!.encoded_plots).toEqual(["base64data"])
  })

  it("should identify the done event", () => {
    const rawChunk = `data: {"type": "done"}\n\n`
    const events = ChatService.parseSSEChunk(rawChunk)
    
    expect(events[0].type).toBe("done")
  })
})

describe("ChatService Message State Management", () => {
  it("should append text chunks to the last message content", () => {
    const messages: MessageResponse[] = [{ id: "1", role: "assistant", content: "Initial", created_at: new Date().toISOString() }]
    const event: ChatEvent = {
      type: "on_chat_model_stream",
      node: "synthesizer",
      data: { chunk: { content: " text" } }
    }
    
    const updatedMessages = ChatService.updateMessagesWithEvent(messages, event)
    expect(updatedMessages[0].content).toBe("Initial text")
  })

  it("should add plots to the message state", () => {
    const messages: MessageResponse[] = [{ id: "1", role: "assistant", content: "Analysis", created_at: new Date().toISOString(), plots: [] }]
    const event: ChatEvent = {
      type: "on_chain_end",
      name: "data_analyst_worker",
      data: { encoded_plots: ["plot1"] }
    }
    
    const updatedMessages = ChatService.updateMessagesWithEvent(messages, event)
    expect(updatedMessages[0].plots).toEqual(["plot1"])
  })
})