import { afterEach, describe, expect, it } from 'vitest'
import { Context } from '@deepseek-ai/cordis'
import LlmRuntime from '@deepseek-ai/dsh-llm'
import * as LlmPiAi from '@deepseek-ai/dsh-llm-pi-ai'
import { bundleProviders, collectStream } from './support.ts'

const contexts: Context[] = []

afterEach(async () => {
  await Promise.all(contexts.splice(0).map(ctx => ctx.fiber.dispose()))
})

async function harness(provider: string): Promise<Context> {
  const profile = bundleProviders()[provider]
  if (profile === undefined) throw new Error(`missing bundle provider ${provider}`)
  const ctx = new Context()
  contexts.push(ctx)
  await ctx.plugin(LlmRuntime)
  await ctx.plugin(LlmPiAi, { providers: { [provider]: profile } })
  return ctx
}

describe.skipIf(!process.env.MAAS_API_KEY)('Ant Digital MaaS live streaming', () => {
  it('streams a Responses answer before the terminal event', async () => {
    const ctx = await harness('ant-maas-responses')
    const result = await collectStream(
      ctx,
      'ant-maas-responses',
      process.env.DSH_MAAS_MODEL ?? 'lingdt-3.0-flash',
    )
    expect(result.text.length).toBeGreaterThan(0)
    expect(result.chunks.some(chunk => chunk.type === 'text-delta')).toBe(true)
    expect(result.chunks.at(-1)?.type).toBe('finish')
    expect(result.finish).toBe('stop')
  })
})

describe.skipIf(!process.env.NVIDIA_API_KEY)('NVIDIA NIM live streaming', () => {
  it('streams a Chat Completions answer before the terminal event', async () => {
    const ctx = await harness('nvidia')
    const result = await collectStream(
      ctx,
      'nvidia',
      process.env.DSH_NVIDIA_NIM_MODEL ?? 'nvidia/nemotron-3-super-120b-a12b',
    )
    expect(result.text.length).toBeGreaterThan(0)
    expect(result.chunks.some(chunk => chunk.type === 'text-delta')).toBe(true)
    expect(result.chunks.at(-1)?.type).toBe('finish')
    expect(result.finish).toBe('stop')
  })
})