// An index of file and folder names, so the model can find where something is without reading
// anything. Only names are kept, never contents. Hidden folders (names starting with a dot),
// app data, system folders and node_modules are never entered, nor anything in the user's list
// of excluded folders.
//
// LM Studio stops plugins between replies, so the index is built in steps that survive a stop:
// progress is saved every few seconds and the next start carries on from there.

import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import { isInside } from "./scope";
import { dataDir, readJson, writeFileAtomic, writeJson } from "./store";

const SKIP_DIRS = new Set(
  [
    "node_modules",
    "appdata",
    "application data",
    "local settings",
    "$recycle.bin",
    "system volume information",
    "windows",
    "program files",
    "program files (x86)",
    "programdata",
    "__pycache__",
  ].concat(process.platform === "darwin" ? ["library"] : [])
);

const MAX_ENTRIES = 1_500_000;
const SAVE_EVERY_MS = 3_000;
export const MAX_AGE_MS = 24 * 60 * 60 * 1000;

export function skipDir(name: string, full: string, excludes: string[]): boolean {
  if (name.startsWith(".")) return true;
  if (SKIP_DIRS.has(name.toLowerCase())) return true;
  return excludes.some((e) => isInside(full, e));
}

interface Meta {
  builtAt: number;
  roots: string[];
  excludes: string[];
  count: number;
  complete: boolean;
}

interface BuildState {
  roots: string[];
  excludes: string[];
  queue: string[];
  count: number;
  startedAt: number;
}

const dir = () => path.join(dataDir(), "names");
const files = () => ({
  index: path.join(dir(), "index.txt"),
  meta: path.join(dir(), "meta.json"),
  build: path.join(dir(), "build.json"),
  partial: path.join(dir(), "partial.txt"),
});

const sameList = (a: string[], b: string[]) => a.length === b.length && a.every((x, i) => x === b[i]);

// The build running in this process, if any. Only one may write the index files at a time, so a
// new build (for example after the folders in the settings changed) stops the old one first.
let running: { owner: NameIndex; stop: boolean; done: Promise<void> } | null = null;

export class NameIndex {
  private building: Promise<void> | null = null;
  private cache: { mtime: number; entries: string[]; lower: string[] } | null = null;

  constructor(
    private roots: string[],
    private excludes: string[],
    private log: (msg: string) => void = () => {}
  ) {}

  meta(): Meta | null {
    return readJson<Meta | null>(files().meta, null);
  }

  isBuilding(): boolean {
    return this.building !== null || fs.existsSync(files().build);
  }

  // Starts or resumes a build when there is none under way and the index is missing, old, or
  // was made for other folders. Never waits for it.
  refreshIfNeeded(force = false): void {
    if (this.building) return;
    const meta = this.meta();
    const state = readJson<BuildState | null>(files().build, null);
    const sameSetup = (m: { roots: string[]; excludes: string[] } | null) => !!m && sameList(m.roots, this.roots) && sameList(m.excludes, this.excludes);
    if (state && sameSetup(state) && !force) {
      this.start(false);
      return;
    }
    if (force || !meta || !sameSetup(meta) || Date.now() - meta.builtAt > MAX_AGE_MS) this.start(true);
  }

  private start(fresh: boolean): void {
    const previous = running;
    if (previous && previous.owner !== this) previous.stop = true;
    const run = { owner: this as NameIndex, stop: false, done: Promise.resolve() };
    run.done = (previous ? previous.done : Promise.resolve())
      .then(() => this.build(fresh, run))
      .catch((e) => this.log(`name index: build failed: ${(e as Error).message}`))
      .finally(() => {
        this.building = null;
        if (running === run) running = null;
      });
    running = run;
    this.building = run.done;
  }

  // For tests and the reindex command: runs the build to the end.
  async buildNow(): Promise<void> {
    if (this.building) await this.building;
    this.start(true);
    await this.building;
  }

