import { LMStudioClient } from "@lmstudio/sdk";
import { homedir } from "node:os";
import { join, resolve, relative, isAbsolute } from "node:path";
import { readFile, writeFile, rename, readdir, unlink } from "node:fs/promises";

export const PRIMER_ID = "strata-vision-primer";
export const PRIMER_MODEL = "ggml-org/SmolVLM-256M-Instruct-GGUF/SmolVLM-256M-Instruct-Q8_0.gguf";
export const PLUGIN_ID = "yuanjumay/strata-endpoint";
// Existing local-install conversations remain eligible during migration.
const LEGACY_PLUGIN_ID = "local/strata-endpoint";

export function primerConversation(raw: string, model = PRIMER_MODEL, identifier = PRIMER_ID): string | undefined {
  const chat = JSON.parse(raw);
  if (!Array.isArray(chat.plugins) || !chat.plugins.some((id: unknown) => id === PLUGIN_ID || id === LEGACY_PLUGIN_ID)) return;
  // Never replace another selected local model. This workaround only fills an empty slot.
  if (chat.lastUsedModel && !(chat.lastUsedModel.indexedModelIdentifier === model && chat.lastUsedModel.identifier === PRIMER_ID && identifier !== PRIMER_ID)) return;
  chat.lastUsedModel = {
    indexedModelIdentifier: model,
    identifier,
    instanceLoadTimeConfig: { fields: [] },
    instanceOperationTimeConfig: { fields: [] },
  };
  return JSON.stringify(chat, null, 2);
}

export function conversationPath(home: string, identifier: string): string {
  const root = resolve(home, "conversations");
  const target = resolve(root, identifier);
  const rel = relative(root, target);
  if (!rel || isAbsolute(rel) || rel.startsWith("..") || !target.endsWith(".conversation.json")) throw new Error("Invalid active conversation path.");
  return target;
}

export async function primeActiveConversations(home: string, model = PRIMER_MODEL, identifier = PRIMER_ID): Promise<number> {
  const stateRoot = join(home, ".internal", "ui-state");
  const windows = (await readdir(stateRoot)).filter(name => /^window-\d+\.json$/.test(name));
  const targets = new Set<string>();
  for (const window of windows) {
    const ui = JSON.parse(await readFile(join(stateRoot, window), "utf8"));
    const id = ui.chat?.activeConversationIdentifier;
    if (typeof id === "string") targets.add(conversationPath(home, id));
  }
  let changed = 0;
  for (const path of targets) {
    const before = await readFile(path, "utf8");
    const after = primerConversation(before, model, identifier);
    if (!after) continue;
    // Back up before touching the active chat; do not use a newest-file heuristic.
    try { await writeFile(`${path}.strata-vision-primer.bak`, before, { flag: "wx" }); }
    catch (error) { if ((error as NodeJS.ErrnoException).code !== "EEXIST") throw error; }
    if (await readFile(path, "utf8") !== before) throw new Error("Conversation changed during priming; retry after editing stops.");
    const temp = `${path}.strata-primer-${process.pid}-${Date.now()}.tmp`;
    try {
      await writeFile(temp, after, { flag: "wx" });
      // Recheck immediately before replacement to avoid overwriting an in-flight edit.
      if (await readFile(path, "utf8") !== before) throw new Error("Conversation changed before priming; retry.");
      await rename(temp, path);
    } finally { await unlink(temp).catch(() => {}); }
    changed++;
  }
  return changed;
}

export async function ensureVisionModel(client: LMStudioClient): Promise<string> {
  const loaded = await Promise.all((await client.llm.listLoaded()).map(model => model.getModelInfo()));
  const own = loaded.find(model => model.identifier === PRIMER_ID);
  if (own && own.path !== PRIMER_MODEL) throw new Error("Primer identifier is already used by another model.");
  // Also reuse a copy the user loaded manually; do not create a second copy.
  const existing = own ?? loaded.find(model => model.path === PRIMER_MODEL && model.vision);
  if (existing) return existing.identifier;
  const installed = await client.system.listDownloadedModels("llm");
  if (!installed.some(model => model.path === PRIMER_MODEL && model.vision)) throw new Error("Install SmolVLM-256M-Instruct Q8_0 with its vision adapter first. No model is downloaded automatically.");
  const model = await client.llm.load(PRIMER_MODEL, {
    identifier: PRIMER_ID,
    config: { gpu: { ratio: "off" }, offloadKVCacheToGpu: false, contextLength: 512 },
    verbose: false,
    signal: AbortSignal.timeout(30_000),
    // No idle TTL: the helper receives no predictions, so a TTL would evict it
    // even while the user is actively chatting with Strata.
  });
  return (await model.getModelInfo()).identifier;
}

