import { Provide, Scope, ScopeEnum, Inject, Logger } from '@midwayjs/core';
import { ILogger } from '@midwayjs/logger';
import { InjectEntityModel } from '@midwayjs/typeorm';
import { Repository, In, IsNull } from 'typeorm';
import { NetaClawMessageEntity } from '../entity/message.js';
import { LLMMessage, LLMMessageWithId } from '../plugins/plugin_entry.js';
import { AuxiliaryLLMClient } from './auxiliary_llm.js';
import { estimateMessageTokens, estimateHistoryTokens } from './token_utils.js';

/** 压缩触发原因 */
export type CompactionReason = 'manual' | 'auto';

/** 压缩结果 */
export interface CompactionResult {
  success: boolean;
  /** 被压缩的消息数 */
  compactedCount: number;
  /** 保留在历史末尾的消息数 */
  keepRecentCount: number;
  /** 生成的摘要消息完整内容 */
  summaryContent: string;
  /** 摘要消息在 DB 中的 id */
  summaryMessageId: number;
  /** 压缩前 token 估算 */
  tokensBefore: number;
  /** 压缩后 token 估算 */
  tokensAfter: number;
  /** 失败原因 */
  errorMessage?: string;
}

/** 工具结果裁剪占位符 */
const TOOL_RESULT_PLACEHOLDER = '[...truncated by compaction...]';

/** 尾部 token 预算占比（阈值 token 的百分比） */
const TAIL_TOKEN_BUDGET_RATIO = 0.3;

@Provide()
@Scope(ScopeEnum.Singleton)
export class CompactionService {
  @Logger()
  logger: ILogger;

  @Inject()
  auxiliaryLLM: AuxiliaryLLMClient;

  @InjectEntityModel(NetaClawMessageEntity)
  messageRepo: Repository<NetaClawMessageEntity>;

  /**
   * 执行压缩
   * @param sessionId 会话ID
   * @param history 带 id 的完整历史（来自 session.loadHistoryWithId()）
   * @param thresholdTokens 压缩阈值对应的 token 数（maxTokens * threshold%）
   * @param auxChannelId Agent 级辅助渠道ID（可选）
   * @param reason 触发原因
   */
  async compact(
    sessionId: string,
    history: LLMMessageWithId[],
    thresholdTokens: number,
    auxChannelId: number | undefined,
    reason: CompactionReason,
  ): Promise<CompactionResult> {
    const tokensBefore = estimateHistoryTokens(history);

    // 查询上一次摘要（用于迭代压缩）
    const previousSummary = await this.getLatestSummary(sessionId);

    // Stage 1: 裁剪工具结果（尾部 tailBudget 内的消息不裁剪）
    const pruned = this.phase1PruneToolResults(history, TAIL_TOKEN_BUDGET_RATIO, thresholdTokens);

    // Stage 2: 边界检测（用 token 预算而非固定条数）
    const { compressEnd } = this.phase2FindBoundaries(pruned, TAIL_TOKEN_BUDGET_RATIO, thresholdTokens);

    const toCompact = history.slice(0, compressEnd);
    const toKeep = history.slice(compressEnd);

    if (toCompact.length === 0) {
      return {
        success: false,
        compactedCount: 0,
        keepRecentCount: toKeep.length,
        summaryContent: '',
        summaryMessageId: 0,
        tokensBefore,
        tokensAfter: tokensBefore,
        errorMessage: '没有可压缩的历史消息',
      };
    }

    // 序列化被压缩部分
    const serialized = this.serializeMessages(pruned.slice(0, compressEnd));

    // Stage 3: LLM 摘要
    const summaryBudget = Math.max(1024, Math.floor(thresholdTokens * 0.15));
    let summaryContent: string;
    let usage: { inputTokens: number; outputTokens: number };
    try {
      const result = await this.auxiliaryLLM.summarize({
        agentAuxChannelId: auxChannelId,
        serialized,
        previousSummary: previousSummary ?? undefined,
        summaryBudget,
      });
      summaryContent = result.summary;
      usage = result.usage;
    } catch (e: any) {
      this.logger.error('[Compaction] 摘要生成失败: %s', e.message);
      return {
        success: false,
        compactedCount: toCompact.length,
        keepRecentCount: toKeep.length,
        summaryContent: '',
        summaryMessageId: 0,
        tokensBefore,
        tokensAfter: tokensBefore,
        errorMessage: `摘要生成失败: ${e.message}`,
      };
    }

    // Stage 4: 装配 — 创建摘要消息，标记旧消息
    const summaryMsg = await this.messageRepo.save({
      sessionId,
      role: 'user',
      content: summaryContent,
      isCompactionSummary: true,
      metadata: {
        reason,
        compactedIds: toCompact.map(m => m.id),
        usage,
        hasPreviousSummary: !!previousSummary,
      } as any,
    });

    const compactedAt = new Date();
    const ids = toCompact.map(m => m.id);
    if (ids.length > 0) {
      await this.messageRepo.update(
        { id: In(ids) },
        { compactedAt, compactedIntoId: summaryMsg.id },
      );
    }

    // 计算压缩后 token
    const summaryLLMMsg: LLMMessage = { role: 'user', content: summaryContent };
    const tokensAfter = estimateMessageTokens(summaryLLMMsg) + estimateHistoryTokens(toKeep);

    return {
      success: true,
      compactedCount: toCompact.length,
      keepRecentCount: toKeep.length,
      summaryContent,
      summaryMessageId: summaryMsg.id,
      tokensBefore,
      tokensAfter,
    };
  }