  private async build(fresh: boolean, run: { stop: boolean }): Promise<void> {
    const f = files();
    fs.mkdirSync(dir(), { recursive: true });
    let state: BuildState;
    if (fresh) {
      state = { roots: this.roots, excludes: this.excludes, queue: this.roots.filter((r) => fs.existsSync(r)), count: 0, startedAt: Date.now() };
      fs.writeFileSync(f.partial, "");
      writeJson(f.build, state);
      this.log("name index: building");
    } else {
      state = readJson<BuildState>(f.build, { roots: this.roots, excludes: this.excludes, queue: [], count: 0, startedAt: Date.now() });
      this.log(`name index: resuming, ${state.count} entries so far`);
    }
    let batch: string[] = [];
    let lastSave = Date.now();
    const save = () => {
      if (batch.length) fs.appendFileSync(f.partial, batch.join("\n") + "\n", "utf8");
      batch = [];
      writeJson(f.build, state);
      lastSave = Date.now();
    };
    while (state.queue.length && state.count < MAX_ENTRIES) {
      if (run.stop) {
        save();
        this.log("name index: build stopped for a newer one");
        return;
      }
      const folder = state.queue.pop()!;
      let handle: fs.Dir;
      try {
        handle = await fs.promises.opendir(folder);
      } catch {
        continue;
      }
      const found: string[] = [];
      try {
        for await (const entry of handle) {
          const full = path.join(folder, entry.name);
          if (entry.isSymbolicLink()) continue;
          if (entry.isDirectory()) {
            if (skipDir(entry.name, full, this.excludes)) continue;
            found.push(full + path.sep);
            state.queue.push(full);
          } else if (entry.isFile()) {
            found.push(full);
          }
        }
      } catch {
        // a folder that cannot be read to the end keeps what was read
      }
      // A folder's entries are saved together with its removal from the queue, so a stop
      // between saves repeats at most the folders since the last save.
      batch.push(...found);
      state.count += found.length;
      if (Date.now() - lastSave > SAVE_EVERY_MS) save();
    }
    save();
    const complete = state.queue.length === 0;
    fs.renameSync(f.partial, f.index);
    writeJson(f.meta, { builtAt: Date.now(), roots: state.roots, excludes: state.excludes, count: state.count, complete } satisfies Meta);
    fs.rmSync(f.build, { force: true });
    this.cache = null;
    this.log(`name index: built, ${state.count} entries${complete ? "" : " (stopped at the size limit)"}`);
  }

  private load(): { entries: string[]; lower: string[] } | null {
    let mtime: number;
    try {
      mtime = fs.statSync(files().index).mtimeMs;
    } catch {
      return null;
    }
    if (!this.cache || this.cache.mtime !== mtime) {
      const entries = fs.readFileSync(files().index, "utf8").split("\n").filter(Boolean);
      this.cache = { mtime, entries, lower: entries.map((e) => e.toLowerCase()) };
    }
    return this.cache;
  }

  // Entries whose path contains every word of the query, best first: words in the name count
  // more than words in the folders above it, an exact name counts most, and shorter paths win
  // ties. "*.md" or ".md" limits the search to that extension. Null when there is no index yet.
  find(query: string, limit = 25): string[] | null {
    const data = this.load();
    if (!data) return null;
    const words = query
      .toLowerCase()
      .split(/[\s\\/]+/)
      .map((w) => w.replace(/^\*+/, ""))
      .filter(Boolean);
    if (!words.length) return [];
    const ext = words.find((w) => /^\.[a-z0-9]+$/.test(w));
    const terms = words.filter((w) => w !== ext);
    const scored: Array<{ i: number; score: number }> = [];
    for (let i = 0; i < data.lower.length; i++) {
      const p = data.lower[i];
      if (ext && !p.endsWith(ext)) continue;
      let ok = true;
      for (const t of terms) {
        if (!p.includes(t)) {
          ok = false;
          break;
        }
      }
      if (!ok) continue;
      const base = path.basename(p.endsWith(path.sep) ? p.slice(0, -1) : p);
      let score = 0;
      for (const t of terms) score += base.includes(t) ? 3 : 1;
      const bare = base.replace(/\.[^.]+$/, "");
      if (terms.length && (base === terms.join(" ") || bare === terms.join(" "))) score += 10;
      scored.push({ i, score: score - p.length / 1000 });
    }
    scored.sort((a, b) => b.score - a.score);
    return scored.slice(0, limit).map((s) => data.entries[s.i]);
  }
}

export function defaultRoots(): string[] {
  return [os.homedir()];
}

