// ---------------------------------------------------------------------------
// Memory Graph — LLM-based consolidation engine
//
// Runs daily (or on demand). Processes nodes in partitions:
// 1. Recency: nodes from last 7 days — merge duplicates, initial narrative
// 2. Significance: top N by significance — update narrative arcs
// 3. Random sample: cross-pollination and pattern detection
//
// Each partition is a separate LLM call. The LLM produces a MemoryDiff
// (same format as extraction) that is applied to the graph.
// ---------------------------------------------------------------------------

import { getConfiguredProvider } from "@vellumai/plugin-api";

import type { AssistantConfig } from "../../../../../config/types.js";
import { getMemoryDb } from "../../../../../persistence/db-connection.js";
import {
  createTrigger,
  deduplicateParagraphs,
  deleteNode,
  getEdgesForNode,
  getTriggersForNode,
  queryNodes,
  recordNodeEdit,
  updateNode,
} from "../../graph/store.js";
import type { MemoryNode } from "../../graph/types.js";
import { isCapabilityNode } from "../../graph/types.js";
import { BackendUnavailableError } from "../../host-utils.js";
import { extractToolUse, userMessage } from "../../llm-helpers.js";
import { getLogger } from "../../logging.js";
import {
  EVENT_DATE_PROMPT_RULES,
  formatAuthoritativeConversationTimestamp,
  parseEpochMs,
} from "./extraction.js";

const log = getLogger("graph-consolidation");

// ---------------------------------------------------------------------------
// Consolidation prompt
// ---------------------------------------------------------------------------

function buildConsolidationPrompt(
  partitionName: string,
  nodes: Array<{
    id: string;
    type: string;
    content: string;
    significance: number;
    fidelity: string;
    reinforcementCount: number;
    created: number;
    eventDate: number | null;
    hasImage: boolean;
  }>,
  edges: Array<{ sourceId: string; targetId: string; relationship: string }>,
): string {
  const nodeList = nodes
    .map((n) => {
      const age = Math.floor((Date.now() - n.created) / (1000 * 60 * 60 * 24));
      const eventStr =
        n.eventDate != null
          ? ` eventDate=${new Date(n.eventDate).toISOString().split("T")[0]}`
          : "";
      const imageStr = n.hasImage ? " [has_image]" : "";
      return `  [${n.id}] type=${n.type} sig=${n.significance.toFixed(
        2,
      )} fidelity=${n.fidelity} reinforced=${
        n.reinforcementCount
      }x age=${age}d${eventStr}${imageStr}\n    ${n.content}`;
    })
    .join("\n\n");

  const edgeList =
    edges.length > 0
      ? edges
          .map((e) => `  ${e.sourceId} --${e.relationship}--> ${e.targetId}`)
          .join("\n")
      : "  (none)";

  const currentTimestamp = formatAuthoritativeConversationTimestamp(Date.now());

  return `You are consolidating the "${partitionName}" partition of a memory graph. These are memories stored by an AI assistant about its conversations and relationship with a user.

## Authoritative Current Timestamp

${currentTimestamp}

${EVENT_DATE_PROMPT_RULES}

## Your Tasks

1. **Merge duplicates**: If two or more nodes describe the exact same specific event or fact with substantially the same details, merge them into one by:
   - Keeping the richer/more complete version (UPDATE it to incorporate details from duplicates)
   - DELETE the duplicates
   - Preserve the highest significance, reinforcement count, and stability from the merged nodes
   - Create a "supersedes" edge from the surviving node to each deleted node
   - Two nodes about the same person or topic but with DIFFERENT details, timestamps, or context are NOT duplicates — leave them both intact.

2. **Rewrite faded content**: For nodes at "faded" or "gist" fidelity, rewrite their content to be shorter and more abstract — like how a real memory fades. A "faded" memory should be 1-2 sentences. A "gist" memory should be one sentence capturing only the essence.

3. **Update narrative roles**: If a node is clearly a turning point, inciting incident, or thesis in a larger story arc, set its narrativeRole and partOfStory.

4. **Resolve stale prospective nodes**: If a prospective node (type=prospective) is older than 7 days and has no "resolved-by" edge, downgrade its fidelity to "gist" and rewrite it as a past observation (e.g. "Had planned to X" instead of "Need to X"). If the node has an event_date in the past, clear it by setting event_date to null.

## Constraints

- Do NOT create new nodes — consolidation only merges, updates, and rewrites
- Do NOT change a node's type
- Do NOT increase fidelity (memories only fade, never sharpen)
- Do NOT delete non-duplicate nodes — only delete the non-survivor in a merge. Fading and eventual cleanup are handled by the decay engine, not consolidation.
- Preserve first-person prose style in content rewrites
- When merging, keep the node with higher reinforcementCount as the survivor

## Current Nodes (${partitionName})

${nodeList}

## Current Edges

${edgeList}

Use the consolidate_diff tool to output your changes.`;
}