export async function startVisionPrimer(home = join(homedir(), ".lmstudio"), followActiveChat = false, suppliedClient?: LMStudioClient): Promise<void> {
  const client = suppliedClient ?? new LMStudioClient({
    baseUrl: process.env.LMS_PLUGIN_BASE_URL,
    clientIdentifier: process.env.LMS_PLUGIN_CLIENT_IDENTIFIER,
    clientPasskey: process.env.LMS_PLUGIN_CLIENT_PASSKEY,
  });
  let busy = false;
  let identifier: string | undefined;
  let nextCheck = 0;
  const tick = async () => {
    if (busy) return;
    busy = true;
    try {
      if (!identifier && Date.now() < nextCheck) return;
      if (Date.now() >= nextCheck) {
        identifier = undefined;
        identifier = await ensureVisionModel(client);
        nextCheck = Date.now() + 30_000;
      }
      if (!identifier) return;
      const changed = await primeActiveConversations(home, PRIMER_MODEL, identifier);
      if (changed) console.info(`[Strata] Vision primer ready (${identifier}); updated ${changed} active Strata conversation(s).`);
    } finally { busy = false; }
  };
  // Install retries before the initial attempt so a transient startup failure
  // does not permanently disable automatic loading for this plugin process.
  if (followActiveChat) {
    const timer = setInterval(() => {
      void tick().catch(error => { nextCheck = Date.now() + 5000; console.debug("[Strata] Vision primer retry:", error.message); });
    }, 2000);
    timer.unref();
  }
  await tick();
  console.info(`[Strata] Vision primer ready (${identifier}).`);
}

import { LMStudioClient } from "@lmstudio/sdk";
import { homedir } from "node:os";
import { join, resolve, relative, isAbsolute } from "node:path";
import { readFile, writeFile, rename, readdir, unlink } from "node:fs/promises";

export const PRIMER_ID = "strata-vision-primer";
export const PRIMER_MODEL = "ggml-org/SmolVLM-256M-Instruct-GGUF/SmolVLM-256M-Instruct-Q8_0.gguf";
export const PLUGIN_ID = "yuanjumay/strata-endpoint";
// Existing local-install conversations remain eligible during migration.
const LEGACY_PLUGIN_ID = "local/strata-endpoint";

export function primerConversation(raw: string, model = PRIMER_MODEL, identifier = PRIMER_ID): string | undefined {
  const chat = JSON.parse(raw);
  if (!Array.isArray(chat.plugins) || !chat.plugins.some((id: unknown) => id === PLUGIN_ID || id === LEGACY_PLUGIN_ID)) return;
  // Never replace another selected local model. This workaround only fills an empty slot.
  if (chat.lastUsedModel && !(chat.lastUsedModel.indexedModelIdentifier === model && chat.lastUsedModel.identifier === PRIMER_ID && identifier !== PRIMER_ID)) return;
  chat.lastUsedModel = {
    indexedModelIdentifier: model,
    identifier,
    instanceLoadTimeConfig: { fields: [] },
    instanceOperationTimeConfig: { fields: [] },
  };
  return JSON.stringify(chat, null, 2);
}

export function conversationPath(home: string, identifier: string): string {
  const root = resolve(home, "conversations");
  const target = resolve(root, identifier);
  const rel = relative(root, target);
  if (!rel || isAbsolute(rel) || rel.startsWith("..") || !target.endsWith(".conversation.json")) throw new Error("Invalid active conversation path.");
  return target;
}

