#!/usr/bin/env tsx
/**
 * BEAM 1M — resumable per-conversation substrate ingest (mind-per-conv).
 *
 * DESIGN (per build-phase brief):
 *   - One persistent .mind file PER conversation → no cross-conversation lock
 *     contention, so N conversations can ingest concurrently.
 *   - RESUMABLE: a `<mind>.done.json` marker is written only after a
 *     conversation fully ingests + indexes. On restart, conversations with a
 *     valid marker are skipped; a .mind file WITHOUT a marker (crash mid-ingest)
 *     is deleted and rebuilt from scratch. This mirrors LongMemEval's
 *     skip-if-complete minds cache.
 *   - Reads turns directly from the BEAM repo's per-conversation chat.json
 *     (~4 MB each), NOT the 3 GB canonical — `extractTurnsFromBeam` does a
 *     readFileSync that would blow V8's string limit on the 1M archive.
 *
 * The gop_id written per frame is `beam_<convId>`, matching the
 * `conversation_id` on every canonical instance — so the answer cells scope to
 * the right conversation without any change to their gopId filter.
 *
 * Embedding: local Ollama `nomic-embed-text` (1024-dim) — free, no API spend.
 * The embedder has a fixed 30 s per-request timeout, so the index batch size is
 * kept modest (default 48) to avoid timing out on long turns.
 *
 * Progress: one stdout line per conversation start/finish, plus a marker file
 * per completed conversation and an appended progress log — all monitorable
 * from the filesystem while this runs in the background.
 *
 * Usage:
 *   tsx benchmarks/harness/scripts/beam-ingest-1m.ts \
 *     [--convs 1-35] [--concurrency 1] [--batch-size 48] \
 *     [--beam-chats D:/Projects/BEAM/chats] [--minds-dir <path>] [--force]
 */

import fs from 'node:fs';
import path from 'node:path';
import url from 'node:url';
import process from 'node:process';

import { createOllamaEmbedder } from '@waggle/core';
import { createSubstrate } from '../src/substrate.js';
import { ingestBeamCorpus } from '../src/ingest-beam.js';
import type { BeamTurn } from '../src/ingest-beam.js';

const CHAT_SIZE = '1M';
const CHAT_SIZE_DIR = '1M';

interface Args {
  convs: number[];
  concurrency: number;
  batchSize: number;
  beamChats: string;
  mindsDir: string;
  force: boolean;
}

function parseConvSpec(spec: string): number[] {
  const out = new Set<number>();
  for (const part of spec.split(',')) {
    const m = part.match(/^(\d+)-(\d+)$/);
    if (m) {
      const a = parseInt(m[1], 10);
      const b = parseInt(m[2], 10);
      for (let i = a; i <= b; i++) out.add(i);
    } else if (/^\d+$/.test(part.trim())) {
      out.add(parseInt(part.trim(), 10));
    }
  }
  return [...out].sort((a, b) => a - b);
}

function parseArgs(): Args {
  const argv = process.argv.slice(2);
  const here = url.fileURLToPath(import.meta.url);
  const repoRoot = path.resolve(path.dirname(here), '..', '..', '..');
  let convs = parseConvSpec('1-35');
  let concurrency = 1;
  let batchSize = 48;
  let beamChats = path.resolve(repoRoot, '..', 'BEAM', 'chats');
  let mindsDir = path.join(repoRoot, 'benchmarks', 'data', 'beam', 'minds-1M');
  let force = false;
  for (let i = 0; i < argv.length; i++) {
    const f = argv[i];
    const next = argv[i + 1];
    if (f === '--convs' && next) { convs = parseConvSpec(next); i++; }
    else if (f === '--concurrency' && next) { concurrency = Math.max(1, parseInt(next, 10)); i++; }
    else if (f === '--batch-size' && next) { batchSize = Math.max(1, parseInt(next, 10)); i++; }
    else if (f === '--beam-chats' && next) { beamChats = path.resolve(next); i++; }
    else if (f === '--minds-dir' && next) { mindsDir = path.resolve(next); i++; }
    else if (f === '--force') { force = true; }
  }
  return { convs, concurrency, batchSize, beamChats, mindsDir, force };
}

interface RawMsg { role?: string; content?: string }
interface RawBatch { turns?: RawMsg[][] }

/**
 * Flatten a BEAM chat.json into ordered {role, content} messages.
 * Replicates build-beam-canonical.ts::flattenBeamTurns EXACTLY (batches →
 * turn-groups → messages, non-empty content only) so the ingested frames match
 * the canonical's conversation turns.
 */
function flattenChatJson(chatJsonPath: string): Array<{ role: string; content: string }> {
  const batches = JSON.parse(fs.readFileSync(chatJsonPath, 'utf-8')) as RawBatch[];
  const out: Array<{ role: string; content: string }> = [];
  for (const batch of batches) {
    if (!Array.isArray(batch.turns)) continue;
    for (const group of batch.turns) {
      if (!Array.isArray(group)) continue;
      for (const msg of group) {
        const role = String(msg.role ?? 'unknown').toLowerCase();
        const content = String(msg.content ?? '').trim();
        if (content) out.push({ role, content });
      }
    }
  }
  return out;
}

interface IngestMarker {
  conv: number;
  gop_id: string;
  chat_size: string;
  frames: number;
  turns_seen: number;
  ingest_ms: number;
  index_ms: number;
  batch_size: number;
  embedder_model: string;
  embedder_dims: number;
  built_at: string;
}

