// Local JSON vector store over the wiki pages + hybrid (cosine + BM25) search.
// Incremental: only re-embeds pages whose file changed since the last build.

import * as fs from "node:fs/promises";
import * as path from "node:path";
import * as crypto from "node:crypto";
import type { EmbeddingDynamicHandle } from "@lmstudio/sdk";
import {
  type WikiLayout,
  listIndexablePages,
  parsePage,
  readText,
} from "./wikiStore";
import { chunkMarkdown, embeddingTextFor, type ChunkOptions } from "./chunk";
import { embedMany } from "./embeddings";

// v2: staleness is keyed on a content hash instead of the file mtime, so the
// index is portable across machines (git checkout / cloud sync rewrite mtimes
// but not content). Bumping the version forces one clean rebuild on upgrade.
const INDEX_VERSION = 2;

/** Content hash used for change detection. Normalised to be insensitive to
 *  CRLF/LF line-ending differences (git autocrlf) so the same page hashes
 *  identically on every machine. */
function hashContent(content: string): string {
  return crypto.createHash("sha1").update(content.replace(/\r\n/g, "\n")).digest("hex");
}

export interface IndexedChunk {
  id: string; // `${relPath}#${i}`
  relPath: string; // relative to wiki/
  pageTitle: string;
  wikilink: string; // relPath without .md, forward slashes
  headingPath: string[];
  text: string;
  vector: number[];
}

interface FileMeta {
  hash: string; // content hash (see hashContent) — portable across machines
  size: number;
  chunkIds: string[];
}

export interface RagIndex {
  version: number;
  model: string;
  dim: number;
  createdAt: string;
  updatedAt: string;
  files: Record<string, FileMeta>;
  chunks: Record<string, IndexedChunk>;
}

function emptyIndex(model: string): RagIndex {
  const now = new Date().toISOString();
  return {
    version: INDEX_VERSION,
    model,
    dim: 0,
    createdAt: now,
    updatedAt: now,
    files: {},
    chunks: {},
  };
}

export async function loadIndex(layout: WikiLayout): Promise<RagIndex | null> {
  try {
    const raw = await readText(layout.ragIndexPath);
    const parsed = JSON.parse(raw) as RagIndex;
    if (parsed.version !== INDEX_VERSION) return null;
    return parsed;
  } catch {
    return null;
  }
}

export async function saveIndex(layout: WikiLayout, index: RagIndex): Promise<void> {
  await fs.mkdir(layout.ragDir, { recursive: true });
  await fs.writeFile(layout.ragIndexPath, JSON.stringify(index), "utf-8");
}

function pageTitleFor(relPath: string, data: Record<string, unknown>): string {
  const t = (data.title ?? data.name ?? data.concept) as string | undefined;
  if (t && typeof t === "string" && t.trim()) return t.trim();
  const base = path.basename(relPath, ".md");
  return base.replace(/-/g, " ");
}

function wikilinkFor(relPath: string): string {
  return relPath.replace(/\\/g, "/").replace(/\.md$/i, "");
}

export interface ReindexStats {
  pagesTotal: number;
  pagesReembedded: number;
  pagesUnchanged: number;
  pagesRemoved: number;
  chunksTotal: number;
  model: string;
  fullRebuild: boolean;
}

export interface ReindexOptions extends ChunkOptions {
  onProgress?: (done: number, total: number, label: string) => void;
  abortSignal?: AbortSignal;
}

/**
 * Rebuild the index incrementally against the current wiki pages.
 * `modelIdentifier` must match the embedder; if it differs from the stored
 * index model, a full rebuild is performed (vectors are not comparable).
 */
