/**
 * scripts/ingest-oracle-docs.js — one-time Oracle docs RAG ingestion.
 *
 * Chunks every PDF in data/oracle-docs/, embeds each chunk with OpenAI
 * (text-embedding-3-small, 1536 dims), and stores the vectors in Neon
 * (oracle_doc_chunks) for TuneBot retrieval.
 *
 * Idempotent: a doc already present in oracle_doc_chunks is skipped, so
 * re-runs only ingest new/failed docs. Run with:  npm run ingest-docs
 *
 * Requires OPENAI_API_KEY and DATABASE_URL in .env (loaded via dotenv BEFORE
 * db/index.js, which throws at import time when DATABASE_URL is unset).
 */

'use strict';

// dotenv MUST load before requiring db/index.js (it throws on missing DATABASE_URL).
require('dotenv').config();

const fs   = require('fs');
const path = require('path');
const pdf  = require('pdf-parse');
const OpenAI = require('openai');

const pool = require('../db/index'); // reuse the app's single Pool — no second pool

// ─── Config ─────────────────────────────────────────────────────────────────
const DOCS_DIR       = path.join(__dirname, '..', 'data', 'oracle-docs');
const EMBED_MODEL    = 'text-embedding-3-small';
const EMBED_DIMS     = 1536;
const TARGET_TOKENS  = 400;   // accumulate paragraphs up to ~this size
const HARD_TOKENS    = 600;   // a chunk must never exceed this
const EMBED_BATCH    = 100;   // chunks per OpenAI embeddings call
const CHARS_PER_TOKEN = 4;    // rough token estimate: Math.ceil(len / 4)

// ─── Helpers ────────────────────────────────────────────────────────────────

/** Rough token estimate — Math.ceil(chars / 4), per spec. */
function estimateTokens(text) {
  return Math.ceil(text.length / CHARS_PER_TOKEN);
}

/**
 * Chunk text by ~400 tokens: split on blank lines, accumulate paragraphs until
 * ~TARGET_TOKENS, and hard-split so no chunk exceeds HARD_TOKENS. A single
 * paragraph larger than HARD_TOKENS is sliced into ~TARGET_TOKENS pieces.
 */
function chunkText(text) {
  const paragraphs = text.split(/\n\n+/).map(p => p.trim()).filter(Boolean);
  const chunks = [];
  let current = '';

  const flush = () => {
    const trimmed = current.trim();
    if (trimmed) chunks.push(trimmed);
    current = '';
  };

  for (const para of paragraphs) {
    // Oversized single paragraph → flush what we have, then hard-slice it.
    if (estimateTokens(para) > HARD_TOKENS) {
      flush();
      const sliceChars = TARGET_TOKENS * CHARS_PER_TOKEN; // ~400 tokens per piece
      for (let i = 0; i < para.length; i += sliceChars) {
        const piece = para.slice(i, i + sliceChars).trim();
        if (piece) chunks.push(piece);
      }
      continue;
    }

    const candidate = current ? current + '\n\n' + para : para;
    const candidateTokens = estimateTokens(candidate);

    if (candidateTokens > HARD_TOKENS) {
      // Adding this paragraph would blow the hard cap → close the current chunk first.
      flush();
      current = para;
    } else if (candidateTokens >= TARGET_TOKENS) {
      // Reached the target size → include the paragraph and close the chunk.
      current = candidate;
      flush();
    } else {
      // Still under target → keep accumulating.
      current = candidate;
    }
  }
  flush();
  return chunks;
}

/** Embed a batch of chunk strings, retrying transient failures with backoff. */
async function embedBatch(openai, inputs, attempt = 1) {
  try {
    const res = await openai.embeddings.create({
      model: EMBED_MODEL,
      input: inputs,
      dimensions: EMBED_DIMS,
    });
    return res.data.map(d => d.embedding);
  } catch (err) {
    if (attempt < 4) {
      const waitMs = 1000 * Math.pow(2, attempt - 1); // 1s, 2s, 4s
      console.log(`  ! embed batch failed (${err.message}) — retry ${attempt}/3 in ${waitMs}ms`);
      await new Promise(r => setTimeout(r, waitMs));
      return embedBatch(openai, inputs, attempt + 1);
    }
    throw err;
  }
}

/** Postgres vector literal for a float array: [0.1,0.2,...]. */
function toVectorLiteral(embedding) {
  return '[' + embedding.join(',') + ']';
}

