/**
 * @license
 * Copyright 2026 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */
import { deriveStableId } from '../../utils/cryptoUtils.js';
import type { JSONSchemaType } from 'ajv';
import type {
  ContextProcessor,
  ProcessArgs,
  BackstopTargetOptions,
} from '../pipeline.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
import { type ConcreteNode, type Snapshot, NodeType } from '../graph/types.js';
import {
  SnapshotGenerator,
  findLatestSnapshotBaseline,
} from '../utils/snapshotGenerator.js';
import { debugLogger } from '../../utils/debugLogger.js';

export interface StateSnapshotProcessorOptions extends BackstopTargetOptions {
  model?: string;
  systemInstruction?: string;
  maxSummaryTurns?: number;
  maxStateTokens?: number;
}

export const StateSnapshotProcessorOptionsSchema: JSONSchemaType<StateSnapshotProcessorOptions> =
  {
    type: 'object',
    properties: {
      target: {
        type: 'string',
        enum: ['incremental', 'freeNTokens', 'max'],
        nullable: true,
      },
      freeTokensTarget: { type: 'number', nullable: true },
      model: { type: 'string', nullable: true },
      systemInstruction: { type: 'string', nullable: true },
      maxSummaryTurns: { type: 'number', nullable: true },
      maxStateTokens: { type: 'number', nullable: true },
    },
    required: [],
  };

