/**
 * @license
 * Copyright 2026 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */
import type { ConcreteNode } from '../graph/types.js';
import { NodeType } from '../graph/types.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
import { LlmRole } from '../../telemetry/llmRole.js';
import { formatNodesForLlm } from './formatNodesForLlm.js';
import { randomUUID } from 'node:crypto';
import { isRecord } from '../../utils/markdownUtils.js';

function isStringArray(value: unknown): value is string[] {
  return (
    Array.isArray(value) && value.every((item) => typeof item === 'string')
  );
}

function isNumberArray(value: unknown): value is number[] {
  return (
    Array.isArray(value) && value.every((item) => typeof item === 'number')
  );
}

function isTaskArray(
  value: unknown,
): value is Array<{ id: string; description: string }> {
  return (
    Array.isArray(value) &&
    value.every((item) => {
      if (!isRecord(item)) return false;
      const id = item['id'];
      const desc = item['description'];
      return typeof id === 'string' && typeof desc === 'string';
    })
  );
}

function isString(value: unknown): value is string {
  return typeof value === 'string';
}

export interface SnapshotState {
  active_tasks: Array<{ id: string; description: string }>;
  discovered_facts: string[];
  constraints_and_preferences: string[];
  recent_arc: string[];
}

import { debugLogger } from '../../utils/debugLogger.js';

export function isSnapshotState(text: string): boolean {
  const trimmed = text.trim();
  if (!trimmed.startsWith('{') || !trimmed.endsWith('}')) {
    return false;
  }
  try {
    const parsed: unknown = JSON.parse(trimmed);
    if (!isRecord(parsed)) return false;
    const isSnap =
      Array.isArray(parsed['active_tasks']) &&
      Array.isArray(parsed['discovered_facts']) &&
      Array.isArray(parsed['constraints_and_preferences']) &&
      Array.isArray(parsed['recent_arc']);
    if (!isSnap) {
      debugLogger.log(
        '[isSnapshotState] FAILED FOR JSON:',
        JSON.stringify(parsed),
      );
    }
    return isSnap;
  } catch {
    debugLogger.log('[isSnapshotState] PARSE FAILED FOR:', trimmed);
    return false;
  }
}

export interface BaselineSnapshotInfo {
  text: string;
  abstractsIds: string[];
  id: string;
  timestamp: number;
}

/**
 * Global Lookback: Scans the target nodes in reverse to find the absolute
 * most recent valid Snapshot node to use as a Delta baseline.
 */
export function findLatestSnapshotBaseline(
  targets: readonly ConcreteNode[],
): BaselineSnapshotInfo | undefined {
  debugLogger.log(
    '[findLatestSnapshotBaseline] Targets:',
    targets.map((t) => ({
      id: t.id,
      type: t.type,
      text:
        t.payload &&
        typeof t.payload === 'object' &&
        'text' in t.payload &&
        typeof t.payload.text === 'string'
          ? t.payload.text.substring(0, 20)
          : '',
    })),
  );
  const lastSnapshotNode = [...targets]
    .reverse()
    .find((n) => n.type === NodeType.SNAPSHOT && n.payload.text);

  if (lastSnapshotNode?.payload.text) {
    return {
      text: lastSnapshotNode.payload.text,
      abstractsIds: lastSnapshotNode.abstractsIds
        ? [...lastSnapshotNode.abstractsIds]
        : [],
      id: lastSnapshotNode.id,
      timestamp: lastSnapshotNode.timestamp,
    };
  }

  return undefined;
}

export class SnapshotGenerator {
  constructor(private readonly env: ContextEnvironment) {}

