/**
 * Context Manager - tracks token usage and manages conversation compression
 *
 * Smart Compression Strategy (beyond Claude Code):
 *   - Uses Haiku for fast, cheap summarization (~0.5s, ~$0.001/compression)
 *   - Preserves "anchor turns": user decisions, file edits, error resolution
 *   - Two-phase compression: first summarize old turns, then if still over, summarize summary
 *   - Tracks compression history for debugging
 */
export class ContextManager {
  constructor(config) {
    this.config = config;
    this.maxTokens = 200000; // Claude context window
    this.maxOutputTokens = config.maxOutputTokens || 16000;
    this.messages = [];
    this.turnCount = 0;
    this.estimatedTokens = 0;
    this.systemTokens = 0;
    this.toolTokens = 0;
    this.compressionThreshold = 0.85; // Compress at 85% usage
    this.client = null;

    // Real token tracking from API response
    this.hasRealUsage = false;
    this.lastInputTokens = 0;
    this.lastOutputTokens = 0;
    this.cacheCreation = 0;
    this.cacheRead = 0;
    this.totalInputTokens = 0;
    this.totalOutputTokens = 0;

    // Compression history
    this.compressionCount = 0;
    this.totalTokensSaved = 0;

    // Context exclusion: set of message indices to skip when sending to LLM
    this.excludedIndices = new Set();
  }

  setSystemTokenEstimate(tokens) {
    this.systemTokens = tokens;
  }

  setToolTokenEstimate(tokens) {
    this.toolTokens = tokens;
  }

  addUserMessage(content) {
    this.messages.push({ role: 'user', content });
    this.estimateTokens();
    this.turnCount++;
  }

  addAssistantMessage(content) {
    this.messages.push({ role: 'assistant', content });
    this.estimateTokens();
  }

  addToolResult(toolUseId, content) {
    // Tool results go into user messages per Anthropic API format.
    // When multiple tool_use blocks are in the same assistant message,
    // their tool_results MUST be in the same user message — otherwise
    // _fixOrphanedToolUse and getMessages() merge logic can mis-pair them.
    const resultBlock = {
      type: 'tool_result',
      tool_use_id: toolUseId,
      content: typeof content === 'string' ? content : JSON.stringify(content),
    };

    // If the last message is already a user message with tool_result blocks,
    // append to it (same tool-call round) instead of creating a new message.
    const last = this.messages[this.messages.length - 1];
    if (last && last.role === 'user' && Array.isArray(last.content)
        && last.content.length > 0 && last.content[0].type === 'tool_result') {
      last.content.push(resultBlock);
    } else {
      this.messages.push({ role: 'user', content: [resultBlock] });
    }
    this.estimateTokens();
  }

