import { describe, expect, it } from "vitest"
import { stopAtStartOf, stopAtStopTokens } from "./charStream"

async function* createMockStream(chunks: string[]): AsyncGenerator<string> {
  for (const chunk of chunks) {
    yield chunk
  }
}

async function streamToString(stream: AsyncGenerator<string>): Promise<string> {
  let result = ""
  for await (const chunk of stream) {
    result += chunk
  }
  return result
}

describe("stopAtStopTokens", () => {
  it("should yield characters until a stop token is encountered", async () => {
    const mockStream = createMockStream(["Hello", " world", "! Stop", "here"])
    const stopTokens = ["Stop"]
    const result = stopAtStopTokens(mockStream, stopTokens)

    const output = []
    for await (const char of result) {
      output.push(char)
    }

    expect(output.join("")).toBe("Hello world! ")
  })

  it("should handle multiple stop tokens", async () => {
    const mockStream = createMockStream(["This", " is a ", "test. END", " of stream"])
    const stopTokens = ["END", "STOP", "HALT"]
    const result = stopAtStopTokens(mockStream, stopTokens)

    expect(await streamToString(result)).toBe("This is a test. ")
  })

  it("should handle stop tokens split across chunks", async () => {
    const mockStream = createMockStream(["Hello", " wo", "r", "ld! ST", "OP now"])
    const stopTokens = ["STOP"]
    const result = stopAtStopTokens(mockStream, stopTokens)

    const output = []
    for await (const char of result) {
      output.push(char)
    }

    expect(output.join("")).toBe("Hello world! ")
  })

  it("should yield all characters if no stop token is encountered", async () => {
    const mockStream = createMockStream(["This", " is ", "a complete", " stream"])
    const stopTokens = ["END"]
    const result = stopAtStopTokens(mockStream, stopTokens)

    expect(await streamToString(result)).toBe("This is a complete stream")
  })

  it("should handle empty chunks", async () => {
    const mockStream = createMockStream(["Hello", "", " world", "", "! STOP"])
    const stopTokens = ["STOP"]
    const result = stopAtStopTokens(mockStream, stopTokens)

    expect(await streamToString(result)).toBe("Hello world! ")
  })

  it("should handle stop token at the beginning of the stream", async () => {
    const mockStream = createMockStream(["STOP", "Hello world"])
    const stopTokens = ["STOP"]
    const result = stopAtStopTokens(mockStream, stopTokens)

    const output = []
    for await (const char of result) {
      output.push(char)
    }

    expect(output.join("")).toBe("")
  })

  it("should handle stop token at the end of the stream", async () => {
    const mockStream = createMockStream(["Hello world", "STOP"])
    const stopTokens = ["STOP"]
    const result = stopAtStopTokens(mockStream, stopTokens)

    expect(await streamToString(result)).toBe("Hello world")
  })

  it("should handle multiple stop tokens of different lengths", async () => {
    const mockStream = createMockStream(["This is a ", "test with ", "multiple STOP", " tokens END"])
    const stopTokens = ["STOP", "END", "HALT"]
    const result = stopAtStopTokens(mockStream, stopTokens)

    expect(await streamToString(result)).toBe("This is a test with multiple ")
  })

  it("should handle an empty stream", async () => {
    const mockStream = createMockStream([])
    const stopTokens = ["STOP"]
    const result = stopAtStopTokens(mockStream, stopTokens)

    expect(await streamToString(result)).toBe("")
  })

  it("should handle an empty stop tokens array", async () => {
    const mockStream = createMockStream(["Hello", " world!"])
    const stopTokens: string[] = []
    const result = stopAtStopTokens(mockStream, stopTokens)

    const output = []
    for await (const char of result) {
      output.push(char)
    }

    expect(output.join("")).toBe("Hello world!")
  })

  it("should handle stop token when remaining buffer is smaller than maximum stop token length", async () => {
    const mockStream = createMockStream(["Hello world!STOP"])
    const stopTokens: string[] = ["STOP", "STOP_TOKEN_THAT_IS_LARGER_THAN_BUFFER"]
    const result = stopAtStopTokens(mockStream, stopTokens)

    expect(await streamToString(result)).toBe("Hello world!")
  })
})

describe("stopAtStartOf", () => {
  const sampleCode = `      {
        method: "GET",
        headers: {
          "Content-Type": "application/json",
          Authorization: \`Bearer \${this.workOsAccessToken}\`,
        },
      },
    );
    const data = await response.json();
    return data.items;
  }

  async getContextItems(
    query: string,
    extras: ContextProviderExtras,
  ): Promise<ContextItem[]> {
    const response = await extras.fetch(
      new URL(
        \`/proxy/context/\${this.options.id}/retrieve\`,
        controlPlaneEnv.CONTROL_PLANE_URL,
      ),
`

  /* Some LLMs, such as Codestral, repeat the suffix of the query. To test our filtering, we cut the sample code at random positions, remove a part of the input
and construct a response, containing the removed part and the suffix. The goal of the stopAtStartOf() method is to detect the start of the suffix in the response */
  it("should stop if the start of the suffix is reached", async () => {
    const suffix = `
  const data = await response.json();
  return data.items;
}`
    const mockStream = createMockStream(sampleCode.split(/(?! )/g))
    const result = stopAtStartOf(mockStream, suffix)

    const resultStr = await streamToString(result)
    expect(resultStr).toBe(`      {
        method: "GET",
        headers: {
          "Content-Type": "application/json",
          Authorization: \`Bearer \${this.workOsAccessToken}\`,
        },
      },
    );
    `)
  })
  it("should stop if the start of the suffix is reached, even if the suffix has a prefix", async () => {
    const suffix = `
  xxxconst data = await response.json();
  return data.items;
}`
    const mockStream = createMockStream(sampleCode.split(/(?! )/g))
    const result = stopAtStartOf(mockStream, suffix)

    const resultStr = await streamToString(result)
    expect(resultStr).toBe(`      {
        method: "GET",
        headers: {
          "Content-Type": "application/json",
          Authorization: \`Bearer \${this.workOsAccessToken}\`,
        },
      },
    );
    `)
  })
})
