import { ref, type Ref } from "vue";
import type { ChatSendOptions } from "./Chat";
import type OpenAI from "openai";
import { ChatMessage } from "../model/Message";
import type { ChatSessionManager } from "../composables/useChatSession";
import { useSlideWindowMemoryStrategy } from "./strategy/ContextSlideWindowStrategy";
import Agent from "@/api/agent/Agent";
import AgentChatFile from "@/api/agent/AgentChatFile";
import { requireNotNull } from "@imengyu/imengyu-utils";

export type ContextMemorySetting = 'slide-window'|'long-summary';
export type ContextMemoryConfig = {
  slideWindowShortRelatedLookback?: number;
  slideWindowShortUnrelatedLookback?: number;
  slideWindowWindowMinRounds?: number;
  slideWindowWindowMaxRounds?: number;
  slideWindowHeadRounds?: number;
  slideWindowTailRounds?: number;
  longSummaryReservedRatio?: number;
  longSummaryShortRounds?: number;
  longSummaryStageRoundInterval?: number;
  longSummaryLongMemoryMaxItems?: number;
  minContextTokens?: number;
  /**
   * 推荐上下文token数量
   * @default 32000
   */
  prefferTokens?: number;
};

export type RoundedMessages = {
  userMessages: ChatMessage[];
  aiMessages: ChatMessage[];
};

export type InnerContextManager =
{
  estimateMessageTokens: (message: ChatMessage) => number;
  estimateMessagesTokens: (messages: ChatMessage[]) => number;
  estimateRoundedMessagesTokens: (roundedMessages: RoundedMessages[]) => number;
  getContextTokenBudget: (sendOptions: ChatSendOptions) => number;
  pickMessagesWithBudgetFromTail: (list: ChatMessage[], maxTokens: number) => ChatMessage[];
  getRoundedMessages: () => RoundedMessages[];
  cutRoundedMessagesFromId: (groupedMessages: RoundedMessages[], id: number) => RoundedMessages[];
  roundedMessagesToMessages: (roundedMessages: RoundedMessages[]) => ChatMessage[];
};

/**
 * 聊天上下文管理
 */
