// ----------------------------------------------------------------------------
// In-memory document store: recursive chunking + cosine retrieval.
//
// For a single-instance deployment this is fine and fast. If you scale to
// multiple instances or want persistence, swap this Map for Redis / a vector DB
// (Qdrant, pgvector). The interface (put/get/retrieve) stays the same.
// ----------------------------------------------------------------------------

const docs = new Map(); // docId -> { chunks: string[], vectors: number[][], at: number }

// --- recursive character splitter (LangChain-style) ------------------------
export function recursiveSplit(
  text,
  chunkSize = 1000,
  chunkOverlap = 150,
  separators = ['\n\n', '\n', '. ', ' ', ''],
) {
  const splitRec = (input, seps) => {
    if (input.length <= chunkSize) return input.trim() ? [input] : [];
    const [sep, ...rest] = seps.length ? seps : [''];
    const parts = sep === '' ? input.split('') : input.split(sep);
    const chunks = [];
    let current = '';
    for (const part of parts) {
      const piece = sep === '' ? part : part + sep;
      if ((current + piece).length > chunkSize) {
        if (current) chunks.push(current);
        if (piece.length > chunkSize && rest.length) {
          chunks.push(...splitRec(piece, rest));
          current = '';
        } else {
          current = piece;
        }
      } else {
        current += piece;
      }
    }
    if (current) chunks.push(current);
    return chunks;
  };

  const raw = splitRec(text, separators).filter((c) => c.trim().length > 0);
  if (chunkOverlap <= 0) return raw;

  const out = [];
  for (let i = 0; i < raw.length; i++) {
    const prevTail = i > 0 ? raw[i - 1].slice(-chunkOverlap) : '';
    out.push((prevTail + raw[i]).trim());
  }
  return out;
}

export function cosine(a, b) {
  let dot = 0;
  let na = 0;
  let nb = 0;
  for (let i = 0; i < a.length; i++) {
    dot += a[i] * b[i];
    na += a[i] * a[i];
    nb += b[i] * b[i];
  }
  return dot / (Math.sqrt(na) * Math.sqrt(nb) + 1e-8);
}

export function hasDoc(docId) {
  return docs.has(docId);
}

export function putDoc(docId, chunks, vectors) {
  docs.set(docId, { chunks, vectors, at: Date.now() });
}

/** Retrieve top-k chunks across one or more docs for a query vector. */
export function retrieve(docIds, queryVec, k = 4) {
  const scored = [];
  for (const id of docIds) {
    const doc = docs.get(id);
    if (!doc) continue;
    for (let i = 0; i < doc.chunks.length; i++) {
      scored.push({ text: doc.chunks[i], score: cosine(queryVec, doc.vectors[i]) });
    }
  }
  return scored.sort((a, b) => b.score - a.score).slice(0, k).map((s) => s.text);
}

// Optional: evict docs older than N ms to cap memory (call periodically).
export function evictOlderThan(ms) {
  const cutoff = Date.now() - ms;
  for (const [id, doc] of docs) {
    if (doc.at < cutoff) docs.delete(id);
  }
}