const CONSOLIDATE_TOOL_SCHEMA = {
  name: "consolidate_diff",
  description: "Output consolidation changes to the memory graph",
  input_schema: {
    type: "object" as const,
    properties: {
      updates: {
        type: "array" as const,
        description:
          "Nodes to update (content rewrites, narrative roles, fidelity downgrades)",
        items: {
          type: "object" as const,
          properties: {
            id: { type: "string" as const },
            content: {
              type: "string" as const,
              description: "New content (if rewriting)",
            },
            fidelity: {
              type: "string" as const,
              enum: ["vivid", "clear", "faded", "gist", "gone"],
            },
            narrativeRole: { type: "string" as const },
            partOfStory: { type: "string" as const },
            event_date: {
              type: ["number", "null"] as const,
              description:
                "Epoch ms of the event this memory describes. Resolve partial dates from the authoritative current timestamp and do not infer a prior year unless stated. Preserve from merged duplicates when the survivor lacks one. Set to null to clear a stale event date.",
            },
          },
          required: ["id"] as const,
        },
      },
      delete_ids: {
        type: "array" as const,
        description: "Node IDs to delete (merged duplicates only)",
        items: { type: "string" as const },
      },
      merge_edges: {
        type: "array" as const,
        description: "Supersedes edges for merged nodes (survivor → deleted)",
        items: {
          type: "object" as const,
          properties: {
            survivor_id: { type: "string" as const },
            deleted_id: { type: "string" as const },
          },
          required: ["survivor_id", "deleted_id"] as const,
        },
      },
    },
    required: ["updates", "delete_ids", "merge_edges"] as const,
  },
};

// ---------------------------------------------------------------------------
// Partition builders
// ---------------------------------------------------------------------------

function getRecentNodes(days: number = 7): MemoryNode[] {
  const cutoff = Date.now() - days * 24 * 60 * 60 * 1000;
  return queryNodes({
    fidelityNot: ["gone"],
    createdAfter: cutoff,
    limit: 10000,
  }).filter((n) => !isCapabilityNode(n));
}

function getTopSignificanceNodes(n: number = 50): MemoryNode[] {
  return queryNodes({
    fidelityNot: ["gone"],
    minSignificance: 0.6,
  })
    .filter((n) => !isCapabilityNode(n))
    .slice(0, n);
}

function getRandomSample(n: number = 30): MemoryNode[] {
  const all = queryNodes({
    fidelityNot: ["gone"],
    limit: 10000,
  }).filter((n) => !isCapabilityNode(n));
  // Fisher-Yates shuffle, take first n
  for (let i = all.length - 1; i > 0; i--) {
    const j = Math.floor(Math.random() * (i + 1));
    [all[i], all[j]] = [all[j], all[i]];
  }
  return all.slice(0, n);
}

// ---------------------------------------------------------------------------
// Run consolidation on a partition
// ---------------------------------------------------------------------------

const CHUNK_SIZE = 25;