export async function reindex(
  layout: WikiLayout,
  embedder: EmbeddingDynamicHandle,
  modelIdentifier: string,
  opts: ReindexOptions,
): Promise<ReindexStats> {
  const existing = await loadIndex(layout);
  const fullRebuild = !existing || existing.model !== modelIdentifier;
  const index: RagIndex = fullRebuild ? emptyIndex(modelIdentifier) : existing!;

  const pageAbsPaths = await listIndexablePages(layout);
  const liveRelPaths = new Set<string>();

  let pagesReembedded = 0;
  let pagesUnchanged = 0;

  for (const abs of pageAbsPaths) {
    if (opts.abortSignal?.aborted) throw new Error("Reindex aborted.");
    const relPath = path.relative(layout.wikiDir, abs).replace(/\\/g, "/");
    liveRelPaths.add(relPath);

    const stat = await fs.stat(abs);
    const content = await readText(abs);
    const hash = hashContent(content);
    const prior = index.files[relPath];
    if (prior && prior.hash === hash && prior.size === stat.size) {
      pagesUnchanged++;
      continue;
    }

    // (Re)build this page's chunks.
    const { data, body } = parsePage(content);
    const pageTitle = pageTitleFor(relPath, data);
    const wikilink = wikilinkFor(relPath);
    const chunks = chunkMarkdown(body, {
      chunkSizeChars: opts.chunkSizeChars,
      chunkOverlapChars: opts.chunkOverlapChars,
    });

    // Drop the page's old chunks before inserting fresh ones.
    if (prior) for (const id of prior.chunkIds) delete index.chunks[id];

    const texts = chunks.map((c) => embeddingTextFor(pageTitle, c));
    opts.onProgress?.(pagesReembedded, pageAbsPaths.length, `Embedding ${relPath} (${texts.length} chunks)`);
    const vectors = texts.length
      ? await embedMany(embedder, texts, { abortSignal: opts.abortSignal })
      : [];

    const chunkIds: string[] = [];
    for (let i = 0; i < chunks.length; i++) {
      const id = `${relPath}#${i}`;
      chunkIds.push(id);
      index.chunks[id] = {
        id,
        relPath,
        pageTitle,
        wikilink,
        headingPath: chunks[i].headingPath,
        text: chunks[i].text,
        vector: vectors[i],
      };
      if (vectors[i]?.length) index.dim = vectors[i].length;
    }
    index.files[relPath] = { hash, size: stat.size, chunkIds };
    pagesReembedded++;
  }

  // Remove files (and their chunks) that no longer exist.
  let pagesRemoved = 0;
  for (const relPath of Object.keys(index.files)) {
    if (!liveRelPaths.has(relPath)) {
      for (const id of index.files[relPath].chunkIds) delete index.chunks[id];
      delete index.files[relPath];
      pagesRemoved++;
    }
  }

  index.updatedAt = new Date().toISOString();
  await saveIndex(layout, index);

  return {
    pagesTotal: pageAbsPaths.length,
    pagesReembedded,
    pagesUnchanged,
    pagesRemoved,
    chunksTotal: Object.keys(index.chunks).length,
    model: modelIdentifier,
    fullRebuild,
  };
}

/** Relative paths of pages whose file changed since the index was built. */
export async function staleRelPaths(layout: WikiLayout, index: RagIndex | null): Promise<string[]> {
  if (!index) return ["(no index)"];
  const pageAbsPaths = await listIndexablePages(layout);
  const stale: string[] = [];
  const live = new Set<string>();
  for (const abs of pageAbsPaths) {
    const relPath = path.relative(layout.wikiDir, abs).replace(/\\/g, "/");
    live.add(relPath);
    const prior = index.files[relPath];
    try {
      const stat = await fs.stat(abs);
      // Size is a cheap pre-check; the content hash is the source of truth.
      if (!prior || prior.size !== stat.size || prior.hash !== hashContent(await readText(abs))) {
        stale.push(relPath);
      }
    } catch {
      /* ignore */
    }
  }
  for (const relPath of Object.keys(index.files)) {
    if (!live.has(relPath)) stale.push(`${relPath} (deleted)`);
  }
  return stale;
}

// ---------------------------------------------------------------------------
// Search
// ---------------------------------------------------------------------------

