import { generateText, type ModelMessage, type LanguageModel } from 'ai';
import { TokenCounter } from './token-counter.js';
import {
  SUMMARY_MARKER,
  CompressionStatus,
  type CompressionConfig,
  type DetailedCompressionResult,
  DEFAULT_COMPRESSION_CONFIG,
} from './types.js';

/**
 * 摘要生成系统提示词
 */
const COMPACTION_SYSTEM_PROMPT = `你是一个专门生成对话摘要的助手。你的任务是将对话历史压缩成一个简洁但信息完整的摘要。

摘要应该包含：
1. 已完成的工作和关键结果
2. 当前正在进行的任务
3. 涉及的重要文件和代码
4. 用户的关键需求和约束
5. 下一步需要做的事情

要求：
- 保留关键技术细节（文件路径、函数名、配置等）
- 使用简洁的列表格式
- 不要遗漏重要信息
- 使用中文回复`;

/**
 * 摘要生成用户提示词
 */
const COMPACTION_USER_PROMPT = `请总结上面的对话。这个摘要将是对话继续时唯一可用的上下文，所以要保留关键信息，包括：完成了什么、正在进行的工作、涉及的文件、下一步计划、以及用户的关键需求或约束。要简洁但足够详细，以便工作可以无缝继续。`;

/**
 * 检查消息是否为摘要消息
 */
export function isSummaryMessage(message: ModelMessage): boolean {
  if (typeof message.content === 'string') {
    return message.content.includes(SUMMARY_MARKER);
  }
  if (Array.isArray(message.content)) {
    return message.content.some(
      (part) =>
        typeof part === 'object' &&
        'text' in part &&
        typeof part.text === 'string' &&
        part.text.includes(SUMMARY_MARKER)
    );
  }
  return false;
}

/**
 * 创建摘要消息
 */
function createSummaryMessage(summary: string): ModelMessage {
  return {
    role: 'assistant',
    content: `${SUMMARY_MARKER}\n## 对话摘要\n\n${summary}\n${SUMMARY_MARKER}`,
  };
}

/**
 * 验证摘要结果
 * @returns 验证结果，包含状态和 token 数
 */
function validateSummary(
  summary: string,
  originalTokens: number
): { valid: boolean; status: CompressionStatus; summaryTokens: number } {
  // 1. 检测空摘要
  if (!summary || summary.trim().length === 0) {
    return {
      valid: false,
      status: CompressionStatus.FAILED_EMPTY_SUMMARY,
      summaryTokens: 0,
    };
  }

  // 2. 检测 token 膨胀（摘要比原消息还大）
  const summaryTokens = TokenCounter.estimateText(summary);
  if (summaryTokens >= originalTokens) {
    return {
      valid: false,
      status: CompressionStatus.FAILED_TOKEN_INFLATED,
      summaryTokens,
    };
  }

  return {
    valid: true,
    status: CompressionStatus.SUCCESS,
    summaryTokens,
  };
}

/**
 * 查找首个 user-assistant 对
 * 用于首条消息保护，确保上下文连贯性
 */
function findFirstUserAssistantPair(messages: ModelMessage[]): {
  firstPair: ModelMessage[];
  rest: ModelMessage[];
} | null {
  // 找到第一个 user 消息
  let userIndex = -1;
  for (let i = 0; i < messages.length; i++) {
    if (messages[i].role === 'user') {
      userIndex = i;
      break;
    }
  }

  if (userIndex === -1) return null;

  // 找到紧随其后的 assistant 消息
  let assistantIndex = -1;
  for (let i = userIndex + 1; i < messages.length; i++) {
    if (messages[i].role === 'assistant') {
      assistantIndex = i;
      break;
    }
  }

  if (assistantIndex === -1) {
    // 只有 user 消息，没有 assistant 响应
    return {
      firstPair: messages.slice(0, userIndex + 1),
      rest: messages.slice(userIndex + 1),
    };
  }

  return {
    firstPair: messages.slice(0, assistantIndex + 1),
    rest: messages.slice(assistantIndex + 1),
  };
}

/**
 * Compaction 选项
 */
export interface CompactOptions {
  /** 是否保护首个 user-assistant 对 */
  protectFirstPair?: boolean;
}

/**
 * Compaction 策略：使用 AI 生成对话摘要
 *
 * 逻辑：
 * 1. 将历史消息（排除最近保护的部分）发送给 AI
 * 2. AI 生成摘要
 * 3. 用摘要消息替换旧消息
 * 4. 保留最近的消息不变
 *
 * @param messages 消息数组
 * @param model 语言模型
 * @param config 压缩配置
 * @param options 压缩选项
 * @returns 详细压缩结果
 */