// ---------------------------------------------------------------------------
// Duplicate detection — fast LLM call on compact listings
// ---------------------------------------------------------------------------

const DUPE_DETECT_TOOL = {
  name: "report_duplicate_groups",
  description:
    "Report groups of nodes that describe the same event, fact, or topic and should be merged",
  input_schema: {
    type: "object" as const,
    properties: {
      groups: {
        type: "array" as const,
        description:
          "Each group is a list of node IDs that are duplicates of each other",
        items: {
          type: "array" as const,
          items: { type: "string" as const },
        },
      },
    },
    required: ["groups"] as const,
  },
};

/**
 * Fast LLM pass to identify duplicate groups from a compact node listing.
 * Uses a latency-optimized model since it only needs to compare one-line
 * summaries, not reason about full prose. Returns groups of MemoryNodes
 * that should be consolidated together.
 */
async function identifyDuplicateGroups(
  nodes: MemoryNode[],
  _config: AssistantConfig,
): Promise<MemoryNode[][]> {
  if (nodes.length < 2) {
    return [];
  }

  const provider = await getConfiguredProvider("memoryConsolidation");
  if (!provider) {
    return [];
  }

  // Compact listing: ID + first 100 chars of content
  const listing = nodes
    .map((n) => {
      const preview =
        n.content.length > 100 ? n.content.slice(0, 100) + "…" : n.content;
      return `[${n.id}] ${preview}`;
    })
    .join("\n");

  const systemPrompt = `You are scanning a list of memory nodes for DUPLICATES — nodes that describe the exact same specific event or fact. Group duplicates together. Two nodes are duplicates ONLY if they describe the same underlying thing with substantially the same details. Be conservative — nodes about the same person or topic but with different details, timestamps, or context are NOT duplicates. Only include nodes that have at least one true duplicate.`;

  const response = await provider.sendMessage([userMessage(listing)], {
    tools: [DUPE_DETECT_TOOL],
    systemPrompt,
    config: {
      callSite: "memoryConsolidation" as const,
      tool_choice: { type: "tool" as const, name: "report_duplicate_groups" },
    },
  });

  const toolBlock = extractToolUse(response);
  if (!toolBlock) {
    return [];
  }

  const input = toolBlock.input as { groups?: string[][] };
  if (!input.groups) {
    return [];
  }

  const nodeMap = new Map(nodes.map((n) => [n.id, n]));

  return (input.groups ?? [])
    .map((ids) =>
      ids.filter((id) => nodeMap.has(id)).map((id) => nodeMap.get(id)!),
    )
    .filter((group) => group.length >= 2);
}

// ---------------------------------------------------------------------------
// Consolidation partition processing
// ---------------------------------------------------------------------------

interface ConsolidationPartitionResult {
  nodesUpdated: number;
  nodesDeleted: number;
  mergeEdgesCreated: number;
}