// An index of file and folder names, so the model can find where something is without reading
// anything. Only names are kept, never contents. Hidden folders (names starting with a dot),
// app data, system folders and node_modules are never entered, nor anything in the user's list
// of excluded folders.
//
// LM Studio stops plugins between replies, so the index is built in steps that survive a stop:
// progress is saved every few seconds and the next start carries on from there.

import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import { isInside } from "./scope";
import { dataDir, readJson, writeFileAtomic, writeJson } from "./store";

const SKIP_DIRS = new Set(
  [
    "node_modules",
    "appdata",
    "application data",
    "local settings",
    "$recycle.bin",
    "system volume information",
    "windows",
    "program files",
    "program files (x86)",
    "programdata",
    "__pycache__",
  ].concat(process.platform === "darwin" ? ["library"] : [])
);

const MAX_ENTRIES = 1_500_000;
const SAVE_EVERY_MS = 3_000;
export const MAX_AGE_MS = 24 * 60 * 60 * 1000;

export function skipDir(name: string, full: string, excludes: string[]): boolean {
  if (name.startsWith(".")) return true;
  if (SKIP_DIRS.has(name.toLowerCase())) return true;
  return excludes.some((e) => isInside(full, e));
}

interface Meta {
  builtAt: number;
  roots: string[];
  excludes: string[];
  count: number;
  complete: boolean;
}

interface BuildState {
  roots: string[];
  excludes: string[];
  queue: string[];
  count: number;
  startedAt: number;
}

const dir = () => path.join(dataDir(), "names");
const files = () => ({
  index: path.join(dir(), "index.txt"),
  meta: path.join(dir(), "meta.json"),
  build: path.join(dir(), "build.json"),
  partial: path.join(dir(), "partial.txt"),
});

const sameList = (a: string[], b: string[]) => a.length === b.length && a.every((x, i) => x === b[i]);

// The build running in this process, if any. Only one may write the index files at a time, so a
// new build (for example after the folders in the settings changed) stops the old one first.
let running: { owner: NameIndex; stop: boolean; done: Promise<void> } | null = null;

export class NameIndex {
  private building: Promise<void> | null = null;
  private cache: { mtime: number; entries: string[]; lower: string[] } | null = null;

  constructor(
    private roots: string[],
    private excludes: string[],
    private log: (msg: string) => void = () => {}
  ) {}

  meta(): Meta | null {
    return readJson<Meta | null>(files().meta, null);
  }

  isBuilding(): boolean {
    return this.building !== null || fs.existsSync(files().build);
  }

  // Starts or resumes a build when there is none under way and the index is missing, old, or
  // was made for other folders. Never waits for it.
  refreshIfNeeded(force = false): void {
    if (this.building) return;
    const meta = this.meta();
    const state = readJson<BuildState | null>(files().build, null);
    const sameSetup = (m: { roots: string[]; excludes: string[] } | null) => !!m && sameList(m.roots, this.roots) && sameList(m.excludes, this.excludes);
    if (state && sameSetup(state) && !force) {
      this.start(false);
      return;
    }
    if (force || !meta || !sameSetup(meta) || Date.now() - meta.builtAt > MAX_AGE_MS) this.start(true);
  }

  private start(fresh: boolean): void {
    const previous = running;
    if (previous && previous.owner !== this) previous.stop = true;
    const run = { owner: this as NameIndex, stop: false, done: Promise.resolve() };
    run.done = (previous ? previous.done : Promise.resolve())
      .then(() => this.build(fresh, run))
      .catch((e) => this.log(`name index: build failed: ${(e as Error).message}`))
      .finally(() => {
        this.building = null;
        if (running === run) running = null;
      });
    running = run;
    this.building = run.done;
  }

  // For tests and the reindex command: runs the build to the end.
  async buildNow(): Promise<void> {
    if (this.building) await this.building;
    this.start(true);
    await this.building;
  }

