/**
 * Semantic chat-memory store on Azure Cosmos DB (NoSQL vector search).
 *
 * Each memory is one item in the `memory` container:
 *  - Partition key `/sessionId` (partition-high-cardinality + partition-immutable-key):
 *    a whole conversation lives in one logical partition, so recall within a
 *    session is a single-partition query.
 *  - `type: "memory"` discriminator (model-type-discriminator) so the container
 *    could hold other shapes later.
 *  - `embedding`: a 1536-dim vector (text-embedding-3-small) indexed by a vector
 *    index; excluded from the regular index to keep write RU/latency low.
 *
 * Recall uses VectorDistance() with ORDER BY + a literal TOP (query-top-literal),
 * parameterized on the query vector (query-parameterize) for plan caching.
 */
import {
  VectorEmbeddingDataType,
  VectorEmbeddingDistanceFunction,
  VectorIndexType,
  type Container,
  type Database,
} from "@azure/cosmos";
import { randomUUID } from "node:crypto";
import { getCosmosClient, cosmosConfig } from "./cosmos";
import { EMBEDDING_DIMENSIONS, embedText } from "./embeddings";

const MEMORY_CONTAINER_ID = process.env.MEMORY_CONTAINER_ID ?? "memory";

export type MemoryRole = "user" | "assistant" | "note";

export interface MemoryItem {
  id: string;
  sessionId: string;
  type: "memory";
  role: MemoryRole;
  content: string;
  embedding: number[];
  createdAt: string;
  schemaVersion: number;
}

export interface MemoryHit {
  content: string;
  role: MemoryRole;
  createdAt: string;
  /** Cosine distance in [0,2]; smaller = more similar. */
  score: number;
}

export const memoryConfig = {
  containerId: MEMORY_CONTAINER_ID,
  partitionKeyPath: "/sessionId",
  embeddingPath: "/embedding",
  dimensions: EMBEDDING_DIMENSIONS,
} as const;

/**
 * Container definition used by the seed script. The vector embedding policy is
 * immutable after creation, so it must be set at create time.
 */
export function memoryContainerDefinition() {
  return {
    id: MEMORY_CONTAINER_ID,
    partitionKey: { paths: [memoryConfig.partitionKeyPath] },
    // Enable TTL (no default expiry). Set a per-item `ttl` to auto-expire
    // ephemeral memories with zero cleanup cost.
    defaultTtl: -1,
    vectorEmbeddingPolicy: {
      vectorEmbeddings: [
        {
          path: memoryConfig.embeddingPath,
          dataType: VectorEmbeddingDataType.Float32,
          dimensions: memoryConfig.dimensions,
          distanceFunction: VectorEmbeddingDistanceFunction.Cosine,
        },
      ],
    },
    indexingPolicy: {
      includedPaths: [{ path: "/*" }],
      // CRITICAL: keep the raw vector out of the regular index.
      excludedPaths: [{ path: `${memoryConfig.embeddingPath}/*` }],
      vectorIndexes: [
        {
          path: memoryConfig.embeddingPath,
          // QuantizedFlat: fast builds, ideal for < 50K vectors.
          type: VectorIndexType.QuantizedFlat,
        },
      ],
    },
  };
}

export async function ensureMemoryContainer(database: Database): Promise<void> {
  await database.containers.createIfNotExists(memoryContainerDefinition());
}

function getMemoryContainer(): Container {
  return getCosmosClient()
    .database(cosmosConfig.databaseId)
    .container(MEMORY_CONTAINER_ID);
}

export interface AppendMemoryInput {
  sessionId: string;
  role: MemoryRole;
  content: string;
  /** Optional TTL in seconds for ephemeral memories. */
  ttlSeconds?: number;
}

/** Store one memory, embedding its content for later semantic recall. */
export async function appendMemory(
  input: AppendMemoryInput,
): Promise<{ id: string }> {
  const embedding = await embedText(input.content);
  const now = new Date().toISOString();
  const item: MemoryItem & { ttl?: number } = {
    id: randomUUID(),
    sessionId: input.sessionId,
    type: "memory",
    role: input.role,
    content: input.content,
    embedding,
    createdAt: now,
    schemaVersion: 1,
    ...(input.ttlSeconds ? { ttl: input.ttlSeconds } : {}),
  };

  const container = getMemoryContainer();
  await container.items.create(item);
  return { id: item.id };
}

export interface RecallInput {
  query: string;
  /** Scope recall to one conversation (single-partition). Omit for global. */
  sessionId?: string;
  /** Max hits to return (clamped to 1..20). */
  k?: number;
}

/** Semantic recall: nearest memories to `query` by cosine distance. */
export async function recallMemories(input: RecallInput): Promise<MemoryHit[]> {
  const k = Math.min(Math.max(Math.trunc(input.k ?? 5), 1), 20);
  const vector = await embedText(input.query);

  const scoped = Boolean(input.sessionId);
  // TOP must be a literal integer (query-top-literal); `k` is a sanitized int.
  const sql =
    `SELECT TOP ${k} c.role, c.content, c.createdAt, ` +
    `VectorDistance(c.embedding, @vec) AS score ` +
    `FROM c WHERE c.type = "memory"` +
    (scoped ? ` AND c.sessionId = @sessionId` : ``) +
    ` ORDER BY VectorDistance(c.embedding, @vec)`;

  const parameters = [
    { name: "@vec", value: vector },
    ...(scoped ? [{ name: "@sessionId", value: input.sessionId! }] : []),
  ];

  const container = getMemoryContainer();
  const { resources } = await container.items
    .query<MemoryHit>(
      { query: sql, parameters },
      // Single-partition when scoped to a session; keeps recall cheap.
      scoped ? { partitionKey: input.sessionId } : undefined,
    )
    .fetchAll();

  return resources;
}