export interface SearchHit {
  chunk: IndexedChunk;
  score: number;
  semantic: number;
  lexical: number;
  snippet: string;
}

function cosine(a: number[], b: number[]): number {
  if (!a?.length || !b?.length || a.length !== b.length) return 0;
  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];
  }
  if (na === 0 || nb === 0) return 0;
  return dot / (Math.sqrt(na) * Math.sqrt(nb));
}

function tokenize(text: string): string[] {
  return (text.toLowerCase().match(/[\p{L}\p{N}]+/gu) ?? []).filter((t) => t.length > 1);
}

/** Lightweight BM25 over the chunk corpus, computed at query time. */
function bm25Scores(chunks: IndexedChunk[], queryTerms: string[]): Map<string, number> {
  const k1 = 1.5;
  const b = 0.75;
  const N = chunks.length;
  const docTokens = chunks.map((c) => tokenize(`${c.pageTitle} ${c.headingPath.join(" ")} ${c.text}`));
  const docLen = docTokens.map((d) => d.length);
  const avgdl = docLen.reduce((s, l) => s + l, 0) / Math.max(1, N);

  const df = new Map<string, number>();
  for (const term of new Set(queryTerms)) {
    let count = 0;
    for (const d of docTokens) if (d.includes(term)) count++;
    df.set(term, count);
  }

  const scores = new Map<string, number>();
  for (let i = 0; i < chunks.length; i++) {
    const tokens = docTokens[i];
    let score = 0;
    for (const term of queryTerms) {
      const n = df.get(term) ?? 0;
      if (n === 0) continue;
      const idf = Math.log(1 + (N - n + 0.5) / (n + 0.5));
      const tf = tokens.filter((t) => t === term).length;
      if (tf === 0) continue;
      const denom = tf + k1 * (1 - b + (b * docLen[i]) / avgdl);
      score += idf * ((tf * (k1 + 1)) / denom);
    }
    if (score > 0) scores.set(chunks[i].id, score);
  }
  return scores;
}

function normalize(map: Map<string, number>): Map<string, number> {
  let max = 0;
  for (const v of map.values()) if (v > max) max = v;
  if (max === 0) return map;
  const out = new Map<string, number>();
  for (const [k, v] of map) out.set(k, v / max);
  return out;
}

function makeSnippet(text: string, max = 320): string {
  const clean = text.replace(/\s+/g, " ").trim();
  return clean.length <= max ? clean : clean.slice(0, max - 1) + "…";
}

export interface SearchOptions {
  topK: number;
  hybrid: boolean;
}

export function search(
  index: RagIndex,
  queryVector: number[],
  queryText: string,
  opts: SearchOptions,
): SearchHit[] {
  const chunks = Object.values(index.chunks);
  if (chunks.length === 0) return [];

  const semantic = new Map<string, number>();
  for (const c of chunks) semantic.set(c.id, cosine(queryVector, c.vector));

  let combined: Map<string, number>;
  let lexicalNorm = new Map<string, number>();
  if (opts.hybrid) {
    const lexical = bm25Scores(chunks, tokenize(queryText));
    lexicalNorm = normalize(lexical);
    const semanticNorm = normalize(semantic);
    combined = new Map<string, number>();
    for (const c of chunks) {
      const s = semanticNorm.get(c.id) ?? 0;
      const l = lexicalNorm.get(c.id) ?? 0;
      combined.set(c.id, 0.6 * s + 0.4 * l);
    }
  } else {
    combined = semantic;
  }

  return chunks
    .map((c) => ({
      chunk: c,
      score: combined.get(c.id) ?? 0,
      semantic: semantic.get(c.id) ?? 0,
      lexical: lexicalNorm.get(c.id) ?? 0,
      snippet: makeSnippet(c.text),
    }))
    .sort((a, b) => b.score - a.score)
    .slice(0, Math.max(1, opts.topK));
}