export async function primeActiveConversations(home: string, model = PRIMER_MODEL, identifier = PRIMER_ID): Promise<number> {
  const stateRoot = join(home, ".internal", "ui-state");
  const windows = (await readdir(stateRoot)).filter(name => /^window-\d+\.json$/.test(name));
  const targets = new Set<string>();
  for (const window of windows) {
    const ui = JSON.parse(await readFile(join(stateRoot, window), "utf8"));
    const id = ui.chat?.activeConversationIdentifier;
    if (typeof id === "string") targets.add(conversationPath(home, id));
  }
  let changed = 0;
  for (const path of targets) {
    const before = await readFile(path, "utf8");
    const after = primerConversation(before, model, identifier);
    if (!after) continue;
    // Back up before touching the active chat; do not use a newest-file heuristic.
    try { await writeFile(`${path}.strata-vision-primer.bak`, before, { flag: "wx" }); }
    catch (error) { if ((error as NodeJS.ErrnoException).code !== "EEXIST") throw error; }
    if (await readFile(path, "utf8") !== before) throw new Error("Conversation changed during priming; retry after editing stops.");
    const temp = `${path}.strata-primer-${process.pid}-${Date.now()}.tmp`;
    try {
      await writeFile(temp, after, { flag: "wx" });
      // Recheck immediately before replacement to avoid overwriting an in-flight edit.
      if (await readFile(path, "utf8") !== before) throw new Error("Conversation changed before priming; retry.");
      await rename(temp, path);
    } finally { await unlink(temp).catch(() => {}); }
    changed++;
  }
  return changed;
}

export async function ensureVisionModel(client: LMStudioClient): Promise<string> {
  const loaded = await Promise.all((await client.llm.listLoaded()).map(model => model.getModelInfo()));
  const own = loaded.find(model => model.identifier === PRIMER_ID);
  if (own && own.path !== PRIMER_MODEL) throw new Error("Primer identifier is already used by another model.");
  // Also reuse a copy the user loaded manually; do not create a second copy.
  const existing = own ?? loaded.find(model => model.path === PRIMER_MODEL && model.vision);
  if (existing) return existing.identifier;
  const installed = await client.system.listDownloadedModels("llm");
  if (!installed.some(model => model.path === PRIMER_MODEL && model.vision)) throw new Error("Install SmolVLM-256M-Instruct Q8_0 with its vision adapter first. No model is downloaded automatically.");
  const model = await client.llm.load(PRIMER_MODEL, {
    identifier: PRIMER_ID,
    config: { gpu: { ratio: "off" }, offloadKVCacheToGpu: false, contextLength: 512 },
    verbose: false,
    signal: AbortSignal.timeout(30_000),
    // No idle TTL: the helper receives no predictions, so a TTL would evict it
    // even while the user is actively chatting with Strata.
  });
  return (await model.getModelInfo()).identifier;
}

export async function startVisionPrimer(home = join(homedir(), ".lmstudio"), followActiveChat = false, suppliedClient?: LMStudioClient): Promise<void> {
  const client = suppliedClient ?? new LMStudioClient({
    baseUrl: process.env.LMS_PLUGIN_BASE_URL,
    clientIdentifier: process.env.LMS_PLUGIN_CLIENT_IDENTIFIER,
    clientPasskey: process.env.LMS_PLUGIN_CLIENT_PASSKEY,
  });
  let busy = false;
  let identifier: string | undefined;
  let nextCheck = 0;
  const tick = async () => {
    if (busy) return;
    busy = true;
    try {
      if (!identifier && Date.now() < nextCheck) return;
      if (Date.now() >= nextCheck) {
        identifier = undefined;
        identifier = await ensureVisionModel(client);
        nextCheck = Date.now() + 30_000;
      }
      if (!identifier) return;
      const changed = await primeActiveConversations(home, PRIMER_MODEL, identifier);
      if (changed) console.info(`[Strata] Vision primer ready (${identifier}); updated ${changed} active Strata conversation(s).`);
    } finally { busy = false; }
  };
  // Install retries before the initial attempt so a transient startup failure
  // does not permanently disable automatic loading for this plugin process.
  if (followActiveChat) {
    const timer = setInterval(() => {
      void tick().catch(error => { nextCheck = Date.now() + 5000; console.debug("[Strata] Vision primer retry:", error.message); });
    }, 2000);
    timer.unref();
  }
  await tick();
  console.info(`[Strata] Vision primer ready (${identifier}).`);
}