import { XStream } from '@ant-design/x'
import { message as antdMessage } from 'antd'
import { uniqueId, throttle } from 'lodash-es'
import { extractJSON } from '@/utils/helper'
import { AIMessage } from './types'

export interface RequestOptions extends RequestInit {
  onUpdate?: (data: any) => void

  onSuccess?: (data: any) => void

  onError?: (error: any) => void
}

function transformFn<T>(data: T, transform?: (data: T) => any): any {
  return transform ? transform(data) : data
}

function mergeFn(messages: AIMessage[], merge?: (messages: AIMessage[]) => AIMessage) {
  return merge
    ? merge(messages)
    : messages.reduce(
        (result, message) => {
          const content = result.content + (message.content || '')
          const reasoning_content = result.reasoning_content + (message.reasoning_content || '')
          let full_content = ''
          if (reasoning_content) {
            full_content = `<div class="ai-thinking">${reasoning_content}`
            if (content) {
              full_content += `</div>${content}`
            }
          } else {
            full_content = content
          }
          return {
            ...result,
            ...message,
            content,
            reasoning_content,
            full_content,
            json: extractJSON(content),
          }
        },
        { content: '', reasoning_content: '' }
      )
}

export async function request(
  url: string,
  options: RequestOptions = {},
  transform?: (data: any) => any,
  merge?: (messages: AIMessage[]) => AIMessage
) {
  const { onUpdate, onSuccess, onError, body, ...fetchOptions } = options

  const id = uniqueId('chat_')
  const baseMessage: AIMessage = { id, role: 'assistant', content: '', status: 'loading' }
  const messages: AIMessage[] = []

  // 创建节流版本的 onUpdate
  const throttledOnUpdate = onUpdate ? throttle(onUpdate, 400, { leading: true, trailing: true }) : undefined

  try {
    messages.push({ ...baseMessage, status: 'loading' })

    const response = await fetch(url, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body,
      ...fetchOptions,
    })

    if (!response.ok) {
      throw new Error(`${response.status}:${response.statusText}`)
    }

    if (response.body && typeof body === 'string' && body.includes('stream')) {
      const stream = XStream({ readableStream: response.body })
      for await (const chunk of stream) {
        const { data } = chunk
        if (data && !data.includes('[DONE]')) {
          messages.push({ ...baseMessage, ...transformFn(JSON.parse(data), transform), status: 'typing' })
          throttledOnUpdate?.(mergeFn(messages, merge))
        }
      }
      messages.push({ ...baseMessage, status: 'success' })
    } else {
      const json = await response.json()
      messages.push({ ...baseMessage, ...transformFn(json, transform), status: 'success' })
    }

    onSuccess?.(mergeFn(messages, merge))
  } catch (error: any) {
    console.error(error)
    const errorMsg = error?.message || '请求失败'
    if (error.name !== 'AbortError') {
      antdMessage.error(errorMsg)
    }
    messages.push({ ...baseMessage, content: errorMsg, status: 'error' })
    onError?.(mergeFn(messages, merge))
  }

  return mergeFn(messages, merge)
}