// Local JSON vector store over the wiki pages + hybrid (cosine + BM25) search.
// Incremental: only re-embeds pages whose file changed since the last build.

import * as fs from "node:fs/promises";
import * as path from "node:path";
import * as crypto from "node:crypto";
import type { EmbeddingDynamicHandle } from "@lmstudio/sdk";
import {
  type WikiLayout,
  listIndexablePages,
  parsePage,
  readText,
} from "./wikiStore";
import { chunkMarkdown, embeddingTextFor, type ChunkOptions } from "./chunk";
import { embedMany } from "./embeddings";

// v2: staleness is keyed on a content hash instead of the file mtime, so the
// index is portable across machines (git checkout / cloud sync rewrite mtimes
// but not content). Bumping the version forces one clean rebuild on upgrade.
const INDEX_VERSION = 2;

/** Content hash used for change detection. Normalised to be insensitive to
 *  CRLF/LF line-ending differences (git autocrlf) so the same page hashes
 *  identically on every machine. */
function hashContent(content: string): string {
  return crypto.createHash("sha1").update(content.replace(/\r\n/g, "\n")).digest("hex");
}

export interface IndexedChunk {
  id: string; // `${relPath}#${i}`
  relPath: string; // relative to wiki/
  pageTitle: string;
  wikilink: string; // relPath without .md, forward slashes
  headingPath: string[];
  text: string;
  vector: number[];
}

interface FileMeta {
  hash: string; // content hash (see hashContent) — portable across machines
  size: number;
  chunkIds: string[];
}

export interface RagIndex {
  version: number;
  model: string;
  dim: number;
  createdAt: string;
  updatedAt: string;
  files: Record<string, FileMeta>;
  chunks: Record<string, IndexedChunk>;
}

function emptyIndex(model: string): RagIndex {
  const now = new Date().toISOString();
  return {
    version: INDEX_VERSION,
    model,
    dim: 0,
    createdAt: now,
    updatedAt: now,
    files: {},
    chunks: {},
  };
}

export async function loadIndex(layout: WikiLayout): Promise<RagIndex | null> {
  try {
    const raw = await readText(layout.ragIndexPath);
    const parsed = JSON.parse(raw) as RagIndex;
    if (parsed.version !== INDEX_VERSION) return null;
    return parsed;
  } catch {
    return null;
  }
}

export async function saveIndex(layout: WikiLayout, index: RagIndex): Promise<void> {
  await fs.mkdir(layout.ragDir, { recursive: true });
  await fs.writeFile(layout.ragIndexPath, JSON.stringify(index), "utf-8");
}

function pageTitleFor(relPath: string, data: Record<string, unknown>): string {
  const t = (data.title ?? data.name ?? data.concept) as string | undefined;
  if (t && typeof t === "string" && t.trim()) return t.trim();
  const base = path.basename(relPath, ".md");
  return base.replace(/-/g, " ");
}

function wikilinkFor(relPath: string): string {
  return relPath.replace(/\\/g, "/").replace(/\.md$/i, "");
}

export interface ReindexStats {
  pagesTotal: number;
  pagesReembedded: number;
  pagesUnchanged: number;
  pagesRemoved: number;
  chunksTotal: number;
  model: string;
  fullRebuild: boolean;
}

export interface ReindexOptions extends ChunkOptions {
  onProgress?: (done: number, total: number, label: string) => void;
  abortSignal?: AbortSignal;
}

/**
 * Rebuild the index incrementally against the current wiki pages.
 * `modelIdentifier` must match the embedder; if it differs from the stored
 * index model, a full rebuild is performed (vectors are not comparable).
 */