  private async build(fresh: boolean, run: { stop: boolean }): Promise<void> {
    const f = files();
    fs.mkdirSync(dir(), { recursive: true });
    let state: BuildState;
    if (fresh) {
      state = { roots: this.roots, excludes: this.excludes, queue: this.roots.filter((r) => fs.existsSync(r)), count: 0, startedAt: Date.now() };
      fs.writeFileSync(f.partial, "");
      writeJson(f.build, state);
      this.log("name index: building");
    } else {
      state = readJson<BuildState>(f.build, { roots: this.roots, excludes: this.excludes, queue: [], count: 0, startedAt: Date.now() });
      this.log(`name index: resuming, ${state.count} entries so far`);
    }
    let batch: string[] = [];
    let lastSave = Date.now();
    const save = () => {
      if (batch.length) fs.appendFileSync(f.partial, batch.join("\n") + "\n", "utf8");
      batch = [];
      writeJson(f.build, state);
      lastSave = Date.now();
    };
    while (state.queue.length && state.count < MAX_ENTRIES) {
      if (run.stop) {
        save();
        this.log("name index: build stopped for a newer one");
        return;
      }
      const folder = state.queue.pop()!;
      let handle: fs.Dir;
      try {
        handle = await fs.promises.opendir(folder);
      } catch {
        continue;
      }
      const found: string[] = [];
      try {
        for await (const entry of handle) {
          const full = path.join(folder, entry.name);
          if (entry.isSymbolicLink()) continue;
          if (entry.isDirectory()) {
            if (skipDir(entry.name, full, this.excludes)) continue;
            found.push(full + path.sep);
            state.queue.push(full);
          } else if (entry.isFile()) {
            found.push(full);
          }
        }
      } catch {
        // a folder that cannot be read to the end keeps what was read
      }
      // A folder's entries are saved together with its removal from the queue, so a stop
      // between saves repeats at most the folders since the last save.
      batch.push(...found);
      state.count += found.length;
      if (Date.now() - lastSave > SAVE_EVERY_MS) save();
    }
    save();
    const complete = state.queue.length === 0;
    fs.renameSync(f.partial, f.index);
    writeJson(f.meta, { builtAt: Date.now(), roots: state.roots, excludes: state.excludes, count: state.count, complete } satisfies Meta);
    fs.rmSync(f.build, { force: true });
    this.cache = null;
    this.log(`name index: built, ${state.count} entries${complete ? "" : " (stopped at the size limit)"}`);
  }

  private load(): { entries: string[]; lower: string[] } | null {
    let mtime: number;
    try {
      mtime = fs.statSync(files().index).mtimeMs;
    } catch {
      return null;
    }
    if (!this.cache || this.cache.mtime !== mtime) {
      const entries = fs.readFileSync(files().index, "utf8").split("\n").filter(Boolean);
      this.cache = { mtime, entries, lower: entries.map((e) => e.toLowerCase()) };
    }
    return this.cache;
  }

  // Entries whose path contains every word of the query, best first: words in the name count
  // more than words in the folders above it, an exact name counts most, and shorter paths win
  // ties. "*.md" or ".md" limits the search to that extension. Null when there is no index yet.
  find(query: string, limit = 25): string[] | null {
    const data = this.load();
    if (!data) return null;
    const words = query
      .toLowerCase()
      .split(/[\s\\/]+/)
      .map((w) => w.replace(/^\*+/, ""))
      .filter(Boolean);
    if (!words.length) return [];
    const ext = words.find((w) => /^\.[a-z0-9]+$/.test(w));
    const terms = words.filter((w) => w !== ext);
    const scored: Array<{ i: number; score: number }> = [];
    for (let i = 0; i < data.lower.length; i++) {
      const p = data.lower[i];
      if (ext && !p.endsWith(ext)) continue;
      let ok = true;
      for (const t of terms) {
        if (!p.includes(t)) {
          ok = false;
          break;
        }
      }
      if (!ok) continue;
      const base = path.basename(p.endsWith(path.sep) ? p.slice(0, -1) : p);
      let score = 0;
      for (const t of terms) score += base.includes(t) ? 3 : 1;
      const bare = base.replace(/\.[^.]+$/, "");
      if (terms.length && (base === terms.join(" ") || bare === terms.join(" "))) score += 10;
      scored.push({ i, score: score - p.length / 1000 });
    }
    scored.sort((a, b) => b.score - a.score);
    return scored.slice(0, limit).map((s) => data.entries[s.i]);
  }
}

export function defaultRoots(): string[] {
  return [os.homedir()];
}