// Generated by continue

import { describe, expect, it, vi } from "vitest"
import { ChatMessage } from "../index"
import { generateLines, streamLines } from "./util"

describe("streamLines", () => {
  it("should split chunks into lines correctly", async () => {
    async function* streamCompletion(): AsyncGenerator<string> {
      yield "line1\nline"
      yield "2\nline3\n"
      yield "line4"
    }

    const resultLines: string[] = []
    for await (const line of streamLines(streamCompletion())) {
      resultLines.push(line)
    }

    expect(resultLines).toEqual(["line1", "line2", "line3", "line4"])
  })

  it("should handle ChatMessage chunks", async () => {
    const messageChunk1: ChatMessage = {
      role: "assistant",
      content: "line1\nline",
    }
    const messageChunk2: ChatMessage = {
      role: "assistant",
      content: "2\nline3\n",
    }
    const messageChunk3: ChatMessage = {
      role: "assistant",
      content: "line4",
    }

    // const spy = vi.spyOn(messageContentModule, "renderChatMessage");

    async function* streamCompletion(): AsyncGenerator<ChatMessage> {
      yield messageChunk1
      yield messageChunk2
      yield messageChunk3
    }

    const resultLines: string[] = []
    for await (const line of streamLines(streamCompletion())) {
      resultLines.push(line)
    }

    expect(resultLines).toEqual(["line1", "line2", "line3", "line4"])
    // expect(spy).toHaveBeenCalledTimes(3);
  })

  it("should log lines if log parameter is true", async () => {
    const consoleSpy = vi.spyOn(console, "log").mockImplementation(() => {})

    async function* streamCompletion(): AsyncGenerator<string> {
      yield "line1\nline2\n"
      yield "line3"
    }

    const resultLines: string[] = []
    for await (const line of streamLines(streamCompletion(), true)) {
      resultLines.push(line)
    }

    expect(resultLines).toEqual(["line1", "line2", "line3"])
    expect(consoleSpy).toHaveBeenCalledWith("Streamed lines: ", "line1\nline2\nline3")

    consoleSpy.mockRestore()
  })
})

describe("generateLines", () => {
  it("should yield the lines provided in the array", async () => {
    const lines = ["line1", "line2", "line3"]
    const resultLines: string[] = []

    for await (const line of generateLines(lines)) {
      resultLines.push(line)
    }

    expect(resultLines).toEqual(lines)
  })
})
