import { loggerService } from '@logger'
import { autoRenameTopic } from '@renderer/hooks/useTopic'
import i18n from '@renderer/i18n'
import { EVENT_NAMES, EventEmitter } from '@renderer/services/EventService'
import { NotificationService } from '@renderer/services/NotificationService'
import { estimateMessagesUsage } from '@renderer/services/TokenService'
import { isTodoWriteBlock, updateOneBlock } from '@renderer/store/messageBlock'
import { selectMessagesForTopic } from '@renderer/store/newMessage'
import { newMessagesActions } from '@renderer/store/newMessage'
import { toolPermissionsActions } from '@renderer/store/toolPermissions'
import type { Assistant } from '@renderer/types'
import { ERROR_I18N_KEY_REQUEST_TIMEOUT, ERROR_I18N_KEY_STREAM_PAUSED } from '@renderer/types/error'
import type {
  MessageBlock,
  PlaceholderMessageBlock,
  Response,
  ThinkingMessageBlock,
  ToolMessageBlock
} from '@renderer/types/newMessage'
import { AssistantMessageStatus, MessageBlockStatus, MessageBlockType } from '@renderer/types/newMessage'
import { uuid } from '@renderer/utils'
import { isAgentSessionTopicId } from '@renderer/utils/agentSession'
import { trackTokenUsage } from '@renderer/utils/analytics'
import { isAbortError, isTimeoutError, serializeError } from '@renderer/utils/error'
import { createBaseMessageBlock, createErrorBlock } from '@renderer/utils/messageUtils/create'
import { findAllBlocks, getMainTextContent } from '@renderer/utils/messageUtils/find'
import { isFocused, isOnHomePage } from '@renderer/utils/window'
import type { AISDKError } from 'ai'
import { NoOutputGeneratedError } from 'ai'

import type { BlockManager } from '../BlockManager'

const logger = loggerService.withContext('BaseCallbacks')
interface BaseCallbacksDependencies {
  blockManager: BlockManager
  dispatch: any
  getState: any
  topicId: string
  assistantMsgId: string
  saveUpdatesToDB: any
  assistant: Assistant
  getCurrentThinkingInfo?: () => { blockId: string | null; millsec: number }
}