async function consolidatePartition(
  partitionName: string,
  nodes: MemoryNode[],
  config: AssistantConfig,
): Promise<ConsolidationPartitionResult> {
  const result: ConsolidationPartitionResult = {
    nodesUpdated: 0,
    nodesDeleted: 0,
    mergeEdgesCreated: 0,
  };

  if (nodes.length === 0) {
    return result;
  }

  // Step 1: Fast LLM call to identify duplicate groups from compact listing
  const dupeGroups = await identifyDuplicateGroups(nodes, config);
  // Collect all nodes that appear in a duplicate group
  const inGroup = new Set(dupeGroups.flat().map((n) => n.id));
  // Singletons: nodes not in any duplicate group
  const singletons = nodes.filter((n) => !inGroup.has(n.id));

  log.info(
    {
      partition: partitionName,
      nodeCount: nodes.length,
      dupeGroups: dupeGroups.length,
      singletons: singletons.length,
    },
    "Identified duplicate groups for consolidation",
  );

  // Step 2: Run full consolidation on each duplicate group
  const deleted = new Set<string>();
  for (let i = 0; i < dupeGroups.length; i++) {
    const chunk = dupeGroups[i].filter((n) => !deleted.has(n.id));
    if (chunk.length < 2) {
      continue;
    }

    const chunkResult = await consolidateChunk(
      `${partitionName} dupes (${i + 1}/${dupeGroups.length})`,
      chunk,
      config,
    );
    result.nodesUpdated += chunkResult.nodesUpdated;
    result.nodesDeleted += chunkResult.nodesDeleted;
    result.mergeEdgesCreated += chunkResult.mergeEdgesCreated;
    for (const id of chunkResult.deletedIds) {
      deleted.add(id);
    }
  }

  // Step 3: Run consolidation on singletons in chunks (for fidelity/narrative updates)
  const remainingSingletons = singletons.filter((n) => !deleted.has(n.id));
  if (remainingSingletons.length >= 2) {
    for (let i = 0; i < remainingSingletons.length; i += CHUNK_SIZE) {
      const chunk = remainingSingletons.slice(i, i + CHUNK_SIZE);
      if (chunk.length < 2) {
        continue;
      }

      const chunkResult = await consolidateChunk(
        `${partitionName} singles (${Math.floor(i / CHUNK_SIZE) + 1})`,
        chunk,
        config,
      );
      result.nodesUpdated += chunkResult.nodesUpdated;
      result.nodesDeleted += chunkResult.nodesDeleted;
      result.mergeEdgesCreated += chunkResult.mergeEdgesCreated;
      for (const id of chunkResult.deletedIds) {
        deleted.add(id);
      }
    }
  }

  return result;
}

interface ChunkResult extends ConsolidationPartitionResult {
  deletedIds: string[];
}

