import { OBSERVATION_CUSTOM_TYPE, OBSERVATIONAL_MEMORY_CONTEXT_CUSTOM_TYPE, OBSERVATIONAL_MEMORY_SNAPSHOT_CUSTOM_TYPE, isObservationEntryData, isReflectionRecord, isRestoredMemorySnapshot, isSupportedMemoryDetails, } from "./types.js";
import { sourceEntryReferenceIds } from "./serialize.js";
import { estimateEntryTokens } from "./tokens.js";
export function isSourceEntry(entry) {
    if (entry.type === "message" || entry.type === "branch_summary")
        return true;
    if (entry.type !== "custom_message")
        return false;
    return entry.customType !== OBSERVATIONAL_MEMORY_CONTEXT_CUSTOM_TYPE;
}
function isObservationEntry(entry) {
    return entry.type === "custom" && entry.customType === OBSERVATION_CUSTOM_TYPE;
}
function isRestoredSnapshotEntry(entry) {
    return entry.type === "custom" && entry.customType === OBSERVATIONAL_MEMORY_SNAPSHOT_CUSTOM_TYPE;
}
export function findLastCompactionIndex(entries) {
    for (let i = entries.length - 1; i >= 0; i--) {
        if (entries[i].type === "compaction")
            return i;
    }
    return -1;
}
export function lastObservationCoverEndIdx(entries) {
    const idToIdx = new Map();
    for (let i = 0; i < entries.length; i++)
        idToIdx.set(entries[i].id, i);
    let maxIdx = -1;
    for (let i = 0; i < entries.length; i++) {
        const entry = entries[i];
        if (!isObservationEntry(entry))
            continue;
        if (!isObservationEntryData(entry.data))
            continue;
        const coverIdx = idToIdx.get(entry.data.coversUpToId);
        if (coverIdx !== undefined && coverIdx > maxIdx)
            maxIdx = coverIdx;
    }
    return maxIdx;
}
function getLatestRestoredSnapshot(entries) {
    for (let i = entries.length - 1; i >= 0; i--) {
        const entry = entries[i];
        if (!isRestoredSnapshotEntry(entry))
            continue;
        if (!isRestoredMemorySnapshot(entry.data))
            continue;
        return entry.data;
    }
    return undefined;
}
function rawTokensFromIndex(entries, startIndex) {
    let total = 0;
    for (let i = Math.max(0, startIndex); i < entries.length; i++) {
        if (isSourceEntry(entries[i]))
            total += estimateEntryTokens(entries[i]);
    }
    return total;
}
function branchIndexAfterCoveredSourceCount(entries, coveredSourceCount) {
    if (coveredSourceCount <= 0)
        return 0;
    let seen = 0;
    for (let i = 0; i < entries.length; i++) {
        if (!isSourceEntry(entries[i]))
            continue;
        seen += 1;
        if (seen >= coveredSourceCount)
            return i + 1;
    }
    return entries.length;
}
export function rawTokensSinceLastBound(entries) {
    const coveredByLiveObservation = lastObservationCoverEndIdx(entries) + 1;
    const restoredSnapshot = getLatestRestoredSnapshot(entries);
    const coveredBySnapshot = branchIndexAfterCoveredSourceCount(entries, restoredSnapshot?.coveredSourceCount ?? 0);
    return rawTokensFromIndex(entries, Math.max(coveredByLiveObservation, coveredBySnapshot));
}
export function rawTokensSinceLastCompaction(entries) {
    const compactionIdx = findLastCompactionIndex(entries);
    if (compactionIdx === -1)
        return rawTokensFromIndex(entries, 0);
    return rawTokensFromIndex(entries, liveTailStartIndex(entries));
}
function liveTailStartIndex(entries) {
    const compactionIdx = findLastCompactionIndex(entries);
    if (compactionIdx === -1)
        return 0;
    const firstKept = entries[compactionIdx].firstKeptEntryId;
    if (!firstKept)
        throw new Error("compaction entry missing firstKeptEntryId");
    const firstKeptIdx = entries.findIndex((e) => e.id === firstKept);
    if (firstKeptIdx === -1)
        throw new Error(`firstKeptEntryId "${firstKept}" not found in entries`);
    return firstKeptIdx;
}
export function firstRawIdAfter(entries, afterIndex) {
    for (let i = Math.max(0, afterIndex + 1); i < entries.length; i++) {
        if (isSourceEntry(entries[i]))
            return entries[i].id;
    }
    return undefined;
}
export function gapRawEntries(entries, newFirstKeptEntryId) {
    const lastBoundIdx = lastObservationCoverEndIdx(entries);
    const newKeptIdx = entries.findIndex((e) => e.id === newFirstKeptEntryId);
    if (newKeptIdx === -1)
        return [];
    const result = [];
    for (let i = lastBoundIdx + 1; i < newKeptIdx; i++) {
        if (isSourceEntry(entries[i]))
            result.push(entries[i]);
    }
    return result;
}
export function rawTailEntriesBetween(entries, fromId, untilId) {
    const fromIdx = entries.findIndex((e) => e.id === fromId);
    const untilIdx = entries.findIndex((e) => e.id === untilId);
    if (fromIdx === -1 || untilIdx === -1 || untilIdx < fromIdx)
        return [];
    const result = [];
    for (let i = fromIdx; i <= untilIdx; i++) {
        if (isSourceEntry(entries[i]))
            result.push(entries[i]);
    }
    return result;
}
function uniqueIds(ids) {
    return Array.from(new Set(ids));
}
function collectIndexedObservations(entries) {
    const observations = [];
    for (let branchIndex = 0; branchIndex < entries.length; branchIndex++) {
        const entry = entries[branchIndex];
        if (!isObservationEntry(entry))
            continue;
        if (!isObservationEntryData(entry.data))
            continue;
        entry.data.records.forEach((observation, observationRecordIndex) => {
            observations.push({ observation, observationEntryId: entry.id, observationRecordIndex, branchIndex });
        });
    }
    return observations;
}
function observationKey(observation) {
    return `${observation.observationEntryId}:${observation.observationRecordIndex}`;
}
function sourceEntryMatchesRequestedId(entry, requested) {
    if (!isSourceEntry(entry))
        return false;
    if (entry.id === requested)
        return true;
    return sourceEntryReferenceIds(entry).includes(requested);
}
function buildSourceEntryResolution(entries, sourceEntryIds) {
    const requested = uniqueIds(sourceEntryIds);
    const matchedEntries = new Map();
    for (const entry of entries) {
        for (const requestedId of requested) {
            if (!sourceEntryMatchesRequestedId(entry, requestedId))
                continue;
            matchedEntries.set(entry.id, entry);
        }
    }
    const missingSourceEntryIds = [];
    const nonSourceEntryIds = [];
    for (const requestedId of requested) {
        const directEntry = entries.find((entry) => entry.id === requestedId);
        if (directEntry && !isSourceEntry(directEntry)) {
            nonSourceEntryIds.push(requestedId);
            continue;
        }
        const matched = entries.some((entry) => sourceEntryMatchesRequestedId(entry, requestedId));
        if (!matched)
            missingSourceEntryIds.push(requestedId);
    }
    const sourceEntries = entries.filter((entry) => matchedEntries.has(entry.id) && isSourceEntry(entry));
    return { requested, sourceEntries, missingSourceEntryIds, nonSourceEntryIds };
}
function resolveSourceEntries(entries, sourceEntryIds) {
    const resolved = buildSourceEntryResolution(entries, sourceEntryIds);
    if (resolved.missingSourceEntryIds.length > 0 || resolved.nonSourceEntryIds.length > 0) {
        return {
            status: "source_unavailable",
            sourceEntryIds: resolved.requested,
            sourceEntries: [],
            missingSourceEntryIds: resolved.missingSourceEntryIds,
            nonSourceEntryIds: resolved.nonSourceEntryIds,
        };
    }
    return {
        status: "ok",
        sourceEntryIds: resolved.requested,
        sourceEntries: resolved.sourceEntries,
        missingSourceEntryIds: [],
        nonSourceEntryIds: [],
    };
}
function resolveSourceEntriesPartial(entries, sourceEntryIds) {
    const resolved = buildSourceEntryResolution(entries, sourceEntryIds);
    return {
        status: resolved.missingSourceEntryIds.length > 0 || resolved.nonSourceEntryIds.length > 0 ? "source_unavailable" : "ok",
        sourceEntryIds: resolved.requested,
        sourceEntries: resolved.sourceEntries,
        missingSourceEntryIds: resolved.missingSourceEntryIds,
        nonSourceEntryIds: resolved.nonSourceEntryIds,
    };
}
function memoryObservationFromIndexed(entries, indexed) {
    const { observation, observationEntryId, observationRecordIndex } = indexed;
    if (!observation.sourceEntryIds || observation.sourceEntryIds.length === 0) {
        return { status: "no_source", observation, observationEntryId, observationRecordIndex };
    }
    const resolved = resolveSourceEntriesPartial(entries, observation.sourceEntryIds);
    if (resolved.status === "source_unavailable") {
        return {
            status: "source_unavailable",
            observation,
            observationEntryId,
            observationRecordIndex,
            sourceEntryIds: resolved.sourceEntryIds,
            sourceEntries: resolved.sourceEntries,
            missingSourceEntryIds: resolved.missingSourceEntryIds,
            nonSourceEntryIds: resolved.nonSourceEntryIds,
        };
    }
    return {
        status: "ok",
        observation,
        observationEntryId,
        observationRecordIndex,
        sourceEntryIds: resolved.sourceEntryIds,
        sourceEntries: resolved.sourceEntries,
    };
}
function uniqueSourceEntriesInBranchOrder(entries, observations) {
    const requested = new Set();
    for (const observation of observations) {
        if (observation.status === "ok" || observation.status === "source_unavailable") {
            for (const entry of observation.sourceEntries)
                requested.add(entry.id);
        }
    }
    return entries.filter((entry) => requested.has(entry.id) && isSourceEntry(entry));
}
function uniqueUnavailableSourceIds(observations, field) {
    const ids = [];
    for (const observation of observations) {
        if (observation.status !== "source_unavailable")
            continue;
        ids.push(...observation[field]);
    }
    return uniqueIds(ids);
}
export function recallObservationSources(entries, observationId) {
    const matches = [];
    for (const entry of entries) {
        if (!isObservationEntry(entry))
            continue;
        if (!isObservationEntryData(entry.data))
            continue;
        for (const observation of entry.data.records) {
            if (observation.id !== observationId)
                continue;
            if (!observation.sourceEntryIds || observation.sourceEntryIds.length === 0) {
                matches.push({ status: "no_source", observation, observationEntryId: entry.id });
                continue;
            }
            const resolved = resolveSourceEntries(entries, observation.sourceEntryIds);
            if (resolved.status === "source_unavailable") {
                matches.push({
                    status: "source_unavailable",
                    observation,
                    observationEntryId: entry.id,
                    sourceEntryIds: resolved.sourceEntryIds,
                    missingSourceEntryIds: resolved.missingSourceEntryIds,
                    nonSourceEntryIds: resolved.nonSourceEntryIds,
                });
                continue;
            }
            matches.push({
                status: "ok",
                observation,
                observationEntryId: entry.id,
                sourceEntryIds: resolved.sourceEntryIds,
                sourceEntries: resolved.sourceEntries,
            });
        }
    }
    if (matches.length === 0)
        return { status: "not_found", observationId, matches: [], collision: false };
    return { status: "found", observationId, matches, collision: matches.length > 1 };
}
function getPriorMemoryDetails(entries) {
    const idx = findLastCompactionIndex(entries);
    if (idx !== -1) {
        const details = entries[idx].details;
        if (isSupportedMemoryDetails(details))
            return details;
    }
    return getLatestRestoredSnapshot(entries)?.details;
}
export function recallMemorySources(entries, memoryId) {
    const indexedObservations = collectIndexedObservations(entries);
    const observationsById = new Map();
    for (const observation of indexedObservations) {
        const matches = observationsById.get(observation.observation.id) ?? [];
        matches.push(observation);
        observationsById.set(observation.observation.id, matches);
    }
    const priorDetails = getPriorMemoryDetails(entries);
    const reflectionMatches = [];
    if (priorDetails) {
        priorDetails.reflections.forEach((reflection, reflectionIndex) => {
            if (!isReflectionRecord(reflection) || reflection.id !== memoryId)
                return;
            reflectionMatches.push({ reflection, reflectionIndex });
        });
    }
    const directIndexedObservations = observationsById.get(memoryId) ?? [];
    const observationsByKey = new Map();
    for (const observation of directIndexedObservations) {
        observationsByKey.set(observationKey(observation), observation);
    }
    const unavailableSupportingObservations = [];
    const unavailableReflectionProvenance = [];
    for (const reflectionMatch of reflectionMatches) {
        if (reflectionMatch.reflection.legacy === true) {
            unavailableReflectionProvenance.push({ ...reflectionMatch, reason: "legacy" });
            continue;
        }
        for (const observationId of reflectionMatch.reflection.supportingObservationIds) {
            const supportingObservations = observationsById.get(observationId);
            if (!supportingObservations || supportingObservations.length === 0) {
                unavailableSupportingObservations.push({ ...reflectionMatch, observationId });
                continue;
            }
            for (const observation of supportingObservations) {
                observationsByKey.set(observationKey(observation), observation);
            }
        }
    }
    const indexedObservationBag = Array.from(observationsByKey.values()).sort((a, b) => {
        if (a.branchIndex !== b.branchIndex)
            return a.branchIndex - b.branchIndex;
        return a.observationRecordIndex - b.observationRecordIndex;
    });
    const observations = indexedObservationBag.map((observation) => memoryObservationFromIndexed(entries, observation));
    const directObservationKeys = new Set(directIndexedObservations.map(observationKey));
    const directObservationMatches = observations.filter((observation) => directObservationKeys.has(observationKey(observation)));
    const sourceEntries = uniqueSourceEntriesInBranchOrder(entries, observations);
    const missingSourceEntryIds = uniqueUnavailableSourceIds(observations, "missingSourceEntryIds");
    const nonSourceEntryIds = uniqueUnavailableSourceIds(observations, "nonSourceEntryIds");
    if (reflectionMatches.length === 0 && directObservationMatches.length === 0) {
        return {
            status: "not_found",
            memoryId,
            reflectionMatches: [],
            directObservationMatches: [],
            observations: [],
            sourceEntries: [],
            unavailableSupportingObservations: [],
            unavailableReflectionProvenance: [],
            missingSourceEntryIds: [],
            nonSourceEntryIds: [],
            collision: false,
            partial: false,
        };
    }
    const hasNoSourceObservationInMemoryRecall = reflectionMatches.length > 0 && observations.some((observation) => observation.status === "no_source");
    return {
        status: "found",
        memoryId,
        reflectionMatches,
        directObservationMatches,
        observations,
        sourceEntries,
        unavailableSupportingObservations,
        unavailableReflectionProvenance,
        missingSourceEntryIds,
        nonSourceEntryIds,
        collision: reflectionMatches.length + directObservationMatches.length > 1,
        partial: unavailableReflectionProvenance.length > 0 ||
            hasNoSourceObservationInMemoryRecall ||
            unavailableSupportingObservations.length > 0 ||
            missingSourceEntryIds.length > 0 ||
            nonSourceEntryIds.length > 0,
    };
}
export function collectObservationsByCoverage(entries, priorFirstKeptEntryId, newFirstKeptEntryId) {
    const idToIdx = new Map();
    for (let i = 0; i < entries.length; i++)
        idToIdx.set(entries[i].id, i);
    const newFKIIdx = idToIdx.get(newFirstKeptEntryId);
    if (newFKIIdx === undefined)
        return [];
    let priorFKIIdx;
    if (priorFirstKeptEntryId === undefined) {
        priorFKIIdx = -1;
    }
    else {
        const idx = idToIdx.get(priorFirstKeptEntryId);
        if (idx === undefined)
            throw new Error(`priorFirstKeptEntryId "${priorFirstKeptEntryId}" not found in entries`);
        priorFKIIdx = idx;
    }
    const result = [];
    for (const entry of entries) {
        if (!isObservationEntry(entry))
            continue;
        if (!isObservationEntryData(entry.data))
            continue;
        const fromIdx = idToIdx.get(entry.data.coversFromId);
        if (fromIdx === undefined)
            continue;
        if (fromIdx >= priorFKIIdx && fromIdx < newFKIIdx)
            result.push(entry.data);
    }
    return result;
}
function collectObservationsPendingNextCompaction(entries) {
    const idToIdx = new Map();
    for (let i = 0; i < entries.length; i++)
        idToIdx.set(entries[i].id, i);
    const priorCompactionIdx = findLastCompactionIndex(entries);
    const restoredSnapshot = getLatestRestoredSnapshot(entries);
    let thresholdStartIdx;
    if (priorCompactionIdx === -1) {
        thresholdStartIdx = branchIndexAfterCoveredSourceCount(entries, restoredSnapshot?.coveredSourceCount ?? 0);
    }
    else {
        const priorFirstKept = entries[priorCompactionIdx].firstKeptEntryId;
        if (!priorFirstKept)
            throw new Error("prior compaction entry missing firstKeptEntryId");
        const idx = idToIdx.get(priorFirstKept);
        if (idx === undefined)
            throw new Error(`prior firstKeptEntryId "${priorFirstKept}" not found in entries`);
        thresholdStartIdx = idx;
    }
    const result = [];
    for (const entry of entries) {
        if (!isObservationEntry(entry))
            continue;
        if (!isObservationEntryData(entry.data))
            continue;
        const fromIdx = idToIdx.get(entry.data.coversFromId);
        if (fromIdx === undefined)
            continue;
        if (fromIdx >= thresholdStartIdx)
            result.push(entry.data);
    }
    return result;
}
export function getMemoryState(entries) {
    const priorDetails = getPriorMemoryDetails(entries);
    const pendingData = collectObservationsPendingNextCompaction(entries);
    return {
        reflections: priorDetails?.reflections ?? [],
        committedObs: priorDetails?.observations ?? [],
        pendingObs: pendingData.flatMap((d) => d.records),
    };
}
