import { getPrefixedLogger } from "../../../util/logger.js";
import { buildModeMessage } from "../../util/mode-message.js";
import { findBreakpoints } from "./caching.js";
import { compact } from "./compactor.js";
import { DEFAULT_OPTIONS, DEFAULT_CACHING } from "./config.js";
import { OperationQueue } from "./operation-queue.js";
import { getTotalInputTokens, getTotalTokens } from "./token-tracker.js";
import { toStoredMessage, toSdkMessages, CACHE_BREAKPOINT_PROVIDER_OPTIONS, normalizeToolMessage, stripUndefinedFromToolOutputs, repairBrokenToolCalls, repairOrphanedToolResults, normalizeReasoningProviderOptions, stripUnsignedReasoning, } from "./types.js";
import { uniquifyFilePartFilenames } from "./uniquify-filenames.js";
const logger = getPrefixedLogger("[context-v2]");
/** Check if a tool message contains an unstripped screenshot media result. */
function hasScreenshotMedia(msg) {
    if (msg.role !== "tool" || !Array.isArray(msg.content)) {
        return false;
    }
    return msg.content.some((part) => typeof part === "object" &&
        part !== null &&
        "type" in part &&
        part.type === "tool-result" &&
        "toolName" in part &&
        part.toolName === "build_captureScreenshot" &&
        "output" in part &&
        part.output?.type === "content");
}
/** Strip base64 media from screenshot tool results, replacing with a text placeholder. */
export function stripScreenshotMediaFromToolMessage(msg) {
    if (msg.role !== "tool" || !Array.isArray(msg.content)) {
        return msg;
    }
    let modified = false;
    const strippedContent = msg.content.map((part) => {
        if (typeof part === "object" &&
            part !== null &&
            "type" in part &&
            part.type === "tool-result" &&
            "toolName" in part &&
            part.toolName === "build_captureScreenshot" &&
            "output" in part) {
            const output = part.output;
            if (output &&
                typeof output === "object" &&
                "type" in output &&
                output.type === "content") {
                modified = true;
                return {
                    ...part,
                    output: {
                        type: "text",
                        value: "[screenshot captured]",
                    },
                };
            }
        }
        return part;
    });
    if (!modified) {
        return msg;
    }
    logger.info("Stripped screenshot media from tool message before storage");
    return { ...msg, content: strippedContent };
}
export class ContextV2 {
    systemPrompt;
    messages = [];
    attachments;
    config;
    caching;
    summarizationModel;
    queue = new OperationQueue();
    onPersist;
    compactionCallbacks;
    tracer;
    rootDir;
    // Pending summaries waiting to be attached to messages
    pendingInputSummaries = new Map();
    pendingOutputSummaries = new Map();
    // Pending protection flags for tool calls (e.g., system skill reads)
    pendingProtection = new Set();
    skillsNeedReload = false;
    getModeState;
    // When disposed, this context may still be referenced by in-flight streams.
    // Keep it inert so stale generations cannot persist cleared history.
    isDisposed = false;
    constructor(options) {
        this.rootDir = options.rootDir;
        this.summarizationModel = options.summarizationModel;
        this.config = { ...DEFAULT_OPTIONS, ...options.config };
        this.caching = { ...DEFAULT_CACHING, ...options.caching };
        this.onPersist = options.onPersist;
        this.compactionCallbacks = options.compactionCallbacks;
        this.tracer = options.tracer;
        this.getModeState = options.getModeState;
    }
    setCompactionCallbacks(callbacks) {
        this.compactionCallbacks = callbacks;
    }
    updateConfig(config) {
        this.config = { ...this.config, ...config };
    }
    dispose() {
        this.isDisposed = true;
        this.onPersist = undefined;
    }
    storeSummary(toolCallId, summary) {
        if (summary.input) {
            this.pendingInputSummaries.set(toolCallId, summary.input);
        }
        this.pendingOutputSummaries.set(toolCallId, summary.output);
        // If the summary indicates protection, store that too
        if (summary.protected) {
            this.pendingProtection.add(toolCallId);
        }
    }
    markSkillsNeedReload() {
        this.skillsNeedReload = true;
    }
    checkAndClearSkillsNeedReload() {
        if (this.skillsNeedReload) {
            this.skillsNeedReload = false;
            return true;
        }
        return false;
    }
    setAttachments(attachments) {
        this.attachments = attachments;
        void this.persist();
    }
    getAttachments() {
        return this.attachments;
    }
    truncateAfterTimestamp(timestamp) {
        void this.queue.enqueue(async () => {
            const initialCount = this.messages.length;
            this.messages = this.messages.filter((msg) => !msg._superblocks?.timestamp ||
                msg._superblocks.timestamp <= timestamp);
            const removed = initialCount - this.messages.length;
            logger.info(`Truncated context after timestamp ${timestamp}. Removed ${removed} messages, ${this.messages.length} remaining.`);
            await this.persist();
        });
    }
    static fromRecord(record, options) {
        const context = new ContextV2(options);
        context.messages = record.messages;
        context.attachments = record.attachments;
        return context;
    }
    toRecord() {
        return {
            version: 2,
            messages: this.messages,
            attachments: this.attachments,
        };
    }
    setSystemPrompt(prompt) {
        this.systemPrompt = toStoredMessage(prompt);
    }
    collectPendingInputSummaries(message) {
        if (message.role !== "assistant" || typeof message.content === "string") {
            return undefined;
        }
        const toolCallIds = message.content
            .filter((p) => p.type === "tool-call" && p.toolCallId)
            .map((p) => p.toolCallId);
        const summaries = {};
        for (const id of toolCallIds) {
            const summary = this.pendingInputSummaries.get(id);
            if (summary) {
                summaries[id] = summary;
                this.pendingInputSummaries.delete(id);
            }
        }
        return Object.keys(summaries).length > 0 ? summaries : undefined;
    }
    collectPendingOutputSummaries(message) {
        if (message.role !== "tool" || !Array.isArray(message.content)) {
            return undefined;
        }
        const toolCallIds = message.content
            .filter((p) => p.type === "tool-result" && p.toolCallId)
            .map((p) => p.toolCallId);
        const summaries = {};
        for (const id of toolCallIds) {
            const summary = this.pendingOutputSummaries.get(id);
            if (summary) {
                summaries[id] = summary;
                this.pendingOutputSummaries.delete(id);
            }
        }
        return Object.keys(summaries).length > 0 ? summaries : undefined;
    }
    /**
     * Collect protection flags for tool results in this message.
     * Returns a record mapping toolCallId to protection status, or undefined if none are protected.
     * Clears the pending protection flags as they are consumed.
     */
    collectPendingProtection(message) {
        if (message.role !== "tool" || !Array.isArray(message.content)) {
            return undefined;
        }
        const toolCallIds = message.content
            .filter((p) => p.type === "tool-result" && p.toolCallId)
            .map((p) => p.toolCallId);
        const protectedResults = {};
        for (const id of toolCallIds) {
            if (this.pendingProtection.has(id)) {
                protectedResults[id] = true;
                this.pendingProtection.delete(id);
            }
        }
        return Object.keys(protectedResults).length > 0
            ? protectedResults
            : undefined;
    }
    async addUserMessage(message) {
        return this.queue.enqueue(async () => {
            this.messages.push(toStoredMessage(message));
        });
    }
    async addResponse(response, usage, options) {
        return this.queue.enqueue(async () => {
            switch (response.role) {
                case "assistant": {
                    // Total input tokens = non-cached + cached (prompt caching splits these)
                    const reportedInputTokens = getTotalInputTokens(usage);
                    // Ensure cumulative tokens never decrease - on cache misses, the LLM may
                    // report artificially low input tokens (just new tokens, not the full context).
                    // Use the higher of reported tokens or last known to maintain monotonicity.
                    const lastKnownTokens = this.getLastKnownInputTokens();
                    const cumulativeInputTokens = Math.max(reportedInputTokens, lastKnownTokens);
                    // Collect pending input summaries for tool-calls in this message
                    const toolSummaries = this.collectPendingInputSummaries(response);
                    normalizeReasoningProviderOptions(response);
                    this.messages.push(toStoredMessage(response, {
                        cumulativeInputTokens,
                        tokenCount: usage.outputTokens,
                        ...(toolSummaries && { toolSummaries }),
                    }));
                    break;
                }
                case "tool": {
                    const normalized = normalizeToolMessage(stripUndefinedFromToolOutputs(response));
                    // Strip media from older screenshots so at most one is kept in stored messages.
                    if (hasScreenshotMedia(normalized)) {
                        let strippedCount = 0;
                        for (let i = 0; i < this.messages.length; i++) {
                            const msg = this.messages[i];
                            if (msg && hasScreenshotMedia(msg)) {
                                this.messages[i] = stripScreenshotMediaFromToolMessage(msg);
                                strippedCount++;
                            }
                        }
                        if (strippedCount > 0) {
                            logger.info(`Stripped media from ${strippedCount} older screenshot(s)`);
                        }
                    }
                    // Collect pending output summaries for tool-results in this message
                    const toolSummaries = this.collectPendingOutputSummaries(normalized);
                    // Collect protection flags for individual tool results
                    const protectedResults = this.collectPendingProtection(normalized);
                    this.messages.push(toStoredMessage(normalized, {
                        ...(toolSummaries && { toolSummaries }),
                        ...(protectedResults && { protected: protectedResults }),
                    }));
                    break;
                }
                default:
                    this.messages.push(toStoredMessage(response));
                    break;
            }
            if (!options?.skipCompaction) {
                await this.compactIfNeeded(usage);
            }
            await this.persist();
        });
    }
    async getMessages() {
        return this.queue.enqueue(async () => {
            const allMessages = this.systemPrompt
                ? [this.systemPrompt, ...this.messages]
                : [...this.messages];
            // Convert to SDK format (removes _superblocks metadata)
            const sanitized = toSdkMessages(allMessages);
            // Normalize malformed tool outputs and repair broken tool calls.
            const normalized = sanitized.map((msg) => normalizeToolMessage(stripUndefinedFromToolOutputs(msg)));
            const repaired = normalized
                .map(repairBrokenToolCalls)
                .map(repairOrphanedToolResults)
                .filter((msg) => {
                // Remove messages left empty after repair (e.g., tool messages with all
                // results orphaned, or assistant messages with all tool-calls orphaned)
                if (msg.role !== "user" &&
                    Array.isArray(msg.content) &&
                    msg.content.length === 0) {
                    logger.debug(`Filtered out empty ${msg.role} message after repair`);
                    return false;
                }
                return true;
            });
            // Normalize reasoning signatures into both provider namespaces so that
            // conversations started on one provider (e.g. Vertex/anthropic) can be
            // replayed on another (e.g. Bedrock) without losing signatures.
            for (const msg of repaired) {
                if (msg.role === "assistant") {
                    normalizeReasoningProviderOptions(msg);
                }
            }
            // Drop reasoning parts that still lack a valid signature after
            // normalization. These can never be replayed successfully.
            const signatureSafe = stripUnsignedReasoning(repaired);
            if (signatureSafe.length < repaired.length) {
                logger.debug(`Stripped ${repaired.length - signatureSafe.length} message(s) with unsigned reasoning`);
            }
            // Ensure file part filenames are globally unique across all messages.
            // Providers like Bedrock reject requests with duplicate document names.
            const deduped = uniquifyFilePartFilenames(signatureSafe);
            // Apply cache breakpoints (recompute indices after filtering since
            // the .filter() step may have removed messages, shifting positions)
            const repairedBreakpoints = findBreakpoints({
                messages: deduped,
                hasSystemPrompt: this.systemPrompt !== undefined,
                config: this.caching,
            });
            for (const idx of repairedBreakpoints) {
                if (deduped[idx]) {
                    deduped[idx].providerOptions = CACHE_BREAKPOINT_PROVIDER_OPTIONS;
                }
            }
            return deduped;
        });
    }
    /**
     * Returns the last known input token count from the most recent assistant message.
     * This represents the cumulative context size at the time of that turn.
     */
    getLastKnownInputTokens() {
        for (let i = this.messages.length - 1; i >= 0; i--) {
            const meta = this.messages[i]?._superblocks;
            if (meta?.cumulativeInputTokens !== undefined) {
                return meta.cumulativeInputTokens;
            }
        }
        return 0;
    }
    async persist() {
        if (this.isDisposed) {
            return;
        }
        if (this.onPersist) {
            try {
                await this.onPersist(this.toRecord());
            }
            catch (err) {
                logger.error(`Failed to persist context: ${err}`);
            }
        }
    }
    async compactIfNeeded(usage) {
        const llmReportedTokens = getTotalTokens(usage);
        const threshold = this.config.contextLimit -
            this.config.maxOutputTokens -
            this.config.contextBufferTokens;
        if (llmReportedTokens <= threshold) {
            return;
        }
        logger.info(`Compaction triggered: ${llmReportedTokens} tokens > ${threshold} threshold ` +
            `(${this.messages.length} messages)`);
        const result = await compact({
            messages: this.messages,
            usage,
            model: this.summarizationModel,
            config: this.config,
            callbacks: this.compactionCallbacks,
            tracer: this.tracer,
            rootDir: this.rootDir,
        });
        let phase = "none";
        if (result.phase2Applied) {
            phase = "phase1+phase2";
        }
        else if (result.phase1Applied) {
            phase = "phase1";
        }
        logger.info(`Compaction complete (${phase}): saved ${result.tokensSaved} tokens, ${this.messages.length}→${result.messages.length} messages`);
        this.messages = result.messages.map((msg) => ({
            ...msg,
            _superblocks: {
                ...msg._superblocks,
                timestamp: msg._superblocks?.timestamp ?? Date.now(),
            },
        }));
        if (result.phase2Applied) {
            // Re-inject mode context that was lost when messages were replaced by the summary.
            try {
                const modeState = this.getModeState?.();
                if (modeState) {
                    this.messages.push(toStoredMessage({
                        role: "user",
                        content: buildModeMessage(modeState.mode, modeState.planApproved),
                    }));
                    logger.info(`Re-injected mode context after phase 2 compaction (mode=${modeState.mode})`);
                }
                else if (this.getModeState) {
                    logger.warn("Phase 2 compaction completed but getModeState() returned undefined; " +
                        "mode context will NOT be re-injected");
                }
            }
            catch (err) {
                logger.error(`Failed to re-inject mode state after phase 2 compaction: ${err}`);
            }
        }
    }
}
//# sourceMappingURL=context.js.map