export function useChatContext(options: {
  message: Ref<ChatMessage[]>,
  /**
   * 上下文记忆设置
   * @default 'slide-window'
   * 
   * * slide-window: 滑动窗口长时记忆 + 首尾保留，尽量记忆所有消息，通常用于对话讨论模式, 同时为了节省会尝试判断当前问题与上几条问题是否有关联，如果无关联，则不记忆。
   * * long-summary: 分层记忆（三层）：短期（最近轮次）/中期（阶段摘要）/长期（重要信息）。
   */
  contextMemorySetting?: ContextMemorySetting;
  /**
   * 上下文记忆参数配置
   */
  contextMemoryConfig?: ContextMemoryConfig,
  /**
   * 会话管理器
   */
  sessionManager: ChatSessionManager,
  /**
   * 程序追加消息，不保存到数据库
   */
  onGetAppendMessages?: () => ChatMessage[];
}) {
  const contextMemorySetting = options.contextMemorySetting || 'slide-window';
  const contextMemoryConfig = {
    prefferTokens: 32000,
    shortRelatedLookback: 8,
    shortUnrelatedLookback: 0,
    slideWindowWindowMinRounds: 8,
    slideWindowWindowMaxRounds: 10,
    slideWindowHeadRounds: 1,
    slideWindowTailRounds: 1,
    longSummaryReservedRatio: 0.35,
    longSummaryShortRounds: 6,
    longSummaryStageRoundInterval: 6,
    longSummaryLongMemoryMaxItems: 12,
    minContextTokens: 512,
    ...(options.contextMemoryConfig || {}),
  } as Required<ContextMemoryConfig>;
  const sessionManager = options.sessionManager;

  /**
   * 估算消息token数量
   * @param message 
   * @returns 
   */
  function estimateMessageTokens(message: ChatMessage): number {
    if (message.type !== 'text')
      return 24;
    const text = `${message.content || ''}\n${message.reasoningContent || ''}\n${message.extra || ''}`;
    const cjkCount = (text.match(/[\u4e00-\u9fff]/g) || []).length;
    const nonCjkText = text.replace(/[\u4e00-\u9fff]/g, '');
    const latinWordCount = (nonCjkText.match(/[A-Za-z0-9_]+/g) || []).length;
    const punctuationCount = (nonCjkText.match(/[^\w\s]/g) || []).length;
    // 粗略估算：中文按 1 字≈1 token，英文按 1 单词≈1 token，外加协议开销
    return Math.max(8, cjkCount + latinWordCount + Math.ceil(punctuationCount / 2) + 6);
  }
  /**
   * 估算消息列表token数量
   * @param messages 消息列表
   * @returns 
   */
  function estimateMessagesTokens(messages: ChatMessage[]): number {
    return messages.reduce((sum, item) => sum + estimateMessageTokens(item), 0);
  }
  /**
   * 估算消息列表token数量
   * @param messages 消息列表
   * @returns 
   */
  function estimateRoundedMessagesTokens(roundedMessages: RoundedMessages[]): number {
    return roundedMessages.reduce((sum, item) => sum + estimateMessagesTokens(item.userMessages) + 
      estimateMessagesTokens(item.aiMessages), 0);
  }
  /**
   * 获取上下文token预算(默认为模型最大token的85%)
   * @param sendOptions 发送选项
   * @returns 
   */
  function getContextTokenBudget(sendOptions: ChatSendOptions): number {
    const modelMaxTokens = sendOptions.modelInfo?.max_tokens || 0;
    if (modelMaxTokens <= 0)
      return 16000;
    return Math.max(contextMemoryConfig.minContextTokens, Math.floor(modelMaxTokens * 0.85));
  }
  /**
   * 从消息列表末尾开始，按token数量从大到小选择消息，直到超过预算为止
   * @param list 消息列表
   * @param maxTokens 预算
   * @returns 消息列表
   */
  function pickMessagesWithBudgetFromTail(list: ChatMessage[], maxTokens: number): ChatMessage[] {
    const selected: ChatMessage[] = [];
    let used = 0;
    for (let i = list.length - 1; i >= 0; i--) {
      const item = list[i];
      const itemTokens = estimateMessageTokens(item);
      if (selected.length > 0 && used + itemTokens > maxTokens)
        break;
      selected.push(item);
      used += itemTokens;
    }
    return selected.reverse();
  }

  /**
   * 根据用户与AI消息的回复关系，将消息分组。
   * 
   * 消息通过parentId分组
   */
  function getRoundedMessages() {
    const messages = options.message.value;
    if (options.onGetAppendMessages)
      messages.push(...options.onGetAppendMessages());
    const groupedMessages: RoundedMessages[] = messages
      .filter(m => m.parentId === 0 && m.isUser)
      .map(m => ({
        userMessages: [m].concat(messages.filter(k => k.parentId === m.id && k.isUser && (k.content))),
        aiMessages: messages.filter(k => k.parentId === m.id && (k.isAssistant || k.isTool) && (k.content || k.reasoningContent)),
      }));
    return groupedMessages;
  }
  /**
   * 从消息分组中截取指定ID开始至末尾的消息分组，包含该消息
   * @param groupedMessages 消息分组
   * @param id 消息ID
   * @returns 消息分组
   */
  function cutRoundedMessagesFromId(groupedMessages: RoundedMessages[], id: number) {
    for (let i = groupedMessages.length - 1; i >= 0; i--) {
      const group = groupedMessages[i];
      if (group.userMessages.some(m => m.id === id) || group.aiMessages.some(m => m.id === id)) {
        return groupedMessages.slice(i);
      }
    }
    return groupedMessages;
  }
  /**
   * 将消息分组转换为消息列表
   * @param roundedMessages 消息分组
   * @returns 消息列表
   */
  function roundedMessagesToMessages(roundedMessages: RoundedMessages[]): ChatMessage[] {
    return roundedMessages.flatMap(group => [...group.userMessages, ...group.aiMessages]);
  }

  const contextManager = {
    estimateMessageTokens,
    estimateMessagesTokens,
    estimateRoundedMessagesTokens,
    getContextTokenBudget,
    pickMessagesWithBudgetFromTail,
    getRoundedMessages,
    cutRoundedMessagesFromId,
    roundedMessagesToMessages,
  };

  /**
   * 估算上下文token数量
   * @default 0
   */
  const estimatedTokens = ref(0);
  /**
   * 估算上下文token使用量(百分比)
   * @default 0
   */
  const estimatedTokenUseage = ref(0);
  /**
   * 估算上下文token最大数量
   * @default 0
   */
  const estimatedTokenMax = ref(0);

  /**
   * 将会话消息转换为AI消息
   * 
   * 2. 如果上下文记忆设置为slide-window，则采用滑动窗口（8–10 轮）+ 首尾保留 模式。
   *    根据当前问题与上几条问题是否有关联，有则只记忆最近几条消息。如果没有关联，则不记忆。
   * 3. 如果上下文记忆设置为long-summary，则记忆所有消息，如果检测超出上下文，
   *    则会在超出位置总结前面信息的摘要，并在截断位置之前创建摘要信息；
   *    只传递到摘要为止，不传递到超出部分。
   * 
   * @param sendOptions 发送选项
   * @param currentUserMessageIds 当前用户消息ID列表
   * @param limitHistoryStartAtMessageId 限制历史消息的起始消息ID，包含该消息
   * @returns AI消息
   */
  async function convertMessagesToAi(sendOptions: ChatSendOptions, currentUserMessageIds: number[], limitHistoryStartAtMessageId?: number) {
    const maxTokens = getContextTokenBudget(sendOptions);
    const modelInfo = requireNotNull(sendOptions.modelInfo, '模型信息不能为空');

    //截取基础信息
    let roundedMessages = getRoundedMessages();
    if (limitHistoryStartAtMessageId)
      roundedMessages = cutRoundedMessagesFromId(roundedMessages, limitHistoryStartAtMessageId);

    //估算消息列表token数量，供界面显示
    estimatedTokenMax.value = maxTokens;
    estimatedTokens.value = estimateRoundedMessagesTokens(roundedMessages);
    estimatedTokenUseage.value = Math.round((estimatedTokens.value / maxTokens) * 100);

    //根据上下文记忆策略，获取用户消息和系统消息
    let systemMessages: ChatMessage[] = [];

    switch (contextMemorySetting) {
      case 'slide-window':
        roundedMessages = await useSlideWindowMemoryStrategy({
          roundedMessages,
          contextMemoryConfig,
          maxTokens,
          contextManager,
        });
        break;
      case 'long-summary': {
        //TODO: 实现长时记忆
        /* userMessages = layeredMemoryResult.userMessages;
        systemMessages = layeredMemoryResult.systemMessages;
        await updateSessionSummary(layeredMemoryResult.summary, layeredMemoryResult.shortMemory, layeredMemoryResult.longMemory); */
        break;
      }
      default:
        throw new Error(`Unsupported context memory setting: ${contextMemorySetting}`);
    }

    //添加会话偏好
    const currentSession = sessionManager.currentSession.value;
    if (currentSession) {
      if (currentSession.summary)
        systemMessages.push(ChatMessage.createSystem(`[会话摘要]：${currentSession.summary}`));
      if (currentSession.shortMemory)
        systemMessages.push(ChatMessage.createSystem(`[短期偏好]：${currentSession.shortMemory}`));
      if (currentSession.longMemory)
        systemMessages.push(ChatMessage.createSystem(`[长期偏好]：${currentSession.longMemory}`));
    }

    // 合并一轮中用户消息多条多模态消息为单条合并。
    // 如果用户上传了文件，则需要将文件上传至AI系统中，并获取文件ID，然后在系统提示词中添加文件ID。
    const finalUserMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [];
    const fileIds: string[] = [];
    const finalMergeSystemMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [];
    for (let i = 0; i < roundedMessages.length; i++) {
      const group = roundedMessages[i];
      const userMessages = group.userMessages.map(m => m.toAi() as OpenAI.Chat.ChatCompletionMessageParam | undefined).filter(item => !!item);
      const hasMulitMedia = group.userMessages.some(m => m.isMedia);
      const finalItem: OpenAI.Chat.ChatCompletionMessageParam = {
        role: 'user',
        content: '',
      }
      if (hasMulitMedia) {
        for (const m of group.userMessages.filter(m => m.type === 'input_file' && m.files?.[0])) {
          const file = m.files?.[0]!;
          m.file_id = (await AgentChatFile.processChatFileToAi(file.id, sendOptions.model)).requireData();
          fileIds.push(m.file_id);
        }
        finalItem.content = group.userMessages.map(m => {
          const item : Record<string, any> = {
            ...m.toAi()
          };
          delete (item as any).role;
          if (!item.type && item.content) {
            item.type = 'text';
            item.text = item.content;
            delete item.content;
          }
          return item as OpenAI.Chat.ChatCompletionContentPart;
        });
        //因为阿里百炼大模型不支持在消息内容中传递文件ID，所以需要将文件ID传递到系统提示词中。此处移除消息内容中的文件ID。
        if (modelInfo.docunment_transfer_type !== 'in_message_content') {
          finalItem.content = finalItem.content.filter(m => (m as any).type !== 'input_file');
          if (finalItem.content.length == 1)
            finalItem.content = (finalItem.content[0] as any).content;
        } 
      } else {
        finalItem.content = userMessages.map(m => m.content).join('\n');
      }
      finalUserMessages.push(finalItem);
      finalUserMessages.push(...group.aiMessages.map(m => m.toAi() as OpenAI.Chat.ChatCompletionMessageParam | undefined).filter(item => !!item));
    }

    //将文件ID传递到系统提示词中。
    if (modelInfo.docunment_transfer_type === 'in_system_prompt') {
      finalMergeSystemMessages.push({
        role: 'system',
        content: `${fileIds.map(id => `${modelInfo.docunment_transfer_key}${id}`).join(',')}`,
      } as OpenAI.Chat.ChatCompletionSystemMessageParam);
    }
    return {
      userMessages: finalUserMessages,
      systemMessages: 
        (systemMessages
          .map((item) => item.toAi() as OpenAI.Chat.ChatCompletionMessageParam | undefined)
          .filter((item) => !!item) as OpenAI.Chat.ChatCompletionMessageParam[]
        )
        .concat(finalMergeSystemMessages),
    }
  }

  function estimateTokenUseage(sendOptions: ChatSendOptions) {
    //估算消息列表token数量，供界面显示
    const maxTokens = getContextTokenBudget(sendOptions);
    const roundedMessages = getRoundedMessages();
    estimatedTokenMax.value = maxTokens;
    estimatedTokens.value = estimateRoundedMessagesTokens(roundedMessages);
    estimatedTokenUseage.value = Math.round((estimatedTokens.value / maxTokens) * 100);
  }

  async function updateSessionSummary(summary: string, shortMemory: string, longMemory: string) {
    const currentSession = options.sessionManager.currentSession.value;
    if (!currentSession)
      return;
    if (currentSession.summary === summary &&
      currentSession.shortMemory === shortMemory &&
      currentSession.longMemory === longMemory)
      return;
    currentSession.summary = summary;
    currentSession.shortMemory = shortMemory;
    currentSession.longMemory = longMemory;
    await options.sessionManager.updateSession(currentSession);
  }

  return {
    estimatedTokens,
    estimatedTokenUseage,
    estimatedTokenMax,
    convertMessagesToAi,
    estimateTokenUseage,
    ...contextManager,
  };
}

export type ContextManager = ReturnType<typeof useChatContext>;