export async function reindex(
  layout: WikiLayout,
  embedder: EmbeddingDynamicHandle,
  modelIdentifier: string,
  opts: ReindexOptions,
): Promise<ReindexStats> {
  const existing = await loadIndex(layout);
  const fullRebuild = !existing || existing.model !== modelIdentifier;
  const index: RagIndex = fullRebuild ? emptyIndex(modelIdentifier) : existing!;

  const pageAbsPaths = await listIndexablePages(layout);
  const liveRelPaths = new Set<string>();

  let pagesReembedded = 0;
  let pagesUnchanged = 0;

  for (const abs of pageAbsPaths) {
    if (opts.abortSignal?.aborted) throw new Error("Reindex aborted.");
    const relPath = path.relative(layout.wikiDir, abs).replace(/\\/g, "/");
    liveRelPaths.add(relPath);

    const stat = await fs.stat(abs);
    const content = await readText(abs);
    const hash = hashContent(content);
    const prior = index.files[relPath];
    if (prior && prior.hash === hash && prior.size === stat.size) {
      pagesUnchanged++;
      continue;
    }

    // (Re)build this page's chunks.
    const { data, body } = parsePage(content);
    const pageTitle = pageTitleFor(relPath, data);
    const wikilink = wikilinkFor(relPath);
    const chunks = chunkMarkdown(body, {
      chunkSizeChars: opts.chunkSizeChars,
      chunkOverlapChars: opts.chunkOverlapChars,
    });

    // Drop the page's old chunks before inserting fresh ones.
    if (prior) for (const id of prior.chunkIds) delete index.chunks[id];

    const texts = chunks.map((c) => embeddingTextFor(pageTitle, c));
    opts.onProgress?.(pagesReembedded, pageAbsPaths.length, `Embedding ${relPath} (${texts.length} chunks)`);
    const vectors = texts.length
      ? await embedMany(embedder, texts, { abortSignal: opts.abortSignal })
      : [];

    const chunkIds: string[] = [];
    for (let i = 0; i < chunks.length; i++) {
      const id = `${relPath}#${i}`;
      chunkIds.push(id);
      index.chunks[id] = {
        id,
        relPath,
        pageTitle,
        wikilink,
        headingPath: chunks[i].headingPath,
        text: chunks[i].text,
        vector: vectors[i],
      };
      if (vectors[i]?.length) index.dim = vectors[i].length;
    }
    index.files[relPath] = { hash, size: stat.size, chunkIds };
    pagesReembedded++;
  }

  // Remove files (and their chunks) that no longer exist.
  let pagesRemoved = 0;
  for (const relPath of Object.keys(index.files)) {
    if (!liveRelPaths.has(relPath)) {
      for (const id of index.files[relPath].chunkIds) delete index.chunks[id];
      delete index.files[relPath];
      pagesRemoved++;
    }
  }

  index.updatedAt = new Date().toISOString();
  await saveIndex(layout, index);

  return {
    pagesTotal: pageAbsPaths.length,
    pagesReembedded,
    pagesUnchanged,
    pagesRemoved,
    chunksTotal: Object.keys(index.chunks).length,
    model: modelIdentifier,
    fullRebuild,
  };
}

/** Relative paths of pages whose file changed since the index was built. */
export async function staleRelPaths(layout: WikiLayout, index: RagIndex | null): Promise<string[]> {
  if (!index) return ["(no index)"];
  const pageAbsPaths = await listIndexablePages(layout);
  const stale: string[] = [];
  const live = new Set<string>();
  for (const abs of pageAbsPaths) {
    const relPath = path.relative(layout.wikiDir, abs).replace(/\\/g, "/");
    live.add(relPath);
    const prior = index.files[relPath];
    try {
      const stat = await fs.stat(abs);
      // Size is a cheap pre-check; the content hash is the source of truth.
      if (!prior || prior.size !== stat.size || prior.hash !== hashContent(await readText(abs))) {
        stale.push(relPath);
      }
    } catch {
      /* ignore */
    }
  }
  for (const relPath of Object.keys(index.files)) {
    if (!live.has(relPath)) stale.push(`${relPath} (deleted)`);
  }
  return stale;
}

