'use strict';

const { estimateTokens, estimateMessagesTokens } = require('../context/token-counter');
const { recordContextCompaction } = require('./execution-trace');
const { newId, textFromContent } = require('./model-message');

const DEFAULT_CONTEXT_WINDOW = 200000;
const DEFAULT_THRESHOLD = 0.75;
const DEFAULT_TAIL_TOKEN_BUDGET = 40000;
const DEFAULT_KEEP_RECENT_USER_TURNS = 4;

// Pruning: the cheap compaction layer. Triggers earlier than summarization
// and only truncates OLD tool outputs — no LLM call, no information about
// decisions lost (the agent can re-run a tool if it needs the full output).
const DEFAULT_PRUNE_THRESHOLD = 0.60;
const DEFAULT_PRUNE_PROTECT_TOKENS = 40000;
const DEFAULT_PRUNE_MAX_TOOL_OUTPUT_CHARS = 2000;
const DEFAULT_PRUNE_PROTECTED_TOOLS = Object.freeze(['skill', 'update_todos', 'task', 'ask_user']);
const PRUNED_MARKER_RE = /\[pruned \d+ chars/;

// Chunked summarization: a single LLM call used to see at most the first
// 60k chars of the head — everything past that silently vanished from the
// summary. Chunk and chain instead; if the head is still too big, keep the
// oldest chunks (original task framing) plus the newest and note the gap.
const SUMMARY_CHUNK_CHARS = 60000;
const SUMMARY_MAX_CHUNKS = 4;

class CompactionService {
  constructor({
    provider,
    model,
    contextWindow = DEFAULT_CONTEXT_WINDOW,
    threshold = DEFAULT_THRESHOLD,
    tailTokenBudget = DEFAULT_TAIL_TOKEN_BUDGET,
    keepRecentUserTurns = DEFAULT_KEEP_RECENT_USER_TURNS,
    pruneThreshold = DEFAULT_PRUNE_THRESHOLD,
    pruneProtectTokens = DEFAULT_PRUNE_PROTECT_TOKENS,
    pruneMaxToolOutputChars = DEFAULT_PRUNE_MAX_TOOL_OUTPUT_CHARS,
    pruneProtectedTools = DEFAULT_PRUNE_PROTECTED_TOOLS,
    now,
  } = {}) {
    this.provider = provider || null;
    this.model = model || provider?.model || '';
    this.contextWindow = contextWindow;
    this.threshold = threshold;
    this.tailTokenBudget = tailTokenBudget;
    this.keepRecentUserTurns = keepRecentUserTurns;
    this.pruneThreshold = pruneThreshold;
    this.pruneProtectTokens = pruneProtectTokens;
    this.pruneMaxToolOutputChars = pruneMaxToolOutputChars;
    this.pruneProtectedTools = new Set(pruneProtectedTools || []);
    this.now = now || (() => new Date().toISOString());
  }

  _measuredTokens({ messages = [], systemTokens = 0, usage = null } = {}) {
    const actualInput = Number.isFinite(usage?.input) ? usage.input
      : Number.isFinite(usage?.inputTokens) ? usage.inputTokens
        : null;
    return actualInput == null
      ? systemTokens + estimateMessagesTokens(messages)
      : actualInput;
  }

  shouldCompact({ messages = [], systemTokens = 0, usage = null, contextWindow = this.contextWindow } = {}) {
    return this._measuredTokens({ messages, systemTokens, usage }) >= contextWindow * this.threshold;
  }

  shouldPrune({ messages = [], systemTokens = 0, usage = null, contextWindow = this.contextWindow } = {}) {
    return this._measuredTokens({ messages, systemTokens, usage }) >= contextWindow * this.pruneThreshold;
  }

  /**
   * Cheap compaction layer: truncate tool outputs in OLD turns, keeping the
   * recent `protectTokens` of conversation untouched. Pure — returns new
   * message objects for anything modified; idempotent (already-pruned blocks
   * carry a marker and are skipped). No LLM call.
   */
  prune(messages = [], {
    protectTokens = this.pruneProtectTokens,
    maxToolOutputChars = this.pruneMaxToolOutputChars,
    protectedTools = this.pruneProtectedTools,
  } = {}) {
    const tokensBefore = estimateMessagesTokens(messages);
    if (!Array.isArray(messages) || messages.length === 0) {
      return { pruned: false, messages: messages || [], prunedBlocks: 0, tokensBefore, tokensAfter: tokensBefore };
    }

    const protectedSet = protectedTools instanceof Set ? protectedTools : new Set(protectedTools || []);
    const toolNameById = new Map();
    for (const msg of messages) {
      if (msg?.role !== 'assistant' || !Array.isArray(msg.content)) continue;
      for (const block of msg.content) {
        if (block?.type === 'tool_use' && block.id) toolNameById.set(block.id, block.name || '');
      }
    }

    // Everything from protectStartIndex onward is untouchable: the most
    // recent messages whose cumulative size fits the protect budget. The
    // message that crosses the budget is old enough to prune.
    let protectStartIndex = messages.length;
    let protectedTokens = 0;
    for (let i = messages.length - 1; i >= 0; i--) {
      protectedTokens += estimateMessagesTokens([messages[i]]);
      if (protectedTokens > protectTokens) break;
      protectStartIndex = i;
    }

    let prunedBlocks = 0;
    const out = messages.map((msg, index) => {
      if (index >= protectStartIndex) return msg;
      if (!msg || !Array.isArray(msg.content)) return msg;
      let changed = false;
      const content = msg.content.map((block) => {
        if (block?.type !== 'tool_result') return block;
        const toolName = toolNameById.get(block.tool_use_id) || '';
        if (protectedSet.has(toolName)) return block;
        const text = toolResultText(block.content);
        if (text.length <= maxToolOutputChars || PRUNED_MARKER_RE.test(text)) return block;
        changed = true;
        prunedBlocks += 1;
        const kept = text.slice(0, maxToolOutputChars);
        return {
          ...block,
          content: `${kept}\n[pruned ${text.length - kept.length} chars of old ${toolName || 'tool'} output to save context — re-run the tool if the full output is needed]`,
        };
      });
      return changed ? { ...msg, content } : msg;
    });

    if (prunedBlocks === 0) {
      return { pruned: false, messages, prunedBlocks: 0, tokensBefore, tokensAfter: tokensBefore };
    }
    return {
      pruned: true,
      messages: out,
      prunedBlocks,
      tokensBefore,
      tokensAfter: estimateMessagesTokens(out),
    };
  }

  selectTail(messages = [], {
    tailTokenBudget = this.tailTokenBudget,
    keepRecentUserTurns = this.keepRecentUserTurns,
  } = {}) {
    if (!Array.isArray(messages) || messages.length === 0) {
      return { head: [], tail: [], tailStartIndex: -1, tailStartId: null, tailTokens: 0, headTokens: 0 };
    }

    let tailStartIndex = messages.length;
    let tailTokens = 0;
    let userTurns = 0;

    for (let i = messages.length - 1; i >= 0; i--) {
      const msg = messages[i];
      const msgTokens = estimateMessagesTokens([msg]);
      const stillProtectingTurns = userTurns < keepRecentUserTurns;
      if (!stillProtectingTurns && tailStartIndex < messages.length && tailTokens + msgTokens > tailTokenBudget) {
        break;
      }

      tailStartIndex = i;
      tailTokens += msgTokens;
      if (msg?.role === 'user') userTurns++;

      if (!stillProtectingTurns && tailTokens >= tailTokenBudget) break;
    }

    if (tailStartIndex <= 0) {
      return {
        head: [],
        tail: cloneMessages(messages),
        tailStartIndex: 0,
        tailStartId: getMessageId(messages[0], 0),
        tailTokens,
        headTokens: 0,
      };
    }

    const head = messages.slice(0, tailStartIndex);
    const tail = messages.slice(tailStartIndex);
    return {
      head: cloneMessages(head),
      tail: cloneMessages(tail),
      tailStartIndex,
      tailStartId: getMessageId(messages[tailStartIndex], tailStartIndex),
      tailTokens,
      headTokens: estimateMessagesTokens(head),
    };
  }

  async compact(messages = [], {
    sessionId = '',
    cwd = '',
    reason = 'context_overflow',
    transcript = null,
    trace = null,
    sessionMemory = null,
    tailTokenBudget,
    keepRecentUserTurns,
    tailMode = 'preserve',
    continuePrompt = '',
  } = {}) {
    if (sessionMemory && typeof sessionMemory.precompact === 'function') {
      await sessionMemory.precompact({ messages, sessionId, cwd, reason });
    }

    const selection = this.selectTail(messages, { tailTokenBudget, keepRecentUserTurns });
    const tokensBefore = estimateMessagesTokens(messages);
    const summarizeAndContinue = tailMode === 'continue';
    const summaryMessages = summarizeAndContinue
      ? cloneMessages(messages)
      : selection.head;
    if (summaryMessages.length === 0) {
      const noCompaction = {
        compacted: false,
        reason: 'no_compactable_head',
        messages: cloneMessages(messages),
        tokensBeforeCompaction: tokensBefore,
        tokensAfterCompaction: tokensBefore,
      };
      recordContextCompaction(trace, noCompaction);
      return noCompaction;
    }

    const compactionId = newId('compaction');
    const summarySelection = { ...selection, head: summaryMessages };
    const summary = await this._summarize(summaryMessages, summarySelection);
    const timestamp = this.now();
    const userMessage = {
      role: 'user',
      content: `[Compaction:${compactionId}] Earlier Wall-E context was compacted. Recent tail starts at ${selection.tailStartId}.`,
    };
    const assistantMessage = {
      role: 'assistant',
      content: `Compaction summary:\n${summary}`,
    };
    const continuationMessage = {
      role: 'user',
      content: continuePrompt || 'Continue the Wall-E coding task from the compaction summary above. Re-read files or rerun commands if any exact detail is missing.',
    };
    const compactedMessages = summarizeAndContinue
      ? [userMessage, assistantMessage, continuationMessage]
      : [userMessage, assistantMessage, ...selection.tail];
    const tokensAfter = estimateMessagesTokens(compactedMessages);
    const metadata = {
      compactionId,
      reason,
      sessionId,
      cwd,
      timestamp,
      compaction_user_id: newId('compaction_user'),
      compaction_assistant_id: newId('compaction_assistant'),
      tail_start_id: selection.tailStartId,
      tail_start_index: selection.tailStartIndex,
      tail_message_ids: selection.tail.map((msg, index) => getMessageId(msg, selection.tailStartIndex + index)),
      tail_mode: summarizeAndContinue ? 'continue' : 'preserve',
      compacted_message_count: summaryMessages.length,
      retained_message_count: summarizeAndContinue ? 1 : selection.tail.length,
      tokens_before: tokensBefore,
      tokens_after: tokensAfter,
      summary,
    };

    if (transcript?.appendPart) {
      await transcript.appendPart({ sessionId, cwd, partType: 'compaction', data: metadata });
    } else if (transcript?.append) {
      await transcript.append({ type: 'walle_part', sessionId, cwd, partType: 'compaction', data: metadata });
    }

    const result = {
      compacted: true,
      messages: compactedMessages,
      summary,
      userMessage,
      assistantMessage,
      tailStartId: selection.tailStartId,
      metadata,
      tokensBeforeCompaction: tokensBefore,
      tokensAfterCompaction: tokensAfter,
    };
    recordContextCompaction(trace, result);
    return result;
  }

  remapTailStartId(compaction, idMap = {}, forkMessages = []) {
    return remapTailStartId(compaction, idMap, forkMessages);
  }

  async _summarize(messages, selection) {
    const historyText = messages.map(formatMessageForSummary).join('\n\n');
    const fallback = fallbackSummary(historyText, selection);
    if (!this.provider || typeof this.provider.chat !== 'function') return fallback;

    try {
      const { chunks, omittedChars } = selectSummaryChunks(historyText, SUMMARY_CHUNK_CHARS, SUMMARY_MAX_CHUNKS);
      let carry = '';
      for (const chunk of chunks) {
        carry = await this._summarizeChunk(chunk, carry);
      }
      if (!carry.trim()) return fallback;
      return omittedChars > 0
        ? `${carry}\n\n(Note: ${omittedChars} chars of mid-history were omitted from this summary.)`
        : carry;
    } catch {
      return fallback;
    }
  }

  async _summarizeChunk(chunkText, carrySummary) {
    const prompt = carrySummary
      ? `You are updating a running summary of a Wall-E coding session. Merge the new history below into the summary. Preserve decisions, file paths, commands, tool findings, failures, and unfinished work. Do not include a preamble.\n\nSUMMARY SO FAR:\n${carrySummary}\n\nNEXT PART OF HISTORY:\n${chunkText}`
      : `Summarize the compacted Wall-E coding history. Preserve decisions, file paths, commands, tool findings, failures, and unfinished work. Do not include a preamble.\n\nCOMPACTED HISTORY:\n${chunkText}`;
    const response = await this.provider.chat({
      model: this.model,
      messages: [{ role: 'user', content: prompt }],
      maxTokens: 1200,
      metadata: { purpose: 'compaction' },
    });
    return textFromContent(response?.content || '').trim();
  }
}

function toolResultText(content) {
  if (typeof content === 'string') return content;
  if (Array.isArray(content)) {
    return content
      .map((block) => (typeof block === 'string' ? block : String(block?.text || '')))
      .join('\n');
  }
  return content == null ? '' : String(content);
}

/**
 * Split history into chunks for chained summarization. When over budget,
 * keep the oldest chunks (original task framing) plus the newest chunk and
 * report how much of the middle was dropped.
 */
function selectSummaryChunks(text, chunkSize = SUMMARY_CHUNK_CHARS, maxChunks = SUMMARY_MAX_CHUNKS) {
  const value = String(text || '');
  if (value.length <= chunkSize) return { chunks: [value], omittedChars: 0 };
  const all = [];
  for (let i = 0; i < value.length; i += chunkSize) all.push(value.slice(i, i + chunkSize));
  if (all.length <= maxChunks) return { chunks: all, omittedChars: 0 };
  const kept = [...all.slice(0, maxChunks - 1), all[all.length - 1]];
  const omittedChars = value.length - kept.reduce((sum, chunk) => sum + chunk.length, 0);
  return { chunks: kept, omittedChars };
}

function remapTailStartId(compaction, idMap = {}, forkMessages = []) {
  const oldTailId = compaction?.tailStartId || compaction?.tail_start_id || compaction?.metadata?.tail_start_id || null;
  if (!oldTailId) return null;
  if (idMap[oldTailId]) return idMap[oldTailId];

  const forkIds = new Set((forkMessages || []).map((msg, index) => getMessageId(msg, index)).filter(Boolean));
  if (forkIds.has(oldTailId)) return oldTailId;

  const tailIds = compaction?.tailMessageIds || compaction?.tail_message_ids || compaction?.metadata?.tail_message_ids || [];
  for (const id of tailIds) {
    if (idMap[id]) return idMap[id];
    if (forkIds.has(id)) return id;
  }
  return forkMessages.length > 0 ? getMessageId(forkMessages[0], 0) : null;
}

function getMessageId(message, index = 0) {
  return message?.id || message?.uuid || message?.messageId || message?.message_id || `msg_${index}`;
}

function cloneMessages(messages) {
  return (messages || []).map((msg) => cloneValue(msg));
}

function cloneValue(value) {
  if (value == null) return value;
  return JSON.parse(JSON.stringify(value));
}

function formatMessageForSummary(message, index) {
  const role = message?.role || 'unknown';
  const id = getMessageId(message, index);
  return `${role} (${id}): ${textFromContent(message?.content || '').slice(0, 12000)}`;
}

function fallbackSummary(historyText, selection) {
  const line = `Compacted ${selection.head.length} older messages before tail ${selection.tailStartId}.`;
  if (!historyText) return line;
  const max = 4000;
  return `${line}\n\n${historyText.length > max ? historyText.slice(-max) : historyText}`;
}

module.exports = {
  CompactionService,
  remapTailStartId,
  getMessageId,
  selectSummaryChunks,
  DEFAULT_CONTEXT_WINDOW,
  DEFAULT_THRESHOLD,
  DEFAULT_TAIL_TOKEN_BUDGET,
  DEFAULT_KEEP_RECENT_USER_TURNS,
  DEFAULT_PRUNE_THRESHOLD,
  DEFAULT_PRUNE_PROTECT_TOKENS,
  DEFAULT_PRUNE_MAX_TOOL_OUTPUT_CHARS,
  DEFAULT_PRUNE_PROTECTED_TOOLS,
  estimateTokens,
};
