/**
 * Copy the legacy banaocosting cluster into the sharded clusters.
 *
 *   node scripts/migrate-legacy.mjs llm_events    --dry-run
 *   node scripts/migrate-legacy.mjs llm_events
 *   node scripts/migrate-legacy.mjs lambda_events
 *
 * Design notes, because this run takes hours and will be interrupted:
 *
 * - READ-ONLY on the source. Nothing is deleted from banaocosting, so a failed
 *   or half-finished run costs nothing but time.
 * - RESUMABLE. Documents are walked in _id order and the last copied _id is
 *   checkpointed to disk, so re-running continues rather than restarting.
 * - IDEMPOTENT. A unique index on event_id is created on each target first, so
 *   re-inserting an already-copied document fails with a duplicate-key error
 *   that is counted and ignored. Interrupting mid-batch is therefore safe.
 * - SHARD-AWARE. Writes go to the first shard with room and roll to the next on
 *   a quota error, matching the app's fill-and-spill behaviour.
 * - THROTTLED. Atlas free clusters cap at 100 operations/second and respond to
 *   bursts by throttling the connection, so batches are paced rather than fired
 *   as fast as the network allows.
 */
import { MongoClient } from 'mongodb'
import { readFileSync, writeFileSync, existsSync } from 'node:fs'

// ---- config -----------------------------------------------------------------
const COLLECTION = process.argv[2]
const DRY = process.argv.includes('--dry-run')
const BATCH = Number(process.env.MIGRATE_BATCH ?? 500)
const OPS_PER_SEC = Number(process.env.MIGRATE_OPS_PER_SEC ?? 80) // under Atlas's 100
const QUOTA_BYTES = 512 * 1024 * 1024
const ROLLOVER_AT = 0.9

if (!['llm_events', 'lambda_events'].includes(COLLECTION)) {
  console.error('usage: node scripts/migrate-legacy.mjs <llm_events|lambda_events> [--dry-run]')
  process.exit(1)
}

const GROUP = COLLECTION === 'llm_events' ? 'CostingDBURL' : 'LambdaDBURL'
const CHECKPOINT = new URL(`../.migrate-${COLLECTION}.json`, import.meta.url)

// ---- env --------------------------------------------------------------------
const env = {}
for (const line of readFileSync(new URL('../.env', import.meta.url), 'utf8').split('\n')) {
  const t = line.trim()
  if (!t || t.startsWith('#')) continue
  const i = t.indexOf('=')
  if (i > 0) env[t.slice(0, i).trim()] = t.slice(i + 1).trim()
}
const DB_NAME = env.MONGO_DB_NAME ?? 'costingDB'
const LEGACY = process.env.LEGACY_DB_URL ?? env.LEGACY_DB_URL
if (!LEGACY) {
  console.error('LEGACY_DB_URL is not set (add it to backend/.env or pass it in the environment)')
  process.exit(1)
}

const targetUris = Object.keys(env)
  .map((k) => ({ k, n: k.match(new RegExp(`^${GROUP}(\\d+)$`, 'i'))?.[1] }))
  .filter((x) => x.n !== undefined)
  .sort((a, b) => Number(a.n) - Number(b.n))
  .map((x) => ({ name: x.k, uri: env[x.k] }))

if (!targetUris.length) {
  console.error(`no ${GROUP}* entries found in backend/.env`)
  process.exit(1)
}

const sleep = (ms) => new Promise((r) => setTimeout(r, ms))
const isQuota = (e) => /over your space quota|quota exceeded|you are over/i.test(String(e))

// ---- connect ----------------------------------------------------------------
console.log(`migrating ${COLLECTION}  ->  ${targetUris.map((t) => t.name).join(', ')}${DRY ? '  [DRY RUN]' : ''}`)

const source = new MongoClient(LEGACY, { maxPoolSize: 4 })
await source.connect()
const src = source.db(DB_NAME).collection(COLLECTION)

const targets = []
for (const t of targetUris) {
  const client = new MongoClient(t.uri, { maxPoolSize: 4 })
  await client.connect()
  targets.push({ ...t, client, col: client.db(DB_NAME).collection(COLLECTION) })
}

async function usedBytes(t) {
  const s = await t.client.db(DB_NAME).command({ dbStats: 1, scale: 1 })
  return (s.storageSize ?? 0) + (s.indexSize ?? 0)
}

// Unique event_id is what makes re-runs safe: duplicates are rejected by the
// index rather than silently doubling the data.
if (!DRY) {
  for (const t of targets) {
    await t.col.createIndex({ event_id: 1 }, { unique: true })
    await t.col.createIndex({ created_at: -1 })
  }
  console.log('ensured event_id + created_at indexes on every target')
}

// ---- resume -----------------------------------------------------------------
let checkpoint = null
if (existsSync(CHECKPOINT)) {
  checkpoint = JSON.parse(readFileSync(CHECKPOINT, 'utf8'))
  console.log(`resuming after _id ${checkpoint.lastId} (${checkpoint.copied} already copied)`)
}

const total = await src.estimatedDocumentCount()
const query = checkpoint ? { _id: { $gt: new (await import('mongodb')).ObjectId(checkpoint.lastId) } } : {}
const cursor = src.find(query).sort({ _id: 1 }).batchSize(BATCH)

let copied = checkpoint?.copied ?? 0
let dupes = 0
let ti = 0
let batch = []
const started = Date.now()

async function pickTarget() {
  while (ti < targets.length) {
    const used = await usedBytes(targets[ti])
    if (used < QUOTA_BYTES * ROLLOVER_AT) return targets[ti]
    console.log(`  ${targets[ti].name} is full (${(used / 1e6).toFixed(0)}MB) — rolling over`)
    ti++
  }
  throw new Error('ALL TARGET SHARDS FULL — add more DB URLs to backend/.env and re-run; progress is checkpointed')
}

async function flush() {
  if (!batch.length) return
  if (DRY) { copied += batch.length; batch = []; return }

  const target = await pickTarget()
  try {
    await target.col.insertMany(batch, { ordered: false })
    copied += batch.length
  } catch (err) {
    if (isQuota(err)) {
      ti++ // this shard filled mid-batch; retry the same batch on the next one
      return flush()
    }
    // Duplicate keys mean those documents were already migrated - expected on
    // resume, and the whole reason the run is safe to interrupt.
    const dupCount = err?.writeErrors?.filter((w) => w.err?.code === 11000).length ?? 0
    const other = (err?.writeErrors?.length ?? 0) - dupCount
    if (other > 0) throw err
    dupes += dupCount
    copied += batch.length - dupCount
  }

  writeFileSync(CHECKPOINT, JSON.stringify({ lastId: String(batch.at(-1)._id), copied }))
  batch = []

  const pct = ((copied / total) * 100).toFixed(1)
  const rate = copied / ((Date.now() - started) / 1000)
  const eta = rate > 0 ? Math.round((total - copied) / rate / 60) : 0
  process.stdout.write(`\r  ${copied.toLocaleString()}/${total.toLocaleString()} (${pct}%)  ${rate.toFixed(0)}/s  eta ${eta}m  dupes ${dupes}   `)

  await sleep((batch.length || BATCH) / OPS_PER_SEC * 1000)
}

for await (const doc of cursor) {
  batch.push(doc)
  if (batch.length >= BATCH) await flush()
}
await flush()

console.log(`\ndone: ${copied.toLocaleString()} copied, ${dupes.toLocaleString()} already present`)
console.log('source cluster was not modified')

await source.close()
await Promise.all(targets.map((t) => t.client.close()))