  /**
   * 查询该 session 最近的压缩摘要内容（用于迭代压缩）
   */
  async getLatestSummary(sessionId: string): Promise<string | null> {
    const msg = await this.messageRepo.findOne({
      where: { sessionId, isCompactionSummary: true, compactedAt: IsNull() },
      order: { createTime: 'DESC' },
    });
    return msg?.content ?? null;
  }

  /**
   * Stage 1: 裁剪工具结果
   * 从末尾反向累计 token，超过 tailBudget 后，前面的 tool 结果消息内容超过 200 字符的替换为占位符
   */
  private phase1PruneToolResults(
    messages: LLMMessageWithId[],
    tailTokenBudgetRatio: number,
    thresholdTokens: number,
  ): LLMMessageWithId[] {
    const tailBudget = Math.floor(thresholdTokens * tailTokenBudgetRatio);
    let tailTokens = 0;
    let cutIndex = messages.length;
    for (let i = messages.length - 1; i >= 0; i--) {
      tailTokens += estimateMessageTokens(messages[i]);
      if (tailTokens >= tailBudget) {
        cutIndex = i;
        break;
      }
    }
    return messages.map((m, i) => {
      if (i < cutIndex && m.role === 'tool' && (m.content?.length ?? 0) > 200) {
        return { ...m, content: TOOL_RESULT_PLACEHOLDER };
      }
      return m;
    });
  }

  /**
   * Stage 2: 边界检测
   * 从末尾反向累积 token 到 tailBudget 耗尽，确定 compressEnd 分割点
   * 同时保证不拆分工具调用对
   */
  private phase2FindBoundaries(
    messages: LLMMessageWithId[],
    tailTokenBudgetRatio: number,
    thresholdTokens: number,
  ): { compressEnd: number } {
    const tailBudget = Math.floor(thresholdTokens * tailTokenBudgetRatio);
    let tailTokens = 0;
    let boundary = messages.length;

    // 从末尾反向累积到 tailBudget
    for (let i = messages.length - 1; i >= 0; i--) {
      tailTokens += estimateMessageTokens(messages[i]);
      if (tailTokens >= tailBudget) {
        boundary = i;
        break;
      }
    }

    if (boundary <= 0 || boundary >= messages.length) return { compressEnd: 0 };

    // 向前扩展：不拆分工具调用对
    while (boundary > 0) {
      const current = messages[boundary];
      const prev = messages[boundary - 1];

      // 如果当前是 tool 消息，前一个必须是带 toolCalls 的 assistant；回退
      if (current.role === 'tool') {
        boundary -= 1;
        continue;
      }
      // 如果前一个 assistant 有未闭合的 toolCalls，回退
      if (prev.role === 'assistant' && prev.toolCalls && prev.toolCalls.length > 0) {
        const expectedIds = new Set(prev.toolCalls.map(tc => tc.id));
        for (let i = boundary; i < messages.length && expectedIds.size > 0; i++) {
          if (messages[i].role === 'tool' && messages[i].toolCallId) {
            expectedIds.delete(messages[i].toolCallId!);
          }
        }
        if (expectedIds.size > 0) {
          boundary -= 1;
          continue;
        }
      }
      break;
    }

    return { compressEnd: boundary };
  }

  /** 将消息列表序列化为文本（供 LLM 摘要使用） */
  private serializeMessages(messages: (LLMMessage | LLMMessageWithId)[]): string {
    return messages.map(m => {
      const prefix = `[${m.role}]`;
      const toolCallInfo = m.toolCalls ? ` (调用工具: ${m.toolCalls.map(t => t.name).join(',')})` : '';
      return `${prefix}${toolCallInfo} ${m.content}`;
    }).join('\n\n');
  }
}