export async function compact(
  messages: ModelMessage[],
  model: LanguageModel,
  config: CompressionConfig = DEFAULT_COMPRESSION_CONFIG,
  options: CompactOptions = {}
): Promise<DetailedCompressionResult> {
  const { pruneProtect } = config;
  const { protectFirstPair = true } = options;

  // 首条消息保护：分离首个 user-assistant 对
  let protectedFirst: ModelMessage[] = [];
  let compressibleMessages = messages;

  if (protectFirstPair) {
    const firstPairResult = findFirstUserAssistantPair(messages);
    if (firstPairResult && firstPairResult.rest.length > 0) {
      protectedFirst = firstPairResult.firstPair;
      compressibleMessages = firstPairResult.rest;
    }
  }

  // 计算需要保护的消息数量（从可压缩部分的末尾算起）
  let protectedTokens = 0;
  let protectedCount = 0;

  for (let i = compressibleMessages.length - 1; i >= 0; i--) {
    const tokens = TokenCounter.estimateMessage(compressibleMessages[i]);
    if (protectedTokens + tokens > pruneProtect) {
      break;
    }
    protectedTokens += tokens;
    protectedCount++;
  }

  // 确保至少保护最后 2 条消息（除非 pruneProtect 为 0，表示强制压缩模式）
  if (pruneProtect > 0) {
    protectedCount = Math.max(protectedCount, 2);
  } else {
    // 强制压缩模式：至少保护 1 条消息
    protectedCount = Math.max(protectedCount, 1);
  }

  // 分割消息：需要压缩的部分 vs 保护的部分
  const toCompact = compressibleMessages.slice(0, compressibleMessages.length - protectedCount);
  const toKeep = compressibleMessages.slice(compressibleMessages.length - protectedCount);

  // 如果没有需要压缩的消息，直接返回
  if (toCompact.length === 0) {
    return {
      messages,
      freedTokens: 0,
      type: 'none',
      status: CompressionStatus.NOOP,
    };
  }

  // 检查是否已有摘要消息
  const existingSummaryIndex = toCompact.findIndex(isSummaryMessage);
  const messagesForSummary =
    existingSummaryIndex >= 0 ? toCompact.slice(existingSummaryIndex) : toCompact;

  // 计算压缩前的 tokens
  const originalTokens = TokenCounter.estimateMessages(toCompact);

  try {
    // 调用 AI 生成摘要
    const result = await generateText({
      model,
      system: COMPACTION_SYSTEM_PROMPT,
      messages: [
        ...messagesForSummary,
        {
          role: 'user',
          content: COMPACTION_USER_PROMPT,
        },
      ],
      maxOutputTokens: 2000,
    });

    // 验证摘要结果
    const validation = validateSummary(result.text, originalTokens);

    if (!validation.valid) {
      console.warn(`摘要验证失败: ${validation.status}`);
      return {
        messages,
        freedTokens: 0,
        type: 'none',
        status: validation.status,
        originalTokens,
        summaryTokens: validation.summaryTokens,
      };
    }

    const summaryMessage = createSummaryMessage(result.text);
    const summaryTokens = TokenCounter.estimateMessage(summaryMessage);
    const freedTokens = originalTokens - summaryTokens;

    // 返回：首条保护 + 摘要 + 末尾保护的消息
    return {
      messages: [...protectedFirst, summaryMessage, ...toKeep],
      freedTokens,
      type: 'compaction',
      status: CompressionStatus.SUCCESS,
      originalTokens,
      summaryTokens,
    };
  } catch (error) {
    console.error('生成摘要失败:', error);
    return {
      messages,
      freedTokens: 0,
      type: 'none',
      status: CompressionStatus.FAILED_ERROR,
      error: error instanceof Error ? error.message : String(error),
      originalTokens,
    };
  }
}

/**
 * 简单压缩：不使用 AI，直接截断旧消息
 * 用于没有模型可用或快速压缩的场景
 */
export function simpleCompact(
  messages: ModelMessage[],
  config: CompressionConfig = DEFAULT_COMPRESSION_CONFIG,
  options: CompactOptions = {}
): DetailedCompressionResult {
  const { pruneProtect } = config;
  const { protectFirstPair = true } = options;

  // 首条消息保护：分离首个 user-assistant 对
  let protectedFirst: ModelMessage[] = [];
  let compressibleMessages = messages;

  if (protectFirstPair) {
    const firstPairResult = findFirstUserAssistantPair(messages);
    if (firstPairResult && firstPairResult.rest.length > 0) {
      protectedFirst = firstPairResult.firstPair;
      compressibleMessages = firstPairResult.rest;
    }
  }

  // 计算需要保留的消息
  let keptTokens = 0;
  let keepFromIndex = compressibleMessages.length;

  for (let i = compressibleMessages.length - 1; i >= 0; i--) {
    const tokens = TokenCounter.estimateMessage(compressibleMessages[i]);
    if (keptTokens + tokens > pruneProtect) {
      break;
    }
    keptTokens += tokens;
    keepFromIndex = i;
  }

  // 确保至少保留最后 N 条消息（强制模式下保留 1 条，否则保留 2 条）
  const minKeep = pruneProtect > 0 ? 2 : 1;
  keepFromIndex = Math.min(keepFromIndex, compressibleMessages.length - minKeep);

  const removed = compressibleMessages.slice(0, keepFromIndex);
  const kept = compressibleMessages.slice(keepFromIndex);

  if (removed.length === 0) {
    return {
      messages,
      freedTokens: 0,
      type: 'none',
      status: CompressionStatus.NOOP,
    };
  }

  // 创建简单摘要
  const simpleSummary: ModelMessage = {
    role: 'assistant',
    content: `${SUMMARY_MARKER}\n[对话历史已压缩，共移除 ${removed.length} 条消息]\n${SUMMARY_MARKER}`,
  };

  const freedTokens = TokenCounter.estimateMessages(removed);

  return {
    messages: [...protectedFirst, simpleSummary, ...kept],
    freedTokens,
    type: 'compaction',
    status: CompressionStatus.SUCCESS,
  };
}