// ---------------------------------------------------------------------------
// Search
// ---------------------------------------------------------------------------

export interface SearchHit {
  chunk: IndexedChunk;
  score: number;
  semantic: number;
  lexical: number;
  snippet: string;
}

function cosine(a: number[], b: number[]): number {
  if (!a?.length || !b?.length || a.length !== b.length) return 0;
  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];
  }
  if (na === 0 || nb === 0) return 0;
  return dot / (Math.sqrt(na) * Math.sqrt(nb));
}

function tokenize(text: string): string[] {
  return (text.toLowerCase().match(/[\p{L}\p{N}]+/gu) ?? []).filter((t) => t.length > 1);
}

/** Lightweight BM25 over the chunk corpus, computed at query time. */
function bm25Scores(chunks: IndexedChunk[], queryTerms: string[]): Map<string, number> {
  const k1 = 1.5;
  const b = 0.75;
  const N = chunks.length;
  const docTokens = chunks.map((c) => tokenize(`${c.pageTitle} ${c.headingPath.join(" ")} ${c.text}`));
  const docLen = docTokens.map((d) => d.length);
  const avgdl = docLen.reduce((s, l) => s + l, 0) / Math.max(1, N);

  const df = new Map<string, number>();
  for (const term of new Set(queryTerms)) {
    let count = 0;
    for (const d of docTokens) if (d.includes(term)) count++;
    df.set(term, count);
  }

  const scores = new Map<string, number>();
  for (let i = 0; i < chunks.length; i++) {
    const tokens = docTokens[i];
    let score = 0;
    for (const term of queryTerms) {
      const n = df.get(term) ?? 0;
      if (n === 0) continue;
      const idf = Math.log(1 + (N - n + 0.5) / (n + 0.5));
      const tf = tokens.filter((t) => t === term).length;
      if (tf === 0) continue;
      const denom = tf + k1 * (1 - b + (b * docLen[i]) / avgdl);
      score += idf * ((tf * (k1 + 1)) / denom);
    }
    if (score > 0) scores.set(chunks[i].id, score);
  }
  return scores;
}

function normalize(map: Map<string, number>): Map<string, number> {
  let max = 0;
  for (const v of map.values()) if (v > max) max = v;
  if (max === 0) return map;
  const out = new Map<string, number>();
  for (const [k, v] of map) out.set(k, v / max);
  return out;
}

function makeSnippet(text: string, max = 320): string {
  const clean = text.replace(/\s+/g, " ").trim();
  return clean.length <= max ? clean : clean.slice(0, max - 1) + "…";
}

export interface SearchOptions {
  topK: number;
  hybrid: boolean;
}

export function search(
  index: RagIndex,
  queryVector: number[],
  queryText: string,
  opts: SearchOptions,
): SearchHit[] {
  const chunks = Object.values(index.chunks);
  if (chunks.length === 0) return [];

  const semantic = new Map<string, number>();
  for (const c of chunks) semantic.set(c.id, cosine(queryVector, c.vector));

  let combined: Map<string, number>;
  let lexicalNorm = new Map<string, number>();
  if (opts.hybrid) {
    const lexical = bm25Scores(chunks, tokenize(queryText));
    lexicalNorm = normalize(lexical);
    const semanticNorm = normalize(semantic);
    combined = new Map<string, number>();
    for (const c of chunks) {
      const s = semanticNorm.get(c.id) ?? 0;
      const l = lexicalNorm.get(c.id) ?? 0;
      combined.set(c.id, 0.6 * s + 0.4 * l);
    }
  } else {
    combined = semantic;
  }

  return chunks
    .map((c) => ({
      chunk: c,
      score: combined.get(c.id) ?? 0,
      semantic: semantic.get(c.id) ?? 0,
      lexical: lexicalNorm.get(c.id) ?? 0,
      snippet: makeSnippet(c.text),
    }))
    .sort((a, b) => b.score - a.score)
    .slice(0, Math.max(1, opts.topK));
}