export function createStateSnapshotProcessor(
  id: string,
  env: ContextEnvironment,
  options: StateSnapshotProcessorOptions,
): ContextProcessor {
  const generator = new SnapshotGenerator(env);

  return {
    id,
    name: 'StateSnapshotProcessor',
    process: async ({ targets, inbox }: ProcessArgs) => {
      if (targets.length === 0) {
        return targets;
      }

      // Determine what mode we are looking for: 'incremental' -> 'point-in-time', 'max' -> 'accumulate'
      const strategy = options.target ?? 'max';
      const expectedType =
        strategy === 'incremental' ? 'point-in-time' : 'accumulate';

      // 1. Check Inbox for a completed Snapshot (The Fast Path)
      const proposedSnapshots = inbox.getMessages<{
        newText: string;
        consumedIds: string[];
        type: string;
        timestamp: number;
      }>('PROPOSED_SNAPSHOT');

      if (proposedSnapshots.length > 0) {
        // Filter for the snapshot type that matches our processor mode
        const matchingSnapshots = proposedSnapshots.filter(
          (s) => s.payload.type === expectedType,
        );

        // Sort by newest timestamp first (we want the most accumulated snapshot)
        const sorted = [...matchingSnapshots].sort(
          (a, b) => b.timestamp - a.timestamp,
        );

        for (const proposed of sorted) {
          const { consumedIds, newText, timestamp } = proposed.payload;

          // Verify all consumed IDs still exist sequentially in targets
          const targetIds = new Set(targets.map((t) => t.id));
          const isValid = consumedIds.every((id) => targetIds.has(id));

          if (isValid) {
            env.tracer.logEvent(
              'StateSnapshotProcessor',
              'Snapshot Spliced from Inbox',
              {
                snapshotText: newText,
              },
            );
            debugLogger.log(
              `[StateSnapshotProcessor] Successfully spliced PROPOSED_SNAPSHOT from Inbox into Graph. Consumed ${consumedIds.length} nodes.`,
            );
            // If valid, apply it!
            const newId = deriveStableId(consumedIds);

            const snapshotNode: Snapshot = {
              id: newId,
              turnId: newId,
              type: NodeType.SNAPSHOT,
              timestamp: timestamp ?? Date.now(),
              role: 'user',
              payload: { text: newText },
              abstractsIds: consumedIds,
            };

            // Remove the consumed nodes and insert the snapshot at the earliest index
            const returnedNodes = targets.filter(
              (t) => !consumedIds.includes(t.id),
            );
            const firstRemovedIdx = targets.findIndex((t) =>
              consumedIds.includes(t.id),
            );

            if (firstRemovedIdx !== -1) {
              const idx = Math.max(0, firstRemovedIdx);
              returnedNodes.splice(idx, 0, snapshotNode);
            } else {
              returnedNodes.unshift(snapshotNode);
            }

            inbox.consume(proposed.id);
            return returnedNodes;
          } else {
            debugLogger.log(
              `[StateSnapshotProcessor] Rejected PROPOSED_SNAPSHOT from Inbox because one or more target IDs were missing from the current graph window.`,
            );
          }
        }
      }

      // 2. The Synchronous Backstop (The Slow Path)
      let targetTokensToRemove = 0;

      if (strategy === 'incremental') {
        targetTokensToRemove = Infinity; // incremental implies removing as much as possible if no state is passed
      } else if (strategy === 'freeNTokens') {
        targetTokensToRemove = options.freeTokensTarget ?? Infinity;
      } else if (strategy === 'max') {
        targetTokensToRemove = Infinity;
      }

      let deficitAccumulator = 0;
      const nodesToSummarize: ConcreteNode[] = [];

      // Scan oldest to newest
      for (const node of targets) {
        nodesToSummarize.push(node);
        deficitAccumulator += env.tokenCalculator.getTokenCost(node);

        if (deficitAccumulator >= targetTokensToRemove) break;
      }

      if (nodesToSummarize.length < 2) return targets; // Not enough context

      let previousStateJson: string | undefined = undefined;
      let baselineIdToConsume: string | undefined = undefined;

      // Global Lookback: Find the absolute most recent snapshot anywhere in the active context
      const baseline = findLatestSnapshotBaseline(targets);

      if (baseline) {
        previousStateJson = baseline.text;
        // If the snapshot happens to be inside our summary window, remove it so the LLM doesn't read it as raw transcript
        const summaryIdx = nodesToSummarize.findIndex(
          (n) => n.id === baseline.id,
        );
        if (summaryIdx !== -1) {
          baselineIdToConsume = baseline.id;
          nodesToSummarize.splice(summaryIdx, 1);
        }
      } else {
        debugLogger.log(
          '[StateSnapshotProcessor] No previous snapshot found in context graph. Initializing new Master State baseline.',
        );
      }

      try {
        const snapshotText = await generator.synthesizeSnapshot(
          nodesToSummarize,
          previousStateJson,
          {
            maxSummaryTurns: options.maxSummaryTurns,
            maxStateTokens: options.maxStateTokens,
          },
        );

        env.tracer.logEvent('StateSnapshotProcessor', 'Snapshot Synthesized', {
          snapshotText,
        });

        const consumedIds = nodesToSummarize.map((n) => n.id);
        if (baselineIdToConsume && !consumedIds.includes(baselineIdToConsume)) {
          consumedIds.push(baselineIdToConsume);
        }
        const newId = deriveStableId(consumedIds);

        const snapshotNode: Snapshot = {
          id: newId,
          turnId: newId,
          type: NodeType.SNAPSHOT,
          timestamp: nodesToSummarize[nodesToSummarize.length - 1].timestamp,
          role: 'user',
          payload: { text: snapshotText },
          abstractsIds: [...consumedIds],
        };

        const returnedNodes = targets.filter(
          (t) => !consumedIds.includes(t.id),
        );
        const firstRemovedIdx = targets.findIndex((t) =>
          consumedIds.includes(t.id),
        );

        if (firstRemovedIdx !== -1) {
          const idx = Math.max(0, firstRemovedIdx);
          returnedNodes.splice(idx, 0, snapshotNode);
        } else {
          returnedNodes.unshift(snapshotNode);
        }

        return returnedNodes;
      } catch (e) {
        debugLogger.error('StateSnapshotProcessor failed sync backstop', e);
        return targets;
      }
    },
  };
}