import type { ModelMessage, LanguageModel } from 'ai';
import { TokenCounter } from './token-counter.js';
import { prune, filterCompacted } from './prune.js';
import { compact, simpleCompact, isSummaryMessage, type CompactOptions } from './compaction.js';
import {
  type TokenUsage,
  type CompressionConfig,
  type DetailedCompressionResult,
  CompressionStatus,
  DEFAULT_COMPRESSION_CONFIG,
} from './types.js';
import type { AgentModelConfig } from '../agent/types.js';
import { getProviderRegistry } from '../provider/index.js';

/**
 * 压缩管理器
 * 统一管理对话上下文的压缩策略
 */
export class CompressionManager {
  private config: CompressionConfig;
  /** 主模型（摘要模型的后备） */
  private model: LanguageModel | null = null;
  /** 专用摘要模型（推荐使用小模型以降低成本） */
  private summaryModel: LanguageModel | null = null;
  /** 是否保护首条 user-assistant 对 */
  private protectFirstPair: boolean = true;

  constructor(config: Partial<CompressionConfig> = {}) {
    this.config = { ...DEFAULT_COMPRESSION_CONFIG, ...config };
  }

  /**
   * 设置用于生成摘要的主模型（后备）
   */
  setModel(model: LanguageModel): void {
    this.model = model;
  }

  /**
   * 设置专用摘要模型（优先使用）
   */
  setSummaryModel(model: LanguageModel): void {
    this.summaryModel = model;
  }

  /**
   * 从 Agent 配置设置摘要模型
   * @param config Agent 的模型配置
   * @param apiKey API Key
   * @param baseUrl 可选的 Base URL
   */
  setSummaryModelFromAgentConfig(
    config: AgentModelConfig,
    apiKey: string,
    baseUrl?: string
  ): void {
    const registry = getProviderRegistry();
    const provider = config.provider || 'anthropic';
    const modelName = config.model || 'claude-3-5-haiku-20241022';

    try {
      const getModel = registry.getModelFactory(provider, {
        apiKey,
        baseUrl,
      });
      this.summaryModel = getModel(modelName);
    } catch (error) {
      console.warn('[CompressionManager] Failed to create summary model:', error);
    }
  }

  /**
   * 获取用于摘要生成的模型
   * 优先使用专用摘要模型，无则使用主模型
   */
  private getSummaryModel(): LanguageModel | null {
    return this.summaryModel ?? this.model;
  }

  /**
   * 设置是否保护首条 user-assistant 对
   */
  setProtectFirstPair(protect: boolean): void {
    this.protectFirstPair = protect;
  }

  /**
   * 获取压缩选项
   */
  private getCompactOptions(): CompactOptions {
    return {
      protectFirstPair: this.protectFirstPair,
    };
  }

  /**
   * 获取当前配置
   */
  getConfig(): CompressionConfig {
    return { ...this.config };
  }

  /**
   * 更新配置
   */
  updateConfig(config: Partial<CompressionConfig>): void {
    this.config = { ...this.config, ...config };
  }

  /**
   * 计算消息数组的 token 使用情况
   */
  calculateUsage(messages: ModelMessage[]): TokenUsage {
    const input = TokenCounter.estimateMessages(messages);
    const { contextLimit, outputReserve } = this.config;
    const available = contextLimit - outputReserve;
    const usagePercent = (input / available) * 100;

    return {
      input,
      contextLimit,
      available,
      usagePercent: Math.min(usagePercent, 100),
    };
  }

  /**
   * 检查是否需要压缩（超过溢出阈值）
   */
  shouldCompress(messages: ModelMessage[]): boolean {
    const usage = this.calculateUsage(messages);
    return usage.usagePercent >= this.config.overflowThreshold * 100;
  }

  /**
   * 检查是否溢出（超过可用空间）
   */
  isOverflow(messages: ModelMessage[]): boolean {
    const usage = this.calculateUsage(messages);
    return usage.input >= usage.available;
  }

  /**
   * 执行 prune 策略
   */
  prune(messages: ModelMessage[]): { messages: ModelMessage[]; freedTokens: number } {
    return prune(messages, this.config);
  }

  /**
   * 执行 compaction 策略
   */
  async compact(messages: ModelMessage[]): Promise<DetailedCompressionResult> {
    const summaryModel = this.getSummaryModel();
    if (summaryModel) {
      return compact(messages, summaryModel, this.config, this.getCompactOptions());
    }
    // 没有模型时使用简单压缩
    return simpleCompact(messages, this.config, this.getCompactOptions());
  }