  /**
   * Get messages with incremental caching (BP3/BP4).
   * Places cache_control on the last assistant message so the API caches
   * all conversation history up to that point. On the next turn, only new
   * messages after the breakpoint are fresh tokens (~85% latency reduction).
   */
  getMessages() {
    // Filter out excluded messages, but ensure valid role alternation.
    // Tool results (user messages containing tool_result blocks) tied to an
    // assistant tool_use cannot be excluded without also removing the tool_use,
    // so we only exclude user text messages and their paired assistant replies.
    let filtered = [];
    for (let i = 0; i < this.messages.length; i++) {
      if (this.excludedIndices.has(i)) continue;
      filtered.push(this.messages[i]);
    }

    // Ensure valid role alternation: consecutive same-role messages break the API.
    // Merge consecutive user messages and drop orphaned assistant messages.
    const merged = [];
    for (const msg of filtered) {
      if (merged.length > 0 && merged[merged.length - 1].role === msg.role) {
        // Same role as previous — merge into previous
        const prev = merged[merged.length - 1];
        const prevBlocks = typeof prev.content === 'string'
          ? [{ type: 'text', text: prev.content }] : (Array.isArray(prev.content) ? prev.content : []);
        const curBlocks = typeof msg.content === 'string'
          ? [{ type: 'text', text: msg.content }] : (Array.isArray(msg.content) ? msg.content : []);
        merged[merged.length - 1] = { ...prev, content: [...prevBlocks, ...curBlocks] };
      } else {
        merged.push(msg);
      }
    }

    // Must start with user message
    while (merged.length > 0 && merged[0].role !== 'user') merged.shift();

    if (merged.length < 2) return merged;

    // Find last assistant message to place cache breakpoint
    let lastAssistantIdx = -1;
    for (let i = merged.length - 1; i >= 0; i--) {
      if (merged[i].role === 'assistant') {
        lastAssistantIdx = i;
        break;
      }
    }

    // No assistant message yet (first turn), return as-is
    if (lastAssistantIdx === -1) return merged;

    // Final validation: ensure tool_use/tool_result pairing across entire messages array
    const validated = this._validateToolPairing(merged);

    // Return messages with cache_control on last assistant message's last content block
    const finalLastAssistantIdx = validated.reduce((idx, msg, i) => msg.role === 'assistant' ? i : idx, -1);
    return validated.map((msg, idx) => {
      if (idx !== finalLastAssistantIdx) return msg;
      return this._addCacheControl(msg);
    });
  }

  /**
   * Final pass: ensure every tool_use has a corresponding tool_result in the next message,
   * and every tool_result has a corresponding tool_use in the previous message.
   */
  _validateToolPairing(messages) {
    const result = [];
    for (let i = 0; i < messages.length; i++) {
      const msg = messages[i];
      result.push(msg);

      // After an assistant message with tool_use blocks, check if next message has tool_results
      if (msg.role === 'assistant' && Array.isArray(msg.content)) {
        const toolUseBlocks = msg.content.filter(b => b.type === 'tool_use');
        if (toolUseBlocks.length === 0) continue;

        const next = messages[i + 1];
        const nextResultIds = new Set(
          next && Array.isArray(next.content)
            ? next.content.filter(b => b.type === 'tool_result').map(b => b.tool_use_id)
            : []
        );
        const unpairedUses = toolUseBlocks.filter(b => !nextResultIds.has(b.id));
        if (unpairedUses.length > 0) {
          // Inject synthetic tool_result message for unmatched tool_use blocks
          result.push({ role: 'user', content: unpairedUses.map(b => ({ type: 'tool_result', tool_use_id: b.id, content: '[Result not available]' })) });
        }
      }
    }

    // Remove/strip orphaned tool_result blocks from user messages
    const allToolUseIds = new Set();
    for (const msg of result) {
      if (msg.role === 'assistant' && Array.isArray(msg.content)) {
        msg.content.filter(b => b.type === 'tool_use').forEach(b => allToolUseIds.add(b.id));
      }
    }
    return result.map(msg => {
      if (msg.role !== 'user' || !Array.isArray(msg.content)) return msg;
      // Check if any tool_result blocks are orphaned
      const hasOrphan = msg.content.some(
        b => b.type === 'tool_result' && !allToolUseIds.has(b.tool_use_id)
      );
      if (!hasOrphan) return msg;
      // Strip orphaned tool_result blocks (keep text + valid tool_results)
      const cleaned = msg.content.filter(
        b => b.type !== 'tool_result' || allToolUseIds.has(b.tool_use_id)
      );
      return cleaned.length > 0 ? { ...msg, content: cleaned } : null;
    }).filter(Boolean).filter(msg => {
      // Remove user messages that became empty or only had orphaned tool_results
      if (msg.role !== 'user' || !Array.isArray(msg.content)) return true;
      return msg.content.length > 0;
    });
  }

