import { randomUUID } from "node:crypto";
import { QdrantClient } from "@qdrant/js-client-rest";

const COLLECTION = "assistant_memory";
const client = new QdrantClient({
  url: process.env.QDRANT_URL!,
  apiKey: process.env.QDRANT_API_KEY!,
});

type MemoryPayload = {
  tenant_id: string;
  account_id: string;
  user_id: string;
  text: string;
  source_url?: string;
};

// Run once during deployment or collection setup.
export async function setupAssistantMemory(vectorSize: number) {
  const { exists } = await client.collectionExists(COLLECTION);
  if (!exists) {
    await client.createCollection(COLLECTION, {
      vectors: { size: vectorSize, distance: "Cosine" },
    });
  }

  const collection = await client.getCollection(COLLECTION);
  const indexes = [
    {
      field_name: "tenant_id",
      field_schema: { type: "keyword" as const, is_tenant: true },
    },
    { field_name: "account_id", field_schema: "keyword" as const },
    { field_name: "user_id", field_schema: "keyword" as const },
  ];

  for (const index of indexes) {
    if (!collection.payload_schema[index.field_name]) {
      await client.createPayloadIndex(COLLECTION, { ...index, wait: true });
    }
  }
}

export async function remember(input: {
  embedding: number[];
  tenantId: string;
  accountId: string;
  userId: string;
  text: string;
  sourceUrl?: string;
}) {
  await client.upsert(COLLECTION, {
    wait: true,
    points: [{
      id: randomUUID(),
      vector: input.embedding,
      payload: {
        tenant_id: input.tenantId,
        account_id: input.accountId,
        user_id: input.userId,
        text: input.text,
        source_url: input.sourceUrl,
      } satisfies MemoryPayload,
    }],
  });
}

export async function recall(input: {
  embedding: number[];
  tenantId: string;
  userId?: string;
  limit?: number;
}) {
  const must = [
    { key: "tenant_id", match: { value: input.tenantId } },
    ...(input.userId
      ? [{ key: "user_id", match: { value: input.userId } }]
      : []),
  ];

  const result = await client.query(COLLECTION, {
    query: input.embedding,
    filter: { must },
    with_payload: true,
    limit: input.limit ?? 8,
  });

  return result.points.map((point) => ({
    id: point.id,
    score: point.score,
    payload: point.payload as MemoryPayload | null,
  }));
}

export async function deleteAccountMemory(input: {
  tenantId: string;
  accountId: string;
}) {
  await client.delete(COLLECTION, {
    wait: true,
    filter: {
      must: [
        { key: "tenant_id", match: { value: input.tenantId } },
        { key: "account_id", match: { value: input.accountId } },
      ],
    },
  });
}