async function consolidateChunk(
  chunkName: string,
  nodes: MemoryNode[],
  _config: AssistantConfig,
): Promise<ChunkResult> {
  const result: ChunkResult = {
    nodesUpdated: 0,
    nodesDeleted: 0,
    mergeEdgesCreated: 0,
    deletedIds: [],
  };

  if (nodes.length === 0) {
    return result;
  }

  // Collect edges between partition nodes
  const nodeIds = new Set(nodes.map((n) => n.id));
  const edges: Array<{
    sourceId: string;
    targetId: string;
    relationship: string;
  }> = [];
  for (const node of nodes) {
    for (const edge of getEdgesForNode(node.id)) {
      if (nodeIds.has(edge.sourceNodeId) && nodeIds.has(edge.targetNodeId)) {
        edges.push({
          sourceId: edge.sourceNodeId,
          targetId: edge.targetNodeId,
          relationship: edge.relationship,
        });
      }
    }
  }

  // Deduplicate edges
  const edgeKeys = new Set<string>();
  const dedupedEdges = edges.filter((e) => {
    const key = `${e.sourceId}-${e.relationship}-${e.targetId}`;
    if (edgeKeys.has(key)) {
      return false;
    }
    edgeKeys.add(key);
    return true;
  });

  const provider = await getConfiguredProvider("memoryConsolidation");
  if (!provider) {
    throw new BackendUnavailableError("Provider unavailable for consolidation");
  }

  const systemPrompt = buildConsolidationPrompt(
    chunkName,
    nodes.map((n) => ({
      id: n.id,
      type: n.type,
      content: n.content,
      significance: n.significance,
      fidelity: n.fidelity,
      reinforcementCount: n.reinforcementCount,
      created: n.created,
      eventDate: n.eventDate,
      hasImage: n.imageRefs != null && n.imageRefs.length > 0,
    })),
    dedupedEdges,
  );

  const response = await provider.sendMessage(
    [
      userMessage(
        "Consolidate this partition. Focus on merging duplicates and fading old memories.",
      ),
    ],
    {
      tools: [CONSOLIDATE_TOOL_SCHEMA],
      systemPrompt,
      config: {
        callSite: "memoryConsolidation" as const,
        tool_choice: { type: "tool" as const, name: "consolidate_diff" },
      },
    },
  );

  const toolBlock = extractToolUse(response);
  if (!toolBlock) {
    log.warn("No tool_use block in consolidation response");
    return result;
  }

  const input = toolBlock.input as {
    updates?: Array<{
      id: string;
      content?: string;
      fidelity?: string;
      narrativeRole?: string;
      partOfStory?: string;
      event_date?: number | null;
    }>;
    delete_ids?: string[];
    merge_edges?: Array<{ survivor_id: string; deleted_id: string }>;
  };

  // Build nodeMap once upfront; patch entries after each updateNode() so
  // later iterations always read fresh in-memory state.
  const nodeMap = new Map(nodes.map((n) => [n.id, n]));

  // Apply updates
  for (const update of input.updates ?? []) {
    if (!nodeIds.has(update.id)) {
      continue;
    } // safety: only update nodes in this partition

    const changes: Partial<MemoryNode> = {};
    if (update.content) {
      changes.content = update.content;
    }
    if (update.fidelity) {
      changes.fidelity = update.fidelity as MemoryNode["fidelity"];
    }
    if (update.narrativeRole !== undefined) {
      changes.narrativeRole = update.narrativeRole || null;
    }
    if (update.partOfStory !== undefined) {
      changes.partOfStory = update.partOfStory || null;
    }
    if (update.event_date !== undefined) {
      changes.eventDate = parseEpochMs(update.event_date);
    }
    changes.lastConsolidated = Date.now();

    if (Object.keys(changes).length > 1) {
      // more than just lastConsolidated

      // Wrap edit recording + node update in a transaction so they are atomic:
      // if updateNode fails, the edit record is rolled back. Both writes land on
      // the memory connection now the graph cluster lives there, so the
      // transaction must be opened on that connection to stay atomic.
      const mdb = getMemoryDb();
      if (!mdb) {
        throw new Error(
          "memory database unavailable — graph consolidation cannot proceed",
        );
      }
      mdb.transaction(() => {
        if (changes.content) {
          const cleanContent = deduplicateParagraphs(changes.content);
          const node = nodeMap.get(update.id);
          if (node && node.content !== cleanContent) {
            recordNodeEdit({
              nodeId: update.id,
              previousContent: node.content,
              newContent: cleanContent,
              source: "consolidation",
            });
          }
        }

        updateNode(update.id, changes);
      });
      result.nodesUpdated++;
      // Sync in-memory state with what updateNode actually wrote to the DB
      // (updateNode deduplicates content before persisting)
      if (changes.content) {
        changes.content = deduplicateParagraphs(changes.content);
      }
      const node = nodeMap.get(update.id);
      if (node) {
        Object.assign(node, changes);
      }
    }
  }

  // Apply merge edges (before deletion so the edge can reference the node)
  const { createEdge } = await import("../../graph/store.js");
  for (const merge of input.merge_edges ?? []) {
    if (!nodeIds.has(merge.survivor_id) || !nodeIds.has(merge.deleted_id)) {
      continue;
    }
    try {
      createEdge({
        sourceNodeId: merge.survivor_id,
        targetNodeId: merge.deleted_id,
        relationship: "supersedes",
        weight: 1.0,
        created: Date.now(),
      });
      result.mergeEdgesCreated++;

      // Preserve eventDate from deleted node if survivor doesn't have one
      const survivor = nodeMap.get(merge.survivor_id);
      const deleted = nodeMap.get(merge.deleted_id);
      if (
        survivor &&
        deleted?.eventDate != null &&
        survivor.eventDate == null
      ) {
        updateNode(merge.survivor_id, { eventDate: deleted.eventDate });
        survivor.eventDate = deleted.eventDate;

        // The deleted node's triggers will be cascade-deleted when the node
        // is removed. Ensure the survivor has an event trigger for the
        // inherited eventDate (updateNode only syncs existing triggers).
        const survivorTriggers = getTriggersForNode(merge.survivor_id);
        if (!survivorTriggers.some((t) => t.type === "event")) {
          createTrigger({
            nodeId: merge.survivor_id,
            type: "event",
            schedule: null,
            condition: null,
            conditionEmbedding: null,
            threshold: null,
            eventDate: deleted.eventDate,
            rampDays: 7,
            followUpDays: 2,
            recurring: false,
            consumed: false,
            cooldownMs: null,
            lastFired: null,
          });
        }
      }

      // Preserve imageRefs from deleted node if survivor doesn't have any
      if (
        survivor &&
        deleted?.imageRefs != null &&
        deleted.imageRefs.length > 0 &&
        (survivor.imageRefs == null || survivor.imageRefs.length === 0)
      ) {
        updateNode(merge.survivor_id, { imageRefs: deleted.imageRefs });
        survivor.imageRefs = deleted.imageRefs;
      }
    } catch (err) {
      log.warn({ err }, "Failed to create merge edge");
    }
  }

  // Apply deletions
  for (const id of input.delete_ids ?? []) {
    if (!nodeIds.has(id)) {
      continue;
    } // safety
    try {
      log.info({ nodeId: id }, "Consolidation deleting node");
      deleteNode(id);
      result.nodesDeleted++;
      result.deletedIds.push(id);
    } catch (err) {
      log.warn({ id, err }, "Failed to delete node during consolidation");
    }
  }

  return result;
}