  /** Add cache_control: ephemeral to the last content block of a message */
  _addCacheControl(msg) {
    const content = msg.content;
    if (typeof content === 'string') {
      return {
        ...msg,
        content: [{ type: 'text', text: content, cache_control: { type: 'ephemeral' } }],
      };
    }
    if (Array.isArray(content) && content.length > 0) {
      const lastIdx = content.length - 1;
      return {
        ...msg,
        content: content.map((block, i) =>
          i === lastIdx ? { ...block, cache_control: { type: 'ephemeral' } } : block
        ),
      };
    }
    return msg;
  }

  /** Update with real token counts from API response */
  updateFromUsage(usage) {
    if (!usage) return;
    this.hasRealUsage = true;
    this.lastInputTokens = usage.input_tokens || 0;
    this.lastOutputTokens = usage.output_tokens || 0;
    this.cacheCreation = usage.cache_creation_input_tokens || 0;
    this.cacheRead = usage.cache_read_input_tokens || 0;
    this.totalInputTokens += this.lastInputTokens;
    this.totalOutputTokens += this.lastOutputTokens;
    // Use real input tokens for context window tracking
    this.estimatedTokens = this.lastInputTokens + this.maxOutputTokens;
    // Save anchor for delta-based estimation between API calls
    this._lastRealMsgCount = this.messages.length;
    this._lastRealTokens = this.estimatedTokens;
  }

  getUsage() {
    const base = {
      usedTokens: this.estimatedTokens,
      maxTokens: this.maxTokens,
      messageCount: this.messages.length,
      turnCount: this.turnCount,
    };
    if (this.hasRealUsage) {
      base.inputTokens = this.lastInputTokens;
      base.outputTokens = this.lastOutputTokens;
      base.cacheCreation = this.cacheCreation;
      base.cacheRead = this.cacheRead;
      base.totalInputTokens = this.totalInputTokens;
      base.totalOutputTokens = this.totalOutputTokens;
    }
    return base;
  }

  /** Rough token estimation (4 chars ≈ 1 token) */
  estimateTokens() {
    // If we have real usage from the API, only add a delta for new messages
    // since the last API call, instead of re-counting everything crudely.
    if (this.hasRealUsage && this._lastRealMsgCount !== undefined) {
      const newMsgs = this.messages.slice(this._lastRealMsgCount);
      let delta = 0;
      for (const msg of newMsgs) {
        if (typeof msg.content === 'string') {
          delta += Math.ceil(msg.content.length / 4);
        } else if (Array.isArray(msg.content)) {
          for (const block of msg.content) {
            const text = block.text || block.content || '';
            delta += Math.ceil(text.length / 4);
          }
        }
      }
      this.estimatedTokens = this._lastRealTokens + delta;
      return;
    }
    // Fallback: crude estimate from scratch (before first API response)
    let total = this.systemTokens + this.toolTokens;
    for (const msg of this.messages) {
      if (typeof msg.content === 'string') {
        total += Math.ceil(msg.content.length / 4);
      } else if (Array.isArray(msg.content)) {
        for (const block of msg.content) {
          if (block.type === 'text' || block.type === 'tool_result') {
            const text = block.text || block.content || '';
            total += Math.ceil(text.length / 4);
          }
        }
      }
    }
    total += this.maxOutputTokens; // Reserve for output
    this.estimatedTokens = total;
  }

  needsCompression() {
    return this.estimatedTokens / this.maxTokens > this.compressionThreshold;
  }