  /**
   * 自动压缩：先 prune，不够再 compact
   */
  async compress(messages: ModelMessage[]): Promise<DetailedCompressionResult> {
    // 检查是否需要压缩
    if (!this.shouldCompress(messages)) {
      return {
        messages,
        freedTokens: 0,
        type: 'none',
        status: CompressionStatus.NOOP,
      };
    }

    let result = [...messages];
    let totalFreed = 0;
    let type: DetailedCompressionResult['type'] = 'none';

    // 第一步：尝试 prune
    const pruneResult = this.prune(result);
    if (pruneResult.freedTokens > 0) {
      result = pruneResult.messages;
      totalFreed += pruneResult.freedTokens;
      type = 'prune';
    }

    // 检查是否还需要进一步压缩
    if (this.shouldCompress(result)) {
      // 第二步：执行 compaction
      const compactResult = await this.compact(result);

      if (compactResult.status === CompressionStatus.SUCCESS && compactResult.freedTokens > 0) {
        result = compactResult.messages;
        totalFreed += compactResult.freedTokens;
        type = pruneResult.freedTokens > 0 ? 'both' : 'compaction';

        return {
          messages: result,
          freedTokens: totalFreed,
          type,
          status: CompressionStatus.SUCCESS,
          originalTokens: compactResult.originalTokens,
          summaryTokens: compactResult.summaryTokens,
        };
      }

      // compaction 失败，返回失败状态
      if (compactResult.status !== CompressionStatus.NOOP) {
        return {
          messages, // 返回原消息
          freedTokens: 0,
          type: 'none',
          status: compactResult.status,
          error: compactResult.error,
        };
      }
    }

    return {
      messages: result,
      freedTokens: totalFreed,
      type,
      status: totalFreed > 0 ? CompressionStatus.SUCCESS : CompressionStatus.NOOP,
    };
  }

  /**
   * 强制压缩（用于 /compact 命令）
   * 无论是否达到阈值都执行压缩
   */
  async forceCompress(messages: ModelMessage[]): Promise<DetailedCompressionResult> {
    // 消息数量太少时不压缩（至少需要 4 条消息）
    if (messages.length <= 4) {
      return {
        messages,
        freedTokens: 0,
        type: 'none',
        status: CompressionStatus.NOOP,
      };
    }

    let result = [...messages];
    let totalFreed = 0;
    let type: DetailedCompressionResult['type'] = 'none';

    // 先尝试 prune（使用强制配置）
    const pruneConfig: CompressionConfig = {
      ...this.config,
      pruneMinimum: 0,
      pruneProtect: Math.min(10_000, TokenCounter.estimateMessages(messages) / 4),
    };

    const pruneResult = prune(result, pruneConfig);

    if (pruneResult.freedTokens > 0) {
      result = pruneResult.messages;
      totalFreed += pruneResult.freedTokens;
      type = 'prune';
    }

    // 强制 compaction：使用强制配置
    const summaryModel = this.getSummaryModel();
    const forceConfig: CompressionConfig = {
      ...this.config,
      pruneProtect: 0, // 强制模式：不保护任何 tokens
    };
    // 强制模式不保护首条消息对
    const forceOptions = { protectFirstPair: false };

    if (summaryModel) {
      const compactResult = await compact(result, summaryModel, forceConfig, forceOptions);

      if (compactResult.status === CompressionStatus.SUCCESS && compactResult.freedTokens > 0) {
        result = compactResult.messages;
        totalFreed += compactResult.freedTokens;
        type = type === 'prune' ? 'both' : 'compaction';

        return {
          messages: result,
          freedTokens: totalFreed,
          type,
          status: CompressionStatus.SUCCESS,
          originalTokens: compactResult.originalTokens,
          summaryTokens: compactResult.summaryTokens,
        };
      }

      // AI 压缩失败，回退到简单压缩
      if (compactResult.status !== CompressionStatus.NOOP) {
        const simpleResult = simpleCompact(result, forceConfig, forceOptions);
        if (simpleResult.freedTokens > 0) {
          result = simpleResult.messages;
          totalFreed += simpleResult.freedTokens;
          type = type === 'prune' ? 'both' : 'compaction';
        }
      }
    } else {
      // 没有模型，使用简单压缩
      const simpleResult = simpleCompact(result, forceConfig, forceOptions);
      if (simpleResult.freedTokens > 0) {
        result = simpleResult.messages;
        totalFreed += simpleResult.freedTokens;
        type = type === 'prune' ? 'both' : 'compaction';
      }
    }

    return {
      messages: result,
      freedTokens: totalFreed,
      type,
      status: totalFreed > 0 ? CompressionStatus.SUCCESS : CompressionStatus.NOOP,
    };
  }

  /**
   * 过滤已压缩的内容（发送给模型前调用）
   */
  filterCompacted(messages: ModelMessage[]): ModelMessage[] {
    return filterCompacted(messages);
  }

  /**
   * 检查消息是否为摘要消息
   */
  isSummaryMessage(message: ModelMessage): boolean {
    return isSummaryMessage(message);
  }

  /**
   * 格式化 token 使用情况（用于 CLI 显示）
   */
  formatUsage(messages: ModelMessage[]): string {
    const usage = this.calculateUsage(messages);
    const used = TokenCounter.format(usage.input);
    const limit = TokenCounter.format(usage.available);
    const percent = usage.usagePercent.toFixed(0);
    return `${used}/${limit} (${percent}%)`;
  }
}

// 导出单例（可选使用）
export const compressionManager = new CompressionManager();