// ---------------------------------------------------------------------------
// Full consolidation run
// ---------------------------------------------------------------------------

export interface ConsolidationResult {
  partitions: Record<string, ConsolidationPartitionResult>;
  totalUpdated: number;
  totalDeleted: number;
  totalMergeEdges: number;
  latencyMs: number;
}

export async function runConsolidation(
  config: AssistantConfig,
): Promise<ConsolidationResult> {
  const start = Date.now();

  const result: ConsolidationResult = {
    partitions: {},
    totalUpdated: 0,
    totalDeleted: 0,
    totalMergeEdges: 0,
    latencyMs: 0,
  };

  // Define partitions
  const partitions: Array<{ name: string; nodes: MemoryNode[] }> = [
    { name: "recent", nodes: getRecentNodes() },
    { name: "significant", nodes: getTopSignificanceNodes() },
    { name: "random", nodes: getRandomSample() },
  ];

  for (const partition of partitions) {
    if (partition.nodes.length === 0) {
      log.info({ partition: partition.name }, "Empty partition, skipping");
      continue;
    }

    log.info(
      { partition: partition.name, nodeCount: partition.nodes.length },
      "Consolidating partition",
    );

    try {
      const partitionResult = await consolidatePartition(
        partition.name,
        partition.nodes,
        config,
      );

      result.partitions[partition.name] = partitionResult;
      result.totalUpdated += partitionResult.nodesUpdated;
      result.totalDeleted += partitionResult.nodesDeleted;
      result.totalMergeEdges += partitionResult.mergeEdgesCreated;

      log.info(
        {
          partition: partition.name,
          updated: partitionResult.nodesUpdated,
          deleted: partitionResult.nodesDeleted,
          mergeEdges: partitionResult.mergeEdgesCreated,
        },
        "Partition consolidation complete",
      );
    } catch (err) {
      log.warn(
        {
          partition: partition.name,
          err: err instanceof Error ? err.message : String(err),
        },
        "Partition consolidation failed",
      );
      result.partitions[partition.name] = {
        nodesUpdated: 0,
        nodesDeleted: 0,
        mergeEdgesCreated: 0,
      };
    }
  }

  result.latencyMs = Date.now() - start;

  log.info(
    {
      totalUpdated: result.totalUpdated,
      totalDeleted: result.totalDeleted,
      totalMergeEdges: result.totalMergeEdges,
      latencyMs: result.latencyMs,
    },
    "Full consolidation complete",
  );

  return result;
}