// ─── Step 1: schema ─────────────────────────────────────────────────────────
async function ensureSchema() {
  await pool.query(`
    CREATE TABLE IF NOT EXISTS oracle_doc_chunks (
      id BIGSERIAL PRIMARY KEY,
      doc_name TEXT NOT NULL,
      chunk_index INTEGER NOT NULL,
      content TEXT NOT NULL,
      token_count INTEGER,
      embedding vector(1536),
      created_at TIMESTAMPTZ DEFAULT NOW()
    )
  `);
  // UNIQUE(doc_name, chunk_index) as an idempotent unique index — works whether the
  // table is fresh or pre-dates this rule (a plain CREATE TABLE IF NOT EXISTS would
  // silently skip an inline constraint on an existing table). This is the arbiter
  // for the ON CONFLICT (doc_name, chunk_index) clause, giving true row-level dedup.
  await pool.query(`
    CREATE UNIQUE INDEX IF NOT EXISTS oracle_doc_chunks_doc_chunk_uidx
      ON oracle_doc_chunks (doc_name, chunk_index)
  `);
  await pool.query(`
    CREATE INDEX IF NOT EXISTS oracle_doc_chunks_embedding_idx
      ON oracle_doc_chunks USING ivfflat (embedding vector_cosine_ops)
      WITH (lists = 100)
  `);
}

async function docAlreadyIngested(docName) {
  const { rows } = await pool.query(
    'SELECT COUNT(*)::int AS n FROM oracle_doc_chunks WHERE doc_name = $1',
    [docName]
  );
  return rows[0].n > 0;
}

// ─── Step 2: ingest one PDF ─────────────────────────────────────────────────
async function ingestDoc(openai, docName) {
  if (await docAlreadyIngested(docName)) {
    console.log(`[${docName}] already ingested — skipping`);
    return { skipped: true };
  }

  const buffer = fs.readFileSync(path.join(DOCS_DIR, docName));
  const parsed = await pdf(buffer);
  // Strip NUL bytes (0x00): pdf-parse can emit them, and Postgres TEXT columns
  // reject them ("invalid byte sequence for encoding UTF8: 0x00"). They carry no
  // meaning in extracted prose, so removing them is safe.
  const text = (parsed.text || '').split(String.fromCharCode(0)).join('');
  const chunks = chunkText(text);

  if (chunks.length === 0) {
    console.log(`[${docName}] no extractable text — skipping`);
    return { skipped: true };
  }

  const total = chunks.length;
  console.log(`[${docName}] ${parsed.numpages} pages → ${total} chunks; embedding in batches of ${EMBED_BATCH}`);

  // All of a doc's rows are written in one transaction: if anything fails partway,
  // the whole doc rolls back and leaves zero rows — so the COUNT(*)>0 skip stays
  // accurate and a re-run cleanly re-ingests it (no half-finished docs).
  const client = await pool.connect();
  try {
    await client.query('BEGIN');
    for (let start = 0; start < total; start += EMBED_BATCH) {
      const batch = chunks.slice(start, start + EMBED_BATCH);
      const embeddings = await embedBatch(openai, batch);

      for (let j = 0; j < batch.length; j++) {
        const chunkIndex = start + j;          // 0-based stored index
        const content    = batch[j];
        const tokenCount = estimateTokens(content);
        await client.query(
          `INSERT INTO oracle_doc_chunks (doc_name, chunk_index, content, token_count, embedding)
           VALUES ($1, $2, $3, $4, $5::vector)
           ON CONFLICT (doc_name, chunk_index) DO NOTHING`,
          [docName, chunkIndex, content, tokenCount, toVectorLiteral(embeddings[j])]
        );
        console.log(`[${docName}] chunk ${chunkIndex + 1}/${total} | ~${tokenCount} tokens | embedded ✓`);
      }
    }
    await client.query('COMMIT');
  } catch (err) {
    await client.query('ROLLBACK').catch(() => {});
    throw err;
  } finally {
    client.release();
  }

  return { skipped: false, chunks: total };
}

// ─── Main ───────────────────────────────────────────────────────────────────
async function main() {
  if (!process.env.OPENAI_API_KEY) {
    console.error('OPENAI_API_KEY is not set — add it to .env before running.');
    process.exit(1);
  }
  if (!fs.existsSync(DOCS_DIR)) {
    console.error(`Docs directory not found: ${DOCS_DIR}`);
    process.exit(1);
  }

  const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });

  await ensureSchema();

  const pdfs = fs.readdirSync(DOCS_DIR).filter(f => f.endsWith('.pdf')).sort();
  if (pdfs.length === 0) {
    console.log('No PDFs found in data/oracle-docs/ — nothing to do.');
    return;
  }
  console.log(`Found ${pdfs.length} PDF(s) to process.\n`);

  let ingested = 0, skipped = 0, totalChunks = 0;
  for (const docName of pdfs) {
    try {
      const r = await ingestDoc(openai, docName);
      if (r.skipped) skipped++;
      else { ingested++; totalChunks += r.chunks; }
    } catch (err) {
      console.error(`[${docName}] FAILED: ${err.message} — leaving this doc for a re-run.`);
    }
    console.log('');
  }

  console.log(`Done. Ingested ${ingested} doc(s) (${totalChunks} chunks), skipped ${skipped}.`);
}

main()
  .catch(err => {
    console.error('Fatal:', err);
    process.exitCode = 1;
  })
  .finally(() => pool.end());
