// ingestEmbeddings.ts
import 'dotenv/config';
import fs from 'fs';
import path from 'path';
import { MongoClient, type Db, type Document, type WithId } from 'mongodb';
import OpenAI from 'openai';
import { finished } from 'stream/promises';
import fsPromises from 'fs/promises';

type Projection = Record<string, 0 | 1>;

type EmbeddingRow = {
  _id: string;
  embedding: number[];
  norm: number;
  fp: string;
  model: string;
};

type PendingRow = {
  id: string;
  text: string;
  fp: string;
};

type CollectionConfig<TDoc extends Document> = {
  name: string;
  projection: Projection;
  toText: (d: WithId<TDoc>) => string; // ✅ WithId
  fpFields: (d: WithId<TDoc>) => unknown[]; // ✅ WithId
};

type AttractionDoc = Document & {
  title?: string | null;
  description?: string | null;
  formatted_address?: string | null;
  venue?: string | null;
  city_level?: string | null;
  country_short?: string | null;
  updated_at?: unknown;
};

type EventDoc = Document & {
  title?: string | null;
  description?: string | null;
  formatted_address?: string | null;
  venue?: string | null;
  start_date?: string | null;
  end_date?: string | null;
  start_time?: string | null;
  end_time?: string | null;
  updated_at?: unknown;
};

const mongoUri = process.env.MONGODB_URI ?? '';
const openAiKey = process.env.OPEN_AI_KEY ?? '';
const EMBED_MODEL = process.env.EMBED_MODEL ?? 'text-embedding-3-small';
const OUT_DIR = path.join(process.cwd(), 'embeddings');

if (!mongoUri) throw new Error('Missing env var: MONGODB_URI');
if (!openAiKey) throw new Error('Missing env var: OPEN_AI_KEY');

const client = new MongoClient(mongoUri);
const openai = new OpenAI({ apiKey: openAiKey });

// ---- Per-collection config: projection + toText + fingerprint extras ----
const COLLECTIONS: Array<CollectionConfig<any>> = [
  {
    name: 'attraction',
    projection: {
      title: 1,
      description: 1,
      formatted_address: 1,
      venue: 1,
      city_level: 1,
      country_short: 1,
      updated_at: 1,
    },
    toText: (d: WithId<AttractionDoc>) =>
      [
        d.title,
        d.description,
        d.formatted_address,
        d.venue,
        d.city_level,
        d.country_short,
      ]
        .filter(Boolean)
        .join('\n'),
    fpFields: (d: WithId<AttractionDoc>) => [d.updated_at ?? null],
  },
  {
    name: 'event',
    projection: {
      title: 1,
      description: 1,
      formatted_address: 1,
      venue: 1,
      start_date: 1,
      end_date: 1,
      start_time: 1,
      end_time: 1,
      updated_at: 1,
    },
    toText: (d: WithId<EventDoc>) => {
      const when = [d.start_date, d.start_time, d.end_date, d.end_time]
        .filter(Boolean)
        .join(' ');
      return [d.title, d.description, d.formatted_address, d.venue, when]
        .filter(Boolean)
        .join('\n');
    },
    fpFields: (d: WithId<EventDoc>) => [
      d.updated_at ?? null,
      d.start_date ?? null,
      d.end_date ?? null,
      d.start_time ?? null,
      d.end_time ?? null,
    ],
  },
];

// ---- utils ----
const l2 = (v: number[]): number => Math.sqrt(v.reduce((s, x) => s + x * x, 0));

const embedBatch = async (texts: string[]): Promise<number[][]> => {
  const r = await openai.embeddings.create({
    model: EMBED_MODEL,
    input: texts,
  });
  return r.data.map((d) => d.embedding);
};

// Compacts the JSONL by writing a fresh file from a Map<string, row>
// Waits for the stream to finish before renaming (avoids ENOENT).
const writeJsonlCompact = async (
  filePath: string,
  map: Map<string, EmbeddingRow>,
): Promise<void> => {
  const dir = path.dirname(filePath);
  await fsPromises.mkdir(dir, { recursive: true });

  const tmp = `${filePath}.tmp`;
  const out = fs.createWriteStream(tmp, { flags: 'w' });

  if (map.size === 0) {
    out.end();
    await finished(out);
  } else {
    for (const row of map.values()) {
      if (!out.write(JSON.stringify(row) + '\n')) {
        await new Promise<void>((resolve) => out.once('drain', resolve));
      }
    }
    out.end();
    await finished(out);
  }

  await fsPromises.rename(tmp, filePath);
};

const ingestCollection = async <TDoc extends Document>(
  db: Db,
  cfg: CollectionConfig<TDoc>,
): Promise<void> => {
  const { name, projection, toText, fpFields } = cfg;

  fs.mkdirSync(OUT_DIR, { recursive: true });
  const outPath = path.join(OUT_DIR, `${name}.jsonl`);

  // Load existing rows into memory (for compaction + updates)
  const rows = new Map<string, EmbeddingRow>(); // _id_str -> row
  if (fs.existsSync(outPath)) {
    const raw = fs.readFileSync(outPath, 'utf8').trim();
    const lines = raw ? raw.split('\n').filter(Boolean) : [];
    for (const line of lines) {
      const row = JSON.parse(line) as EmbeddingRow;
      rows.set(row._id, row);
    }
  }

  const coll = db.collection<TDoc>(name);
  const cursor = coll.find({}, { projection });

  // Gather docs that need embedding (new or changed)
  const pending: PendingRow[] = [];

  while (await cursor.hasNext()) {
    const doc = await cursor.next(); // doc: WithId<TDoc> | null
    if (!doc) break;

    const text = toText(doc).trim();
    if (!text) continue;

    const id = String(doc._id);

    const fp = JSON.stringify([
      ...fpFields(doc), // important change markers
      text.length,
      text.slice(0, 64), // small content signature
    ]);

    if (rows.get(id)?.fp === fp) continue; // unchanged → skip
    pending.push({ id, text, fp });
  }

  if (!pending.length) {
    console.log(`[${name}] up to date.`);
    await writeJsonlCompact(outPath, rows);
    return;
  }

  // Embed in batches
  for (let i = 0; i < pending.length; i += 128) {
    const batch = pending.slice(i, i + 128);
    const vecs = await embedBatch(batch.map((b) => b.text));

    vecs.forEach((v, j) => {
      const id = batch[j].id;
      const fp = batch[j].fp;
      rows.set(id, {
        _id: id,
        embedding: v,
        norm: l2(v),
        fp,
        model: EMBED_MODEL,
      });
    });

    console.log(
      `[${name}] ${Math.min(i + 128, pending.length)}/${pending.length}`,
    );
  }

  await writeJsonlCompact(outPath, rows);
  console.log(`[${name}] wrote ${rows.size} rows.`);
};

const main = async (): Promise<void> => {
  try {
    await client.connect();
    const db = client.db();

    for (const c of COLLECTIONS) {
      console.log('Ingesting', c.name);
      await ingestCollection(db, c);
    }

    console.log('Ingest complete.');
  } catch (e) {
    console.error(e);
  } finally {
    await client.close().catch(() => {});
  }
};

main().catch((e) => {
  console.error(e);
  process.exit(1);
});