import assert from 'node:assert/strict'
import { createServer } from 'node:http'
import { Gateway, GatewayError, type GatewayDelta } from '../src/main/gateway'

const requests = new Map<string, Array<Record<string, unknown>>>()
const abortConnections = new Set<import('node:http').ServerResponse>()

function attempts(model: string): Array<Record<string, unknown>> {
  return requests.get(model) ?? []
}

function recordAttempt(body: Record<string, unknown>): number {
  const model = String(body.model)
  const items = requests.get(model) ?? []
  items.push(body)
  requests.set(model, items)
  return items.length
}

function successSse(response: import('node:http').ServerResponse, usage: Record<string, unknown>): void {
  response.writeHead(200, { 'content-type': 'text/event-stream' })
  response.end([
    { choices: [{ delta: { content: 'ok' }, finish_reason: null }] },
    { choices: [{ delta: {}, finish_reason: 'stop' }] },
    { choices: [], usage }
  ].map((frame) => `data: ${JSON.stringify(frame)}\n\n`).join('') + 'data: [DONE]\n\n')
}

const server = createServer(async (request, response) => {
  const chunks: Buffer[] = []
  for await (const chunk of request) chunks.push(Buffer.from(chunk))
  const body = JSON.parse(Buffer.concat(chunks).toString()) as Record<string, unknown>
  const model = String(body.model)
  const attempt = recordAttempt(body)
  const validUsage = {
    prompt_tokens: 128,
    completion_tokens: 2,
    total_tokens: 130,
    prompt_tokens_details: { cached_tokens: 32 }
  }

  if (model === 'valid' || model === 'invalid-usage') {
    successSse(response, model === 'invalid-usage'
      ? { prompt_tokens: -1, completion_tokens: 2, total_tokens: 1 }
      : validUsage)
    return
  }

  if (model === 'legacy-400') {
    if (attempt === 1 && body.stream_options) {
      response.writeHead(400, { 'content-type': 'application/json' })
      response.end(JSON.stringify({ error: { message: 'Unsupported parameter: stream_options' } }))
      return
    }
    successSse(response, validUsage)
    return
  }

  if (model === 'legacy-422') {
    if (attempt === 1 && body.stream_options) {
      response.writeHead(422, { 'content-type': 'application/json' })
      response.end(JSON.stringify({ error: { message: 'include_usage is not supported by this server' } }))
      return
    }
    successSse(response, validUsage)
    return
  }

  if (model === 'ordinary-400') {
    response.writeHead(400, { 'content-type': 'application/json' })
    response.end(JSON.stringify({
      error: { message: 'request validation failed while parsing stream_options' },
      apiKey: 'UPSTREAM_SECRET_MUST_NOT_ECHO'
    }))
    return
  }

  if (model === 'auth-401') {
    response.writeHead(401, { 'content-type': 'application/json' })
    response.end(JSON.stringify({
      error: { message: 'invalid api key' },
      authorization: 'Bearer UPSTREAM_SECRET_MUST_NOT_ECHO'
    }))
    return
  }

  if (model === 'server-500-options') {
    response.writeHead(500, { 'content-type': 'application/json' })
    response.end(JSON.stringify({ error: { message: 'stream_options unsupported due to server failure' } }))
    return
  }

  if (model === 'started-sse-error') {
    response.writeHead(200, { 'content-type': 'text/event-stream' })
    response.write(`data: ${JSON.stringify({ choices: [{ delta: { content: 'before-error' }, finish_reason: null }] })}\n\n`)
    response.write(`data: ${JSON.stringify({ error: { message: 'stream_options unsupported' } })}\n\n`)
    response.end()
    return
  }

  if (model === 'abort-retry') {
    if (attempt === 1 && body.stream_options) {
      response.writeHead(400, { 'content-type': 'application/json' })
      response.end(JSON.stringify({ error: { message: 'unknown field: stream_options' } }))
      return
    }
    abortConnections.add(response)
    response.writeHead(200, {
      'content-type': 'text/event-stream',
      'cache-control': 'no-cache',
      connection: 'keep-alive'
    })
    response.flushHeaders()
    const keepalive = setInterval(() => {
      if (!response.destroyed) response.write(': keepalive\n\n')
    }, 100)
    response.on('close', () => {
      clearInterval(keepalive)
      abortConnections.delete(response)
    })
    return
  }

  response.writeHead(400, { 'content-type': 'application/json' })
  response.end(JSON.stringify({ error: { message: `unknown QA model: ${model}` } }))
})