  async synthesizeSnapshot(
    nodes: readonly ConcreteNode[],
    previousStateJson?: string,
    options: { maxSummaryTurns?: number; maxStateTokens?: number } = {},
  ): Promise<string> {
    const emptyState: SnapshotState = {
      active_tasks: [],
      discovered_facts: [],
      constraints_and_preferences: [],
      recent_arc: [],
    };

    let previousState = emptyState;
    if (previousStateJson) {
      try {
        const parsed = JSON.parse(previousStateJson) as unknown;
        if (isRecord(parsed)) {
          let loadedArc: string[] = [];
          if (isStringArray(parsed['recent_arc'])) {
            loadedArc = parsed['recent_arc'];
          } else if (isString(parsed['summary']) && parsed['summary']) {
            // Migrate legacy v1 summary to V2 recent_arc array
            loadedArc = [parsed['summary']];
          }

          previousState = {
            active_tasks: isTaskArray(parsed['active_tasks'])
              ? parsed['active_tasks']
              : [],
            discovered_facts: isStringArray(parsed['discovered_facts'])
              ? parsed['discovered_facts']
              : [],
            constraints_and_preferences: isStringArray(
              parsed['constraints_and_preferences'],
            )
              ? parsed['constraints_and_preferences']
              : [],
            recent_arc: loadedArc,
          };
        }
      } catch {
        // Fallback to empty if parse fails
      }
    }
    let pressureWarning = '';
    const stateString = JSON.stringify(previousState);
    const estimatedTokens =
      this.env.tokenCalculator.estimateTokensForString(stateString);
    const maxTokens = options.maxStateTokens ?? 4000;

    if (estimatedTokens > maxTokens * 0.8) {
      pressureWarning = `\n\n[CRITICAL WARNING]: The Master State is currently at ${((estimatedTokens / maxTokens) * 100).toFixed(0)}% of its maximum capacity! You MUST aggressively prune obsolete, irrelevant, or overly granular facts and constraints using \`obsolete_fact_indices\` and \`obsolete_constraint_indices\`.`;
    }

    const systemPrompt = `You are an expert Context Memory Manager. You maintain the long-term "Master State" of the AI agent's memory.
You will be provided with the CURRENT Master State and a raw transcript of new conversation turns.
Your task is to generate a JSON Delta Patch representing what has changed in the transcript.${pressureWarning}

CRITICAL OPERATIONAL RULES:
1. FACTS: Extract explicit empirical facts (file paths, exact error codes, specific configs).
2. PRUNING: Keep facts dense. Use obsolete indices to aggressively delete facts that are no longer relevant to the current objective.
3. TASKS: Add any new active user requests to "new_tasks".
4. TASK RESOLUTION: A task may ONLY be placed in "resolved_task_ids" if a success message or explicit confirmation was provided in the transcript. If the task was being worked on but no final confirmation exists, it MUST remain active. Do not prematurely resolve tasks.`;

    const userPromptText = `CURRENT MASTER STATE:
${JSON.stringify(previousState, null, 2)}

TRANSCRIPT OF NEW TURNS:
${formatNodesForLlm(nodes)}`;

    const patchSchema = {
      type: 'object',
      properties: {
        new_facts: {
          type: 'array',
          items: { type: 'string' },
          description:
            'New specific, empirical facts discovered in this transcript chunk.',
        },
        new_constraints: {
          type: 'array',
          items: { type: 'string' },
          description:
            'New specific rules or instructions provided by the user in this chunk.',
        },
        new_tasks: {
          type: 'array',
          items: {
            type: 'object',
            properties: {
              description: {
                type: 'string',
                description: 'The task goal/description.',
              },
            },
          },
        },
        resolved_task_ids: {
          type: 'array',
          items: { type: 'string' },
          description:
            'IDs of tasks from the CURRENT MASTER STATE that were explicitly completed or abandoned in this transcript chunk.',
        },
        obsolete_fact_indices: {
          type: 'array',
          items: { type: 'number' },
          description:
            'Array indices of facts from CURRENT MASTER STATE that are no longer true or relevant and should be deleted.',
        },
        obsolete_constraint_indices: {
          type: 'array',
          items: { type: 'number' },
          description:
            'Array indices of constraints from CURRENT MASTER STATE that are no longer true or relevant and should be deleted.',
        },
        chronological_summary: {
          type: 'string',
          description:
            'A 1-2 sentence summary of the mechanical actions taken in this transcript chunk.',
        },
      },
    };

    let patch: Record<string, unknown> = {};
    try {
      const result = await this.env.llmClient.generateJson({
        role: LlmRole.UTILITY_STATE_SNAPSHOT_PROCESSOR,
        modelConfigKey: { model: 'context-snapshotter' },
        contents: [{ role: 'user', parts: [{ text: userPromptText }] }],
        systemInstruction: { role: 'system', parts: [{ text: systemPrompt }] },
        schema: patchSchema,
        promptId: this.env.promptId,
        abortSignal: new AbortController().signal,
      });
      if (isRecord(result)) {
        patch = result;
      }
    } catch {
      // If generateJson fails, return the unmodified previous state
      return JSON.stringify(previousState);
    }

    // Merging Application Logic (The Safeguard)
    const newState: SnapshotState = {
      active_tasks: [...previousState.active_tasks],
      discovered_facts: [...previousState.discovered_facts],
      constraints_and_preferences: [
        ...previousState.constraints_and_preferences,
      ],
      recent_arc: [...previousState.recent_arc],
    };

    // 1. Process Deletions (Resolved Tasks & Obsolete Items)
    const resolvedIds = patch['resolved_task_ids'];
    if (isStringArray(resolvedIds)) {
      const resolvedSet = new Set(resolvedIds);
      newState.active_tasks = newState.active_tasks.filter(
        (t) => !resolvedSet.has(t.id),
      );
    }

    const obsFacts = patch['obsolete_fact_indices'];
    if (isNumberArray(obsFacts)) {
      const dropSet = new Set(obsFacts);
      newState.discovered_facts = newState.discovered_facts.filter(
        (_, i) => !dropSet.has(i),
      );
    }

    const obsConstraints = patch['obsolete_constraint_indices'];
    if (isNumberArray(obsConstraints)) {
      const dropSet = new Set(obsConstraints);
      newState.constraints_and_preferences =
        newState.constraints_and_preferences.filter((_, i) => !dropSet.has(i));
    }

    // 2. Process Additions
    const newTasks = patch['new_tasks'];
    if (Array.isArray(newTasks)) {
      for (const t of newTasks) {
        if (isRecord(t)) {
          const desc = t['description'];
          if (typeof desc === 'string' && desc) {
            newState.active_tasks.push({
              id: `task_${randomUUID().slice(0, 8)}`,
              description: desc,
            });
          }
        }
      }
    }

    const newFacts = patch['new_facts'];
    if (isStringArray(newFacts)) {
      newState.discovered_facts.push(...newFacts);
    }

    const newConstraints = patch['new_constraints'];
    if (isStringArray(newConstraints)) {
      newState.constraints_and_preferences.push(...newConstraints);
    }

    // 3. Update Summary (Rolling Window)
    const chronoSummary = patch['chronological_summary'];
    if (typeof chronoSummary === 'string' && chronoSummary) {
      newState.recent_arc.push(chronoSummary);
      const maxTurns = options.maxSummaryTurns ?? 5;
      if (newState.recent_arc.length > maxTurns) {
        newState.recent_arc = newState.recent_arc.slice(-maxTurns);
      }
    }

    // 4. Enforce Token Budget (Structured Pruning Backstop)
    let currentTokens = this.env.tokenCalculator.estimateTokensForString(
      JSON.stringify(newState),
    );
    while (currentTokens > maxTokens) {
      // Priority 1: Drop oldest facts
      if (newState.discovered_facts.length > 0) {
        newState.discovered_facts.shift();
      }
      // Priority 2: Drop oldest constraints
      else if (newState.constraints_and_preferences.length > 0) {
        newState.constraints_and_preferences.shift();
      }
      // Priority 3: Drop oldest narrative arc
      else if (newState.recent_arc.length > 0) {
        newState.recent_arc.shift();
      }
      // Priority 4: Drop oldest active tasks (Pathological emergency)
      else if (newState.active_tasks.length > 0) {
        newState.active_tasks.shift();
      }
      // Priority 5: The state is completely empty, break to avoid infinite loop
      else {
        break;
      }

      currentTokens = this.env.tokenCalculator.estimateTokensForString(
        JSON.stringify(newState),
      );
    }

    return JSON.stringify(newState);
  }
}