export const createBaseCallbacks = (deps: BaseCallbacksDependencies) => {
  const {
    blockManager,
    dispatch,
    getState,
    topicId,
    assistantMsgId,
    saveUpdatesToDB,
    assistant,
    getCurrentThinkingInfo
  } = deps

  const startTime = Date.now()
  const notificationService = NotificationService.getInstance()

  // 通用的 block 查找函数
  const findBlockIdForCompletion = (message?: any) => {
    // 优先使用 BlockManager 中的 activeBlockInfo
    const activeBlockInfo = blockManager.activeBlockInfo

    if (activeBlockInfo) {
      return activeBlockInfo.id
    }

    // 如果没有活跃的block，从message中查找最新的block作为备选
    const targetMessage = message || getState().messages.entities[assistantMsgId]
    if (targetMessage) {
      const allBlocks = findAllBlocks(targetMessage)
      if (allBlocks.length > 0) {
        return allBlocks[allBlocks.length - 1].id // 返回最新的block
      }
    }

    // 最后的备选方案：从 blockManager 获取占位符块ID
    return blockManager.initialPlaceholderBlockId
  }

  /**
   * Mark in_progress todos as completed when stream ends,
   * since the model will no longer update them.
   */
  const cleanupInProgressTodos = (): string[] => {
    const currentMessage = getState().messages.entities[assistantMsgId]
    if (!currentMessage) return []

    const allBlockRefs = findAllBlocks(currentMessage)
    const blockState = getState().messageBlocks
    const cleanedBlockIds: string[] = []

    for (const blockRef of allBlockRefs) {
      const block = blockState.entities[blockRef.id]
      if (!isTodoWriteBlock(block)) continue

      const toolResponse = block.metadata.rawMcpToolResponse
      const todos = toolResponse.arguments.todos
      if (!todos.some((todo) => todo.status === 'in_progress')) continue

      const updatedTodos = todos.map((todo) =>
        todo.status === 'in_progress' ? { ...todo, status: 'completed' as const } : todo
      )

      dispatch(
        updateOneBlock({
          id: block.id,
          changes: {
            metadata: {
              ...block.metadata,
              rawMcpToolResponse: {
                ...toolResponse,
                arguments: { todos: updatedTodos }
              }
            }
          }
        })
      )
      cleanedBlockIds.push(block.id)
    }

    return cleanedBlockIds
  }

  return {
    onLLMResponseCreated: async () => {
      const baseBlock = createBaseMessageBlock(assistantMsgId, MessageBlockType.UNKNOWN, {
        status: MessageBlockStatus.PROCESSING
      })
      await blockManager.handleBlockTransition(baseBlock as PlaceholderMessageBlock, MessageBlockType.UNKNOWN)
    },

    onError: async (error: AISDKError) => {
      logger.debug('onError', error)
      if (NoOutputGeneratedError.isInstance(error)) {
        return
      }
      const isErrorTypeAbort = isAbortError(error)
      const isErrorTypeTimeout = isTimeoutError(error)
      const serializableError = serializeError(error)
      if (isErrorTypeAbort) {
        serializableError.i18nKey = ERROR_I18N_KEY_STREAM_PAUSED
      } else if (isErrorTypeTimeout) {
        serializableError.i18nKey = ERROR_I18N_KEY_REQUEST_TIMEOUT
      }

      const duration = Date.now() - startTime
      // 发送错误通知（除了中止错误）
      if (!isErrorTypeAbort) {
        const timeOut = duration > 30 * 1000
        if ((!isOnHomePage() && timeOut) || (!isFocused() && timeOut)) {
          await notificationService.send({
            id: uuid(),
            type: 'error',
            title: i18n.t('notification.assistant'),
            message: serializableError.message ?? '',
            silent: false,
            timestamp: Date.now(),
            source: 'assistant'
          })
        }
      }

      const possibleBlockId = findBlockIdForCompletion()

      if (possibleBlockId) {
        // 更改上一个block的状态为ERROR/PAUSED
        const changes: Partial<ThinkingMessageBlock> = {
          status: isErrorTypeAbort ? MessageBlockStatus.PAUSED : MessageBlockStatus.ERROR
        }
        // 如果是 thinking block，保留实际思考时间
        if (blockManager.lastBlockType === MessageBlockType.THINKING) {
          const thinkingInfo = getCurrentThinkingInfo?.()
          if (thinkingInfo?.blockId === possibleBlockId && thinkingInfo?.millsec && thinkingInfo.millsec > 0) {
            changes.thinking_millsec = thinkingInfo.millsec
          }
        }
        blockManager.smartBlockUpdate(possibleBlockId, changes, blockManager.lastBlockType!, true)
      }

      // Fix: 更新所有仍处于 STREAMING 状态的 blocks 为 PAUSED/ERROR
      // 这修复了停止回复时思考计时器继续运行的问题
      const currentMessage = getState().messages.entities[assistantMsgId]
      const updatedBlockIds: string[] = []
      if (currentMessage) {
        const allBlockRefs = findAllBlocks(currentMessage)
        const blockState = getState().messageBlocks
        // 获取当前思考信息（如果有），用于保留实际思考时间
        const thinkingInfo = getCurrentThinkingInfo?.()
        for (const blockRef of allBlockRefs) {
          const block = blockState.entities[blockRef.id]
          if (!block) continue

          // 更新非 possibleBlockId 的 STREAMING blocks（possibleBlockId 已在上面处理）
          // 跳过 TOOL 类型 blocks，它们在下面的 tool block 分支中统一处理
          if (
            block.id !== possibleBlockId &&
            block.status === MessageBlockStatus.STREAMING &&
            block.type !== MessageBlockType.TOOL
          ) {
            const changes: Partial<ThinkingMessageBlock> = {
              status: isErrorTypeAbort ? MessageBlockStatus.PAUSED : MessageBlockStatus.ERROR
            }
            if (
              block.type === MessageBlockType.THINKING &&
              thinkingInfo?.blockId === block.id &&
              thinkingInfo?.millsec &&
              thinkingInfo.millsec > 0
            ) {
              changes.thinking_millsec = thinkingInfo.millsec
            }
            dispatch(updateOneBlock({ id: block.id, changes }))
            updatedBlockIds.push(block.id)
          }

          // Fix: 更新所有仍处于非完成状态的 tool blocks 的 rawMcpToolResponse.status
          // 当用户点击停止时，tool blocks 的 UI 状态依赖 rawMcpToolResponse.status，
          // 而不是 MessageBlockStatus，所以需要单独更新
          if (block.type === MessageBlockType.TOOL) {
            const toolBlock = block as ToolMessageBlock
            const toolResponse = toolBlock.metadata?.rawMcpToolResponse
            const toolStatus = toolResponse?.status
            if (
              toolResponse &&
              toolStatus &&
              toolStatus !== 'done' &&
              toolStatus !== 'error' &&
              toolStatus !== 'cancelled'
            ) {
              dispatch(
                updateOneBlock({
                  id: block.id,
                  changes: {
                    status: isErrorTypeAbort ? MessageBlockStatus.PAUSED : MessageBlockStatus.ERROR,
                    metadata: {
                      ...toolBlock.metadata,
                      rawMcpToolResponse: {
                        ...toolResponse,
                        status: isErrorTypeAbort ? 'cancelled' : 'error'
                      }
                    }
                  }
                })
              )
              updatedBlockIds.push(block.id)
            }
          }
        }
      }

      // Clean up pending/submitting tool permission requests from this stream.
      // Preserve 'invoking' entries as they may belong to concurrent streams.
      dispatch(toolPermissionsActions.clearPending())

      // Mark in_progress todos as completed since stream ended
      const todoCleanupIds = cleanupInProgressTodos()
      updatedBlockIds.push(...todoCleanupIds)

      const errorBlock = createErrorBlock(assistantMsgId, serializableError, { status: MessageBlockStatus.SUCCESS })
      await blockManager.handleBlockTransition(errorBlock, MessageBlockType.ERROR)
      const messageErrorUpdate = {
        status: isErrorTypeAbort ? AssistantMessageStatus.SUCCESS : AssistantMessageStatus.ERROR
      }
      dispatch(
        newMessagesActions.updateMessage({
          topicId,
          messageId: assistantMsgId,
          updates: messageErrorUpdate
        })
      )

      // 从更新后的 state 中获取需要持久化的 blocks
      const blocksToSave = updatedBlockIds
        .map((id) => getState().messageBlocks.entities[id])
        .filter(Boolean) as MessageBlock[]
      await saveUpdatesToDB(assistantMsgId, topicId, messageErrorUpdate, blocksToSave)

      void EventEmitter.emit(EVENT_NAMES.MESSAGE_COMPLETE, {
        id: assistantMsgId,
        topicId,
        status: isErrorTypeAbort ? 'pause' : 'error',
        error: error.message
      })
    },

    onComplete: async (status: AssistantMessageStatus, response?: Response) => {
      const finalStateOnComplete = getState()
      const finalAssistantMsg = finalStateOnComplete.messages.entities[assistantMsgId]

      if (status === 'success' && finalAssistantMsg) {
        const userMsgId = finalAssistantMsg.askId
        const orderedMsgs = selectMessagesForTopic(finalStateOnComplete, topicId)
        const userMsgIndex = orderedMsgs.findIndex((m) => m.id === userMsgId)
        const contextForUsage = userMsgIndex !== -1 ? orderedMsgs.slice(0, userMsgIndex + 1) : []
        const finalContextWithAssistant = [...contextForUsage, finalAssistantMsg]

        const possibleBlockId = findBlockIdForCompletion(finalAssistantMsg)

        if (possibleBlockId) {
          const changes = {
            status: MessageBlockStatus.SUCCESS
          }
          blockManager.smartBlockUpdate(possibleBlockId, changes, blockManager.lastBlockType!, true)
        }

        const duration = Date.now() - startTime
        const content = getMainTextContent(finalAssistantMsg)

        const timeOut = duration > 30 * 1000
        // 发送长时间运行消息的成功通知
        if ((!isOnHomePage() && timeOut) || (!isFocused() && timeOut)) {
          await notificationService.send({
            id: uuid(),
            type: 'success',
            title: i18n.t('notification.assistant'),
            message: content.length > 50 ? content.slice(0, 47) + '...' : content,
            silent: false,
            timestamp: Date.now(),
            source: 'assistant',
            channel: 'system'
          })
        }

        // 更新topic的name
        void autoRenameTopic(assistant, topicId)

        // 处理usage估算
        // For OpenRouter, always use the accurate usage data from API, don't estimate
        const isOpenRouter = assistant.model?.provider === 'openrouter'
        if (
          !isOpenRouter &&
          response &&
          (response.usage?.total_tokens === 0 ||
            response?.usage?.prompt_tokens === 0 ||
            response?.usage?.completion_tokens === 0)
        ) {
          const usage = await estimateMessagesUsage({ assistant, messages: finalContextWithAssistant })
          response.usage = usage
        }
      }

      if (response && response.metrics) {
        if (response.metrics.completion_tokens === 0 && response.usage?.completion_tokens) {
          response = {
            ...response,
            metrics: {
              ...response.metrics,
              completion_tokens: response.usage.completion_tokens
            }
          }
        }
      }

      // Mark in_progress todos as completed since stream ended
      const todoCleanupIds = cleanupInProgressTodos()
      const todoBlocksToSave = todoCleanupIds
        .map((id) => getState().messageBlocks.entities[id])
        .filter(Boolean) as MessageBlock[]

      const messageUpdates = { status, metrics: response?.metrics, usage: response?.usage }
      dispatch(
        newMessagesActions.updateMessage({
          topicId,
          messageId: assistantMsgId,
          updates: messageUpdates
        })
      )
      await saveUpdatesToDB(assistantMsgId, topicId, messageUpdates, todoBlocksToSave)

      // Track token usage for agent sessions (chat sessions are tracked in fetchChatCompletion)
      if (status === 'success' && isAgentSessionTopicId(topicId)) {
        trackTokenUsage({ usage: response?.usage, model: assistant?.model, source: 'agent' })
      }

      void EventEmitter.emit(EVENT_NAMES.MESSAGE_COMPLETE, { id: assistantMsgId, topicId, status })
      logger.debug('onComplete finished')
    }
  }
}