function gatewayFor(model: string): Gateway {
  const address = server.address()
  assert.ok(address && typeof address === 'object')
  return new Gateway({
    baseUrl: `http://127.0.0.1:${address.port}/v1`,
    apiKey: '',
    authMode: 'none',
    model,
    temperature: 0,
    maxToolRounds: 1
  })
}

async function collect(model: string, signal?: AbortSignal): Promise<GatewayDelta[]> {
  const events: GatewayDelta[] = []
  for await (const event of gatewayFor(model).stream({
    messages: [{ role: 'user', content: 'test' }],
    signal
  })) {
    events.push(event)
  }
  return events
}

async function captureError(model: string): Promise<GatewayError> {
  try {
    await collect(model)
  } catch (error) {
    assert.ok(error instanceof GatewayError, `${model} must reject with GatewayError`)
    return error
  }
  assert.fail(`${model} should reject`)
}

async function waitForAttempts(model: string, count: number): Promise<void> {
  for (let index = 0; index < 250; index++) {
    if (attempts(model).length === count) return
    await new Promise((resolve) => setTimeout(resolve, 10))
  }
  assert.fail(`${model} did not reach ${count} requests`)
}

await new Promise<void>((resolve) => server.listen(0, '127.0.0.1', resolve))
try {
  const valid = await collect('valid')
  assert.deepEqual(valid.filter((event) => event.type === 'text'), [{ type: 'text', text: 'ok' }])
  assert.deepEqual(valid.filter((event) => event.type === 'finish'), [{ type: 'finish', reason: 'stop' }])
  assert.deepEqual(valid.filter((event) => event.type === 'usage'), [{
    type: 'usage',
    usage: { inputTokens: 128, outputTokens: 2, totalTokens: 130, cachedTokens: 32 }
  }])
  assert.equal(attempts('valid').length, 1)
  assert.deepEqual(attempts('valid')[0]?.stream_options, { include_usage: true })

  const invalidUsage = await collect('invalid-usage')
  assert.equal(invalidUsage.filter((event) => event.type === 'usage').length, 0,
    'invalid usage must not emit a usage event')

  for (const model of ['legacy-400', 'legacy-422']) {
    const events = await collect(model)
    assert.deepEqual(events.filter((event) => event.type === 'text'), [{ type: 'text', text: 'ok' }])
    assert.deepEqual(events.filter((event) => event.type === 'usage'), [{
      type: 'usage',
      usage: { inputTokens: 128, outputTokens: 2, totalTokens: 130, cachedTokens: 32 }
    }])
    const tries = attempts(model)
    assert.equal(tries.length, 2, `${model} must retry exactly once`)
    assert.deepEqual(tries[0]?.stream_options, { include_usage: true })
    assert.equal(Object.hasOwn(tries[1] ?? {}, 'stream_options'), false,
      `${model} retry must omit stream_options`)
  }

  const ordinary = await captureError('ordinary-400')
  assert.equal(ordinary.status, 400)
  assert.equal(attempts('ordinary-400').length, 1, 'an ordinary 400 must not retry')
  assert.equal(ordinary.message.includes('UPSTREAM_SECRET_MUST_NOT_ECHO'), false,
    'upstream body content must not be copied into the renderer-facing message')

  const auth = await captureError('auth-401')
  assert.equal(auth.status, 401)
  assert.equal(attempts('auth-401').length, 1, 'authentication failures must not retry')
  assert.equal(auth.message.includes('UPSTREAM_SECRET_MUST_NOT_ECHO'), false,
    'authorization material must not be copied into the renderer-facing message')

  const serverFailure = await captureError('server-500-options')
  assert.equal(serverFailure.status, 500)
  assert.equal(attempts('server-500-options').length, 1,
    'a 5xx response must not retry even when its body names stream_options')

  const startedSse = await captureError('started-sse-error')
  assert.equal(attempts('started-sse-error').length, 1,
    'an error after SSE has started must not trigger another request')
  assert.match(startedSse.message, /上游错误/)

  const controller = new AbortController()
  const aborted = collect('abort-retry', controller.signal)
  await waitForAttempts('abort-retry', 2)
  controller.abort()
  await assert.rejects(aborted, (error: unknown) =>
    error instanceof Error && error.name === 'AbortError')
  assert.equal(attempts('abort-retry').length, 2,
    'the shared AbortSignal must cancel the in-flight compatibility retry')

  console.log('chat-completions-usage-regression: all checks passed')
} finally {
  for (const response of abortConnections) response.destroy()
  server.closeAllConnections?.()
  await new Promise<void>((resolve) => server.close(() => resolve()))
}