/**
 * Shared logic for working with Snapshot node state.
 */
export class SnapshotStateHelper {
  /**
   * Flatten nested abstract IDs to only the "pristine" (non-snapshot) IDs.
   */
  static flattenAbstracts(
    nodes: ConcreteNode[],
    abstractsIds: readonly string[],
  ): string[] {
    const pristineIds: string[] = [];
    const nodeMap = new Map(nodes.map((n) => [n.id, n]));

    const walk = (ids: readonly string[]) => {
      for (const id of ids) {
        const node = nodeMap.get(id);
        if (!node) {
          // Fallback: if node not in map, treat as pristine ID
          pristineIds.push(id);
          continue;
        }

        if (node.type === NodeType.SNAPSHOT && node.abstractsIds) {
          walk(node.abstractsIds);
        } else {
          pristineIds.push(id);
        }
      }
    };

    walk(abstractsIds);
    return Array.from(new Set(pristineIds)); // Dedupe
  }

  /**
   * Helper to extract state from the most recent snapshot in a list of nodes.
   */
  static exportState(nodes: ConcreteNode[]): {
    snapshot?: { text: string; consumedIds: string[] };
  } {
    const baseline = findLatestSnapshotBaseline(nodes);
    if (!baseline) return {};

    const node = nodes.find((n) => n.id === baseline.id);
    if (!node || node.type !== NodeType.SNAPSHOT) return {};

    const consumedIds = this.flattenAbstracts(nodes, node.abstractsIds || []);

    return {
      snapshot: {
        text: baseline.text,
        consumedIds,
      },
    };
  }
}