function markerPath(mindsDir: string, conv: number): string {
  return path.join(mindsDir, `beam_1M_${conv}.done.json`);
}
function mindPath(mindsDir: string, conv: number): string {
  return path.join(mindsDir, `beam_1M_${conv}.mind`);
}

function isComplete(mindsDir: string, conv: number): IngestMarker | null {
  const mp = markerPath(mindsDir, conv);
  if (!fs.existsSync(mp)) return null;
  try {
    const m = JSON.parse(fs.readFileSync(mp, 'utf-8')) as IngestMarker;
    if (m && m.frames > 0 && fs.existsSync(mindPath(mindsDir, conv))) return m;
  } catch { /* fall through */ }
  return null;
}

function logProgress(mindsDir: string, line: string): void {
  const stamped = `${new Date().toISOString()} ${line}`;
  process.stdout.write(stamped + '\n');
  try {
    fs.appendFileSync(path.join(mindsDir, '_ingest-progress.log'), stamped + '\n');
  } catch { /* best-effort */ }
}

async function ingestOne(args: Args, conv: number): Promise<IngestMarker> {
  const gopId = `beam_${conv}`;
  const chatJsonPath = path.join(args.beamChats, CHAT_SIZE_DIR, String(conv), 'chat.json');
  if (!fs.existsSync(chatJsonPath)) {
    throw new Error(`chat.json not found for conv ${conv}: ${chatJsonPath}`);
  }

  // Clean any partial .mind left by a prior crash (no valid marker present).
  const mp = mindPath(args.mindsDir, conv);
  for (const suffix of ['', '-wal', '-shm']) {
    const f = mp + suffix;
    if (fs.existsSync(f)) fs.rmSync(f, { force: true });
  }

  const messages = flattenChatJson(chatJsonPath);
  const turns: BeamTurn[] = messages.map((m, i) => ({
    gopId,
    messageIndex: i,
    role: m.role === 'assistant' ? 'assistant' : 'user',
    content: m.content,
    formattedContent: `${m.role}: ${m.content}`,
    chatSize: CHAT_SIZE,
  }));

  logProgress(args.mindsDir, `[ingest][conv ${conv}] START turns=${turns.length} db=${path.basename(mp)}`);

  const embedder = createOllamaEmbedder();
  const substrate = createSubstrate({ dbPath: mp, embedder });
  try {
    const stats = await ingestBeamCorpus(
      substrate.db, substrate.search, substrate.frames, substrate.sessions,
      turns, { batchSize: args.batchSize },
    );
    const marker: IngestMarker = {
      conv,
      gop_id: gopId,
      chat_size: CHAT_SIZE,
      frames: stats.count,
      turns_seen: turns.length,
      ingest_ms: stats.ingestMs,
      index_ms: stats.indexMs,
      batch_size: args.batchSize,
      embedder_model: 'nomic-embed-text',
      embedder_dims: embedder.dimensions,
      built_at: new Date().toISOString(),
    };
    fs.writeFileSync(markerPath(args.mindsDir, conv), JSON.stringify(marker, null, 2) + '\n', 'utf-8');
    return marker;
  } finally {
    substrate.close();
  }
}

async function runPool(args: Args, convs: number[]): Promise<void> {
  let cursor = 0;
  let done = 0;
  const total = convs.length;
  const startAll = Date.now();

  async function worker(): Promise<void> {
    while (true) {
      const idx = cursor++;
      if (idx >= convs.length) return;
      const conv = convs[idx];
      const existing = args.force ? null : isComplete(args.mindsDir, conv);
      if (existing) {
        done++;
        logProgress(args.mindsDir, `[ingest][conv ${conv}] SKIP (already complete, frames=${existing.frames}) [${done}/${total}]`);
        continue;
      }
      const t0 = Date.now();
      try {
        const m = await ingestOne(args, conv);
        done++;
        const secs = ((Date.now() - t0) / 1000).toFixed(1);
        const elapsedMin = ((Date.now() - startAll) / 60000).toFixed(1);
        const rate = done / ((Date.now() - startAll) / 60000);
        const etaMin = rate > 0 ? ((total - done) / rate).toFixed(1) : '?';
        logProgress(
          args.mindsDir,
          `[ingest][conv ${conv}] DONE frames=${m.frames} index_ms=${m.index_ms} took=${secs}s ` +
          `[${done}/${total}] elapsed=${elapsedMin}m eta=${etaMin}m`,
        );
      } catch (err) {
        logProgress(args.mindsDir, `[ingest][conv ${conv}] ERROR ${(err as Error).message}`);
        throw err;
      }
    }
  }

  const workers = Array.from({ length: Math.min(args.concurrency, convs.length) }, () => worker());
  await Promise.all(workers);
  const totalMin = ((Date.now() - startAll) / 60000).toFixed(1);
  logProgress(args.mindsDir, `[ingest] ALL DONE ${done}/${total} conversations in ${totalMin}m`);
}

async function main(): Promise<void> {
  const args = parseArgs();
  fs.mkdirSync(args.mindsDir, { recursive: true });
  logProgress(
    args.mindsDir,
    `[ingest] start convs=${args.convs[0]}..${args.convs[args.convs.length - 1]} (n=${args.convs.length}) ` +
    `concurrency=${args.concurrency} batch=${args.batchSize} chats=${args.beamChats}`,
  );
  await runPool(args, args.convs);
}

main().catch(err => {
  console.error('[beam-ingest-1m] FATAL:', err);
  process.exit(1);
});