  /**
   * Smart compression: uses Haiku for fast summarization, preserves anchor turns
   * Anchor turns = messages containing: file edits, user decisions, error resolutions
   */
  async compress(client) {
    if (this.messages.length < 6) return false;

    const beforeTokens = this.estimatedTokens;
    const keepRecent = 6; // Keep last 6 messages intact (more than basic)
    const toCompress = this.messages.slice(0, -keepRecent);
    const recentMessages = this.messages.slice(-keepRecent);

    // Identify anchor turns (important messages to preserve verbatim)
    const anchors = [];
    const compressible = [];

    for (const msg of toCompress) {
      if (this._isAnchorTurn(msg)) {
        anchors.push(msg);
      } else {
        compressible.push(msg);
      }
    }

    // Build summary of compressible messages
    const summaryContent = compressible.map(m => {
      const role = m.role;
      const text = this._extractText(m);
      return `[${role}]: ${text.substring(0, 600)}`;
    }).join('\n---\n');

    if (!summaryContent.trim()) {
      // Only anchors, keep as-is but truncate tool results
      this.messages = [...this._truncateAnchors(anchors), ...recentMessages];
      this.estimateTokens();
      return true;
    }

    try {
      // Use Haiku for fast, cheap compression (~0.5s, ~$0.001)
      const compressionModel = 'claude-haiku-4-5-20251001';
      const response = await client.messages.create({
        model: compressionModel,
        max_tokens: 2000,
        system: `You are a conversation compressor for an AI coding assistant. Summarize the conversation preserving:
1. Current task and status (in progress / completed / blocked)
2. Key decisions made and their reasoning
3. File paths modified and what was changed
4. Unresolved errors or pending tasks
5. User preferences expressed during the conversation

Format as a structured summary with sections. Be concise but keep all actionable details.`,
        messages: [{ role: 'user', content: summaryContent }],
      });

      const summary = response.content[0]?.text || 'Summary unavailable';

      // Reconstruct: summary + preserved anchors + recent messages
      const newMessages = [
        { role: 'user', content: `<system-reminder>\n[Compressed conversation history — ${this.compressionCount + 1} compressions so far]\n${summary}\n</system-reminder>\nPlease continue from where we left off.` },
        { role: 'assistant', content: 'I have the context from the compressed history. Continuing.' },
      ];

      // Re-insert anchor turns (truncated), with tool_use/tool_result pairing validated
      const safeAnchors = this._sanitizeAnchorsForPairing(this._truncateAnchors(anchors).slice(-4));
      for (const anchor of safeAnchors) {
        newMessages.push(anchor);
      }

      this.messages = [...newMessages, ...recentMessages];
      this.estimateTokens();
      this.compressionCount++;
      this.totalTokensSaved += Math.max(0, beforeTokens - this.estimatedTokens);
      return true;
    } catch {
      // Fallback: simple truncation
      this.messages = [
        { role: 'user', content: '[Earlier conversation was truncated due to context limits]' },
        { role: 'assistant', content: 'Understood. I\'ll continue with the available context.' },
        ...recentMessages,
      ];
      this.estimateTokens();
      this.compressionCount++;
      return true;
    }
  }

  /** Check if a message is an "anchor turn" worth preserving */
  _isAnchorTurn(msg) {
    const text = this._extractText(msg);
    // User decisions (approval, rejection, explicit instructions)
    if (msg.role === 'user' && /\/(approve|cancel|yes|no)\b/i.test(text)) return true;
    // File edits (tool results from EditFile/WriteFile)
    if (text.includes('EditFile') || text.includes('WriteFile')) return true;
    // Error resolutions
    if (text.includes('Error:') && text.includes('fix')) return true;
    // Todo updates
    if (text.includes('TodoWrite') || text.includes('Todos updated')) return true;
    return false;
  }

  /** Extract text content from a message */
  _extractText(msg) {
    if (typeof msg.content === 'string') return msg.content;
    if (Array.isArray(msg.content)) {
      return msg.content.map(b => b.text || b.content || '').join('\n');
    }
    return '';
  }

  /** Truncate anchor turns to save space (keep key info, trim verbose tool results) */
  _truncateAnchors(anchors) {
    return anchors.map(msg => {
      if (typeof msg.content === 'string' && msg.content.length > 800) {
        return { ...msg, content: msg.content.substring(0, 800) + '\n[...truncated]' };
      }
      if (Array.isArray(msg.content)) {
        return {
          ...msg,
          content: msg.content.map(b => {
            if ((b.type === 'text' || b.type === 'tool_result') && (b.text || b.content || '').length > 500) {
              const field = b.text ? 'text' : 'content';
              return { ...b, [field]: (b[field] || '').substring(0, 500) + '\n[...truncated]' };
            }
            return b;
          }),
        };
      }
      return msg;
    });
  }

  /**
   * Ensure re-inserted anchor turns have valid tool_use/tool_result pairing.
   * Drop orphaned tool_result messages and inject synthetic tool_results
   * for assistant tool_use blocks whose results were compressed away.
   */
  _sanitizeAnchorsForPairing(anchors) {
    // Collect all tool_use IDs from assistant anchors
    const toolUseIds = new Set();
    // Collect all tool_result IDs from user anchors
    const toolResultIds = new Set();

    for (const msg of anchors) {
      if (!Array.isArray(msg.content)) continue;
      for (const block of msg.content) {
        if (block.type === 'tool_use' && block.id) toolUseIds.add(block.id);
        if (block.type === 'tool_result' && block.tool_use_id) toolResultIds.add(block.tool_use_id);
      }
    }

    // Find unpaired: tool_use without tool_result, and vice versa
    const unpairedUseIds = [...toolUseIds].filter(id => !toolResultIds.has(id));
    const unpairedResultIds = new Set([...toolResultIds].filter(id => !toolUseIds.has(id)));

    // Filter out user messages that only contain orphaned tool_results
    const filtered = anchors.filter(msg => {
      if (msg.role !== 'user' || !Array.isArray(msg.content)) return true;
      const hasOnlyOrphanedResults = msg.content.every(
        b => b.type === 'tool_result' && unpairedResultIds.has(b.tool_use_id)
      );
      return !hasOnlyOrphanedResults;
    });

    // For unpaired tool_use IDs, inject synthetic tool_result after the assistant message
    if (unpairedUseIds.length > 0) {
      const result = [];
      for (const msg of filtered) {
        result.push(msg);
        if (msg.role === 'assistant' && Array.isArray(msg.content)) {
          const needsResults = msg.content
            .filter(b => b.type === 'tool_use' && unpairedUseIds.includes(b.id))
            .map(b => ({ type: 'tool_result', tool_use_id: b.id, content: '[Result compressed]' }));
          if (needsResults.length > 0) {
            result.push({ role: 'user', content: needsResults });
          }
        }
      }
      return result;
    }

    return filtered;
  }

  /** Toggle exclusion for a range of message indices (a "turn") */
  toggleExclude(startIdx, endIdx) {
    if (startIdx < 0 || startIdx >= this.messages.length) return false;
    const end = Math.min(endIdx ?? startIdx, this.messages.length - 1);
    // Check if any in range are excluded
    let anyExcluded = false;
    for (let i = startIdx; i <= end; i++) {
      if (this.excludedIndices.has(i)) { anyExcluded = true; break; }
    }
    // Toggle: if any excluded, re-include all; otherwise exclude all
    for (let i = startIdx; i <= end; i++) {
      if (anyExcluded) {
        this.excludedIndices.delete(i);
      } else {
        this.excludedIndices.add(i);
      }
    }
    this.estimateTokens();
    return !anyExcluded; // returns new excluded state
  }

  /** Get which indices are excluded */
  getExcludedIndices() {
    return [...this.excludedIndices];
  }

  /** Get current message count (for UI turn tracking) */
  getMessageCount() {
    return this.messages.length;
  }

  /** Get message list with metadata (for UI display of context status) */
  getMessageList() {
    return this.messages.map((msg, i) => ({
      index: i,
      role: msg.role,
      excluded: this.excludedIndices.has(i),
      preview: this._extractText(msg).substring(0, 120),
    }));
  }

  getCompressionStats() {
    return {
      compressions: this.compressionCount,
      tokensSaved: this.totalTokensSaved,
    };
  }
}
