import {
  text,
  type Chat,
  type ChatMessage,
  type FileHandle,
  type LLMDynamicHandle,
  type PredictionProcessStatusController,
  type PromptPreprocessorController,
} from "@lmstudio/sdk";
import { readFile, writeFile, readdir, stat } from "fs/promises";
import { dirname, join, extname } from "path";
import { pluginConfigSchematics } from "./config";
import { buildFirstTurnDocs } from "./firstTurnInjection";
import { getPersistedState, savePersistedState, getRuntimeWorkingDirectory } from "./stateManager";
import { getDict } from "./locales/i18n";
import { onPredictionStart, injectPluginStatus } from "./contextInjector";

type DocumentContextInjectionStrategy = "none" | "inject-full-content" | "retrieval";

export function getSubAgentDocsCandidatePaths(currentWorkingDirectory: string): string[] {
  return [
    join(dirname(__dirname), "subagent_docs.md"),
    join(dirname(__dirname), "instructions", "subagent_docs.md"),
    join(currentWorkingDirectory, "instructions", "subagent_docs.md"),
    join(currentWorkingDirectory, "subagent_docs.md"),
  ];
}

export async function promptPreprocessor(ctl: PromptPreprocessorController, userMessage: ChatMessage) {
  const userPrompt = userMessage.getText();
  
  const pluginConfig = ctl.getPluginConfig(pluginConfigSchematics);
  
  const history = await ctl.pullHistory();

  let isFirstTurn = false;
  if (Array.isArray(history)) {
    isFirstTurn = history.length === 0;
  } else if ("messages" in history && Array.isArray((history as any).messages)) {
    isFirstTurn = (history as any).messages.length === 0;
  } else if ("length" in history && typeof (history as any).length === "number") {
    isFirstTurn = (history as any).length === 0;
  } else {
    isFirstTurn = true; 
  }

  history.append(userMessage);
  
  const newFiles = userMessage.getFiles(ctl.client).filter(f => f.type !== "image");
  const files = history.getAllFiles(ctl.client).filter(f => f.type !== "image");

  let processingResult: string | ChatMessage | null = null;

  const enableLocalRag = pluginConfig.get("localRag.enabled");
  const unrestricted = pluginConfig.get("general.unrestrictedMode") === true;

  if (enableLocalRag || unrestricted) {
      if (newFiles.length > 0) {
          const strategy = await chooseContextInjectionStrategy(ctl, userPrompt, newFiles);
          if (strategy === "inject-full-content") {
              processingResult = await prepareDocumentContextInjection(ctl, userMessage);
          } else if (strategy === "retrieval") {
              processingResult = await prepareRetrievalResultsContextInjection(ctl, userPrompt, files);
          }
      } else if (files.length > 0) {
          processingResult = await prepareRetrievalResultsContextInjection(ctl, userPrompt, files);
      }
  }

  let currentContent: string;
  if (processingResult) {
      if (typeof processingResult === 'string') {
          currentContent = processingResult;
      } else {
          currentContent = processingResult.getText();
      }
  } else {
      currentContent = userPrompt;
  }

  const defaultWorkspacePath = pluginConfig.get("general.defaultWorkspacePath");
  const lmStudioCwd = getRuntimeWorkingDirectory(ctl);

  const frequency = pluginConfig.get("agent.subAgentFrequency");
  const debugMode = pluginConfig.get("debug.enableDebugMode");

  const messageLanguage = pluginConfig.get("general.messageLanguage");
  const rt = getDict(messageLanguage).runtime;

  const uiLanguageOverride = pluginConfig.get("general.uiLanguageOverride");
  try {
    const state = await getPersistedState(defaultWorkspacePath, lmStudioCwd);
    if (state.uiLanguageOverride !== uiLanguageOverride) {
      state.uiLanguageOverride = uiLanguageOverride;
      await savePersistedState(state);
      ctl.debug(`[i18n] uiLanguageOverride saved: "${uiLanguageOverride}". Restart the plugin to apply the new UI language.`);
    }
  } catch (e) {
    ctl.debug("[i18n] Failed to persist uiLanguageOverride.", e);
  }

  const planMode = pluginConfig.get("general.planMode");
  
  let planHint = "";
  
  if (planMode === "always") {
      planHint = rt.planHintAlways;
  } else if (planMode === "when_useful") {
      planHint = rt.planHintWhenUseful;
  }

  // Delegation hints are intentionally NOT auto-appended.
  // Sub-agent is only used when the user asks to delegate or the model
  // calls consult_secondary_agent / agent_run_with_retry explicitly.
  void frequency;
  void debugMode;

  if (planHint) {
      currentContent += planHint;
  }

  const enableSecondary = pluginConfig.get("agent.enableSecondaryAgent");
  const state = await getPersistedState(defaultWorkspacePath, lmStudioCwd);

  if (isFirstTurn) {
      state.subAgentDocsInjected = false;
      await savePersistedState(state);
  }

  if (enableSecondary && !state.subAgentDocsInjected) {
      const { currentWorkingDirectory } = state;
      const candidatePaths = getSubAgentDocsCandidatePaths(currentWorkingDirectory);

      let docsInjected = false;
      for (const subAgentDocsPath of candidatePaths) {
          try {
              const docsContent = await readFile(subAgentDocsPath, "utf-8");
              if (docsContent && docsContent.trim().length > 0) {
                  currentContent += `\n\n---\n\n${docsContent}\n\n---\n\n`;
                  ctl.debug(`subagent_docs.md injected into context from: ${subAgentDocsPath}`);
                  state.subAgentDocsInjected = true;
                  await savePersistedState(state);
                  docsInjected = true;
                  break;
              }
          } catch (e) {
              // Keep trying fallback paths.
          }
      }

      if (!docsInjected) {
          ctl.debug("subagent_docs.md not found or failed to load from plugin/workspace paths. Skipping injection.");
      }
  }

  if (isFirstTurn) {
    const simpleSystemPrompt = pluginConfig.get("debug.simpleSystemPrompt");
    const { content: docsBlock, usedLite, categories } = buildFirstTurnDocs(
      userPrompt,
      !!simpleSystemPrompt,
      6000,
    );
    ctl.debug(
      `First-turn tool docs: ${usedLite ? "LITE" : "FILTERED"} ` +
        `(categories: ${categories.join(", ") || "core"})`,
    );
    let injectionContent: string = docsBlock;

    try {
        const { currentWorkingDirectory } = state;
        const candidateStartupPaths = [
            join(currentWorkingDirectory, ".beledarian", "startup.md"),
            join(currentWorkingDirectory, "instructions", "startup.md"),
            join(currentWorkingDirectory, "startup.md"),
        ];

        let startupContent = "";
        let usedStartupPath = "";
        for (const startupPath of candidateStartupPaths) {
            try {
                startupContent = await readFile(startupPath, "utf-8");
                usedStartupPath = dirname(startupPath);
                ctl.debug(`startup.md loaded from: ${startupPath}`);
                break;
            } catch (e) {
                // Keep trying
            }
        }

        if (startupContent) {
            const filesToRead = startupContent.split('\n').map(f => f.trim()).filter(f => f);

            for (const file of filesToRead) {
                const candidateFilePaths = [
                    join(usedStartupPath, file),
                    join(currentWorkingDirectory, file),
                ];

                let loaded = false;
                for (const filePath of candidateFilePaths) {
                    try {
                        const fileContent = await readFile(filePath, "utf-8");
                        if (fileContent.trim().length > 0) {
                            injectionContent = `\n\n---\n\n${fileContent}\n\n---\n\n${injectionContent}`;
                            ctl.debug(`${file} loaded and injected into context from ${filePath}.`);
                            loaded = true;
                            break;
                        }
                    } catch (e) {
                        // Keep trying
                    }
                }
                if (!loaded) {
                    ctl.debug(`Failed to load ${file} from startup.md.`);
                }
            }
        }
    } catch (e) {
        ctl.debug("No startup.md file found or failed to load.");
    }

    // --- Workspace .md auto-intake (the part that was missing) ---
    // Previously ONLY startup.md-listed files were injected and ONLY on first turn.
    // Now: memory.md + common project .md files are read from the LIVE workspace
    // CWD and appended so the model actually sees them.
    try {
        const { currentWorkingDirectory: liveCwd } = state;
        const mdCandidates: string[] = [
            join(liveCwd, "memory.md"),
            join(liveCwd, ".beledarian", "startup.md"),
            join(liveCwd, "instructions", "startup.md"),
            join(liveCwd, "startup.md"),
            join(liveCwd, "AGENTS.md"),
            join(liveCwd, "MUSE.md"),
            join(liveCwd, "GEMINI.md"),
            join(liveCwd, "timbers-toolbox.md"),
        ];
        // Also pick up any other small top-level .md files (project notes, specs).
        try {
            const topEntries = await readdir(liveCwd, { withFileTypes: true });
            for (const e of topEntries) {
                if (!e.isFile()) continue;
                if (extname(e.name).toLowerCase() !== ".md") continue;
                const full = join(liveCwd, e.name);
                if (mdCandidates.includes(full)) continue;
                try {
                    const s = await stat(full);
                    // Keep intake bounded: skip huge docs here (RAG/file tools cover those).
                    if (s.size > 64 * 1024) {
                        ctl.debug(`Skipping large .md for auto-intake (>64KB): ${e.name}`);
                        continue;
                    }
                    mdCandidates.push(full);
                } catch { /* ignore */ }
            }
        } catch (e) {
            ctl.debug("Workspace .md scan failed.", e);
        }
        const seen = new Set<string>();
        for (const mdPath of mdCandidates) {
            if (seen.has(mdPath)) continue;
            seen.add(mdPath);
            try {
                const mdContent = await readFile(mdPath, "utf-8");
                if (mdContent.trim().length === 0) continue;
                // Cap per-file so one big memory.md can't blow the context.
                const capped = mdContent.length > 12000
                    ? mdContent.slice(0, 12000) + "\n\n[... truncated for context ...]"
                    : mdContent;
                injectionContent = `\n\n---\n\n## Workspace file: ${mdPath}\n\n${capped}\n\n---\n\n${injectionContent}`;
                ctl.debug(`Workspace .md injected: ${mdPath} (${mdContent.length} chars)`);
            } catch {
                // Missing file — expected for most candidates.
            }
        }
    } catch (e) {
        ctl.debug("Workspace .md auto-intake failed.", e);
    }

    currentContent = `${injectionContent}\n\n---\n\n${currentContent}`;
  }

  onPredictionStart();
  const statusBlock = injectPluginStatus("");

  if (currentContent !== userPrompt) {
      return `${currentContent}\n\n${statusBlock}`;
  }

  try {
    state.messageCount++;
    await savePersistedState(state);
  } catch (e) {
    ctl.debug("Failed to update message count or memory.", e);
  }
  
  return `${userMessage.getText()}\n\n${statusBlock}`;
}

async function prepareRetrievalResultsContextInjection(
  ctl: PromptPreprocessorController,
  originalUserPrompt: string,
  files: Array<FileHandle>,
): Promise<string> {
  const pluginConfig = ctl.getPluginConfig(pluginConfigSchematics);
  const retrievalLimit = pluginConfig.get("retrieval.limit");
  const retrievalAffinityThreshold = pluginConfig.get("retrieval.threshold");

  const statusSteps = new Map<FileHandle, PredictionProcessStatusController>();

  const rtRetrieve = getDict(
    ctl.getPluginConfig(pluginConfigSchematics).get("general.messageLanguage")
  ).runtime;

  const retrievingStatus = ctl.createStatus({
    status: "loading",
    text: rtRetrieve.statusLoadingEmbeddingModel,
  });
  
  const customModel = pluginConfig.get("retrieval.embeddingModel");
  const modelName = (customModel && customModel.trim() !== "") ? customModel : "nomic-ai/nomic-embed-text-v1.5-GGUF";

  const model = await ctl.client.embedding.model(modelName, {
    signal: ctl.abortSignal,
  });
  
  retrievingStatus.setState({
    status: "loading",
    text: rtRetrieve.statusRetrievingCitations,
  });
  const result = await ctl.client.files.retrieve(originalUserPrompt, files, {
    embeddingModel: model,
    limit: retrievalLimit,
    signal: ctl.abortSignal,
    onFileProcessList(filesToProcess) {
      for (const file of filesToProcess) {
        statusSteps.set(
          file,
          retrievingStatus.addSubStatus({
            status: "waiting",
            text: `Process ${file.name} for retrieval`,
          }),
        );
      }
    },
    onFileProcessingStart(file) {
      statusSteps
        .get(file)!
        .setState({ status: "loading", text: `Processing ${file.name} for retrieval` });
    },
    onFileProcessingEnd(file) {
      statusSteps
        .get(file)!
        .setState({ status: "done", text: `Processed ${file.name} for retrieval` });
    },
    onFileProcessingStepProgress(file, step, progressInStep) {
      const verb = step === "loading" ? "Loading" : step === "chunking" ? "Chunking" : "Embedding";
      statusSteps.get(file)!.setState({
        status: "loading",
        text: `${verb} ${file.name} for retrieval (${(progressInStep * 100).toFixed(1)}%)`,
      });
    },
  });

  result.entries = result.entries.filter(entry => entry.score > retrievalAffinityThreshold);

  let processedContent = "";
  const numRetrievals = result.entries.length;
  if (numRetrievals > 0) {
    retrievingStatus.setState({
      status: "done",
      text: rtRetrieve.statusRetrievedCitations(numRetrievals),
    });
    ctl.debug("Retrieval results", result);
    processedContent += rtRetrieve.citationPrefix;
    let citationNumber = 1;
    result.entries.forEach(result => {
      const completeText = result.content;
      processedContent += rtRetrieve.citationEntry(citationNumber, completeText);
      citationNumber++;
    });
    await ctl.addCitations(result);
    processedContent += rtRetrieve.citationSuffix(originalUserPrompt);
  } else {
    retrievingStatus.setState({
      status: "canceled",
      text: rtRetrieve.statusNoRelevantCitations,
    });
    ctl.debug("No relevant citations found for user query");
    processedContent =
      rtRetrieve.noRelevantCitationsNote + `\n\nUser Query:\n\n${originalUserPrompt}`;
  }
  ctl.debug("Processed content", processedContent);

  return processedContent;
}

async function prepareDocumentContextInjection(
  ctl: PromptPreprocessorController,
  input: ChatMessage,
): Promise<ChatMessage> {
  const documentInjectionSnippets: Map<FileHandle, string> = new Map();
  const files = input.consumeFiles(ctl.client, file => file.type !== "image");
  for (const file of files) {
    const { content } = await ctl.client.files.parseDocument(file, {
      signal: ctl.abortSignal,
    });

    ctl.debug(text`
      Strategy: inject-full-content. Injecting full content of file '${file}' into the
      context. Length: ${content.length}.
    `);
    documentInjectionSnippets.set(file, content);
  }

  let formattedFinalUserPrompt = "";

  const rtDoc = getDict(
    ctl.getPluginConfig(pluginConfigSchematics).get("general.messageLanguage")
  ).runtime;

  if (documentInjectionSnippets.size > 0) {
    formattedFinalUserPrompt += rtDoc.documentInjectionHeader;

    for (const [fileHandle, snippet] of documentInjectionSnippets) {
      formattedFinalUserPrompt += rtDoc.documentInjectionFileBlock(fileHandle.name, snippet);
    }

    formattedFinalUserPrompt += rtDoc.documentInjectionSuffix(input.getText());
  }

  input.replaceText(formattedFinalUserPrompt);
  return input;
}

async function measureContextWindow(ctx: Chat, model: LLMDynamicHandle) {
  let currentContextFormatted: string;
  
  try {
    const result = await Promise.resolve(model.applyPromptTemplate(ctx));
    currentContextFormatted = result;
  } catch (err) {
    console.warn("[PromptPreprocessor] applyPromptTemplate failed, using fallback:", err instanceof Error ? err.message : String(err));
    
    let combinedText = "";
    try {
      const chatObj = ctx as any;
      if (chatObj.messages && Array.isArray(chatObj.messages)) {
        for (const msg of chatObj.messages) {
          combinedText += (msg.getText ? msg.getText() : (msg.text || "")) + "\n";
        }
      } else if (chatObj.pullHistory) {
        const messages = await chatObj.pullHistory();
        if (Array.isArray(messages)) {
          for (const msg of messages) {
            combinedText += (msg.getText ? msg.getText() : (msg.text || "")) + "\n";
          }
        }
      }
    } catch (innerErr) {
      console.warn("[PromptPreprocessor] Failed to extract messages from Chat object:", innerErr instanceof Error ? innerErr.message : String(innerErr));
      combinedText = "";
    }
    currentContextFormatted = combinedText;
  }
  
  const totalTokensInContext = await model.countTokens(currentContextFormatted);
  const modelContextLength = await model.getContextLength();
  const modelRemainingContextLength = modelContextLength - totalTokensInContext;
  const contextOccupiedPercent = (totalTokensInContext / modelContextLength) * 100;
  return {
    totalTokensInContext,
    modelContextLength,
    modelRemainingContextLength,
    contextOccupiedPercent,
  };
}

async function chooseContextInjectionStrategy(
  ctl: PromptPreprocessorController,
  originalUserPrompt: string,
  files: Array<FileHandle>,
): Promise<DocumentContextInjectionStrategy> {
  const rtStrategy = getDict(
    ctl.getPluginConfig(pluginConfigSchematics).get("general.messageLanguage")
  ).runtime;

  const status = ctl.createStatus({
    status: "loading",
    text: rtStrategy.statusDecidingStrategy,
  });

  const model = await ctl.client.llm.model();
  const ctx = await ctl.pullHistory();

  let contextMeasurement: { totalTokensInContext: number; modelContextLength: number; modelRemainingContextLength: number; contextOccupiedPercent: number };
  try {
    contextMeasurement = await measureContextWindow(ctx, model);
  } catch (err) {
    console.warn("[PromptPreprocessor] measureContextWindow failed:", err instanceof Error ? err.message : String(err));
    status.setState({ status: "canceled", text: "Could not measure context; using retrieval" });
    return "retrieval";
  }

  // Estimate document token cost roughly
  let estimatedDocTokens = 0;
  for (const file of files) {
    try {
      const { content } = await ctl.client.files.parseDocument(file, { signal: ctl.abortSignal });
      estimatedDocTokens += Math.ceil(content.length / 4);
    } catch {
      estimatedDocTokens += 4000;
    }
  }

  if (estimatedDocTokens < contextMeasurement.modelRemainingContextLength * 0.5) {
    status.setState({ status: "done", text: rtStrategy.statusInjectFullContent || "Injecting full document content" });
    return "inject-full-content";
  }

  status.setState({ status: "done", text: rtStrategy.statusUseRetrieval || "Using retrieval" });
  return "retrieval";
}

import {
  text,
  type Chat,
  type ChatMessage,
  type FileHandle,
  type LLMDynamicHandle,
  type PredictionProcessStatusController,
  type PromptPreprocessorController,
} from "@lmstudio/sdk";
import { readFile, writeFile, readdir, stat } from "fs/promises";
import { dirname, join, extname } from "path";
import { pluginConfigSchematics } from "./config";
import { buildFirstTurnDocs } from "./firstTurnInjection";
import { getPersistedState, savePersistedState, getRuntimeWorkingDirectory } from "./stateManager";
import { getDict } from "./locales/i18n";
import { onPredictionStart, injectPluginStatus } from "./contextInjector";

type DocumentContextInjectionStrategy = "none" | "inject-full-content" | "retrieval";

export function getSubAgentDocsCandidatePaths(currentWorkingDirectory: string): string[] {
  return [
    join(dirname(__dirname), "subagent_docs.md"),
    join(dirname(__dirname), "instructions", "subagent_docs.md"),
    join(currentWorkingDirectory, "instructions", "subagent_docs.md"),
    join(currentWorkingDirectory, "subagent_docs.md"),
  ];
}

export async function promptPreprocessor(ctl: PromptPreprocessorController, userMessage: ChatMessage) {
  const userPrompt = userMessage.getText();
  
  const pluginConfig = ctl.getPluginConfig(pluginConfigSchematics);
  
  const history = await ctl.pullHistory();

  let isFirstTurn = false;
  if (Array.isArray(history)) {
    isFirstTurn = history.length === 0;
  } else if ("messages" in history && Array.isArray((history as any).messages)) {
    isFirstTurn = (history as any).messages.length === 0;
  } else if ("length" in history && typeof (history as any).length === "number") {
    isFirstTurn = (history as any).length === 0;
  } else {
    isFirstTurn = true; 
  }

  history.append(userMessage);
  
  const newFiles = userMessage.getFiles(ctl.client).filter(f => f.type !== "image");
  const files = history.getAllFiles(ctl.client).filter(f => f.type !== "image");

  let processingResult: string | ChatMessage | null = null;

  const enableLocalRag = pluginConfig.get("localRag.enabled");
  const unrestricted = pluginConfig.get("general.unrestrictedMode") === true;

  if (enableLocalRag || unrestricted) {
      if (newFiles.length > 0) {
          const strategy = await chooseContextInjectionStrategy(ctl, userPrompt, newFiles);
          if (strategy === "inject-full-content") {
              processingResult = await prepareDocumentContextInjection(ctl, userMessage);
          } else if (strategy === "retrieval") {
              processingResult = await prepareRetrievalResultsContextInjection(ctl, userPrompt, files);
          }
      } else if (files.length > 0) {
          processingResult = await prepareRetrievalResultsContextInjection(ctl, userPrompt, files);
      }
  }

  let currentContent: string;
  if (processingResult) {
      if (typeof processingResult === 'string') {
          currentContent = processingResult;
      } else {
          currentContent = processingResult.getText();
      }
  } else {
      currentContent = userPrompt;
  }

  const defaultWorkspacePath = pluginConfig.get("general.defaultWorkspacePath");
  const lmStudioCwd = getRuntimeWorkingDirectory(ctl);

  const frequency = pluginConfig.get("agent.subAgentFrequency");
  const debugMode = pluginConfig.get("debug.enableDebugMode");

  const messageLanguage = pluginConfig.get("general.messageLanguage");
  const rt = getDict(messageLanguage).runtime;

  const uiLanguageOverride = pluginConfig.get("general.uiLanguageOverride");
  try {
    const state = await getPersistedState(defaultWorkspacePath, lmStudioCwd);
    if (state.uiLanguageOverride !== uiLanguageOverride) {
      state.uiLanguageOverride = uiLanguageOverride;
      await savePersistedState(state);
      ctl.debug(`[i18n] uiLanguageOverride saved: "${uiLanguageOverride}". Restart the plugin to apply the new UI language.`);
    }
  } catch (e) {
    ctl.debug("[i18n] Failed to persist uiLanguageOverride.", e);
  }

  const planMode = pluginConfig.get("general.planMode");
  
  let planHint = "";
  
  if (planMode === "always") {
      planHint = rt.planHintAlways;
  } else if (planMode === "when_useful") {
      planHint = rt.planHintWhenUseful;
  }

  // Delegation hints are intentionally NOT auto-appended.
  // Sub-agent is only used when the user asks to delegate or the model
  // calls consult_secondary_agent / agent_run_with_retry explicitly.
  void frequency;
  void debugMode;

  if (planHint) {
      currentContent += planHint;
  }

  const enableSecondary = pluginConfig.get("agent.enableSecondaryAgent");
  const state = await getPersistedState(defaultWorkspacePath, lmStudioCwd);

  if (isFirstTurn) {
      state.subAgentDocsInjected = false;
      await savePersistedState(state);
  }

  if (enableSecondary && !state.subAgentDocsInjected) {
      const { currentWorkingDirectory } = state;
      const candidatePaths = getSubAgentDocsCandidatePaths(currentWorkingDirectory);

      let docsInjected = false;
      for (const subAgentDocsPath of candidatePaths) {
          try {
              const docsContent = await readFile(subAgentDocsPath, "utf-8");
              if (docsContent && docsContent.trim().length > 0) {
                  currentContent += `\n\n---\n\n${docsContent}\n\n---\n\n`;
                  ctl.debug(`subagent_docs.md injected into context from: ${subAgentDocsPath}`);
                  state.subAgentDocsInjected = true;
                  await savePersistedState(state);
                  docsInjected = true;
                  break;
              }
          } catch (e) {
              // Keep trying fallback paths.
          }
      }

      if (!docsInjected) {
          ctl.debug("subagent_docs.md not found or failed to load from plugin/workspace paths. Skipping injection.");
      }
  }

  if (isFirstTurn) {
    const simpleSystemPrompt = pluginConfig.get("debug.simpleSystemPrompt");
    const { content: docsBlock, usedLite, categories } = buildFirstTurnDocs(
      userPrompt,
      !!simpleSystemPrompt,
      6000,
    );
    ctl.debug(
      `First-turn tool docs: ${usedLite ? "LITE" : "FILTERED"} ` +
        `(categories: ${categories.join(", ") || "core"})`,
    );
    let injectionContent: string = docsBlock;

    try {
        const { currentWorkingDirectory } = state;
        const candidateStartupPaths = [
            join(currentWorkingDirectory, ".beledarian", "startup.md"),
            join(currentWorkingDirectory, "instructions", "startup.md"),
            join(currentWorkingDirectory, "startup.md"),
        ];

        let startupContent = "";
        let usedStartupPath = "";
        for (const startupPath of candidateStartupPaths) {
            try {
                startupContent = await readFile(startupPath, "utf-8");
                usedStartupPath = dirname(startupPath);
                ctl.debug(`startup.md loaded from: ${startupPath}`);
                break;
            } catch (e) {
                // Keep trying
            }
        }

        if (startupContent) {
            const filesToRead = startupContent.split('\n').map(f => f.trim()).filter(f => f);

            for (const file of filesToRead) {
                const candidateFilePaths = [
                    join(usedStartupPath, file),
                    join(currentWorkingDirectory, file),
                ];

                let loaded = false;
                for (const filePath of candidateFilePaths) {
                    try {
                        const fileContent = await readFile(filePath, "utf-8");
                        if (fileContent.trim().length > 0) {
                            injectionContent = `\n\n---\n\n${fileContent}\n\n---\n\n${injectionContent}`;
                            ctl.debug(`${file} loaded and injected into context from ${filePath}.`);
                            loaded = true;
                            break;
                        }
                    } catch (e) {
                        // Keep trying
                    }
                }
                if (!loaded) {
                    ctl.debug(`Failed to load ${file} from startup.md.`);
                }
            }
        }
    } catch (e) {
        ctl.debug("No startup.md file found or failed to load.");
    }

    // --- Workspace .md auto-intake (the part that was missing) ---
    // Previously ONLY startup.md-listed files were injected and ONLY on first turn.
    // Now: memory.md + common project .md files are read from the LIVE workspace
    // CWD and appended so the model actually sees them.
    try {
        const { currentWorkingDirectory: liveCwd } = state;
        const mdCandidates: string[] = [
            join(liveCwd, "memory.md"),
            join(liveCwd, ".beledarian", "startup.md"),
            join(liveCwd, "instructions", "startup.md"),
            join(liveCwd, "startup.md"),
            join(liveCwd, "AGENTS.md"),
            join(liveCwd, "MUSE.md"),
            join(liveCwd, "GEMINI.md"),
            join(liveCwd, "timbers-toolbox.md"),
        ];
        // Also pick up any other small top-level .md files (project notes, specs).
        try {
            const topEntries = await readdir(liveCwd, { withFileTypes: true });
            for (const e of topEntries) {
                if (!e.isFile()) continue;
                if (extname(e.name).toLowerCase() !== ".md") continue;
                const full = join(liveCwd, e.name);
                if (mdCandidates.includes(full)) continue;
                try {
                    const s = await stat(full);
                    // Keep intake bounded: skip huge docs here (RAG/file tools cover those).
                    if (s.size > 64 * 1024) {
                        ctl.debug(`Skipping large .md for auto-intake (>64KB): ${e.name}`);
                        continue;
                    }
                    mdCandidates.push(full);
                } catch { /* ignore */ }
            }
        } catch (e) {
            ctl.debug("Workspace .md scan failed.", e);
        }
        const seen = new Set<string>();
        for (const mdPath of mdCandidates) {
            if (seen.has(mdPath)) continue;
            seen.add(mdPath);
            try {
                const mdContent = await readFile(mdPath, "utf-8");
                if (mdContent.trim().length === 0) continue;
                // Cap per-file so one big memory.md can't blow the context.
                const capped = mdContent.length > 12000
                    ? mdContent.slice(0, 12000) + "\n\n[... truncated for context ...]"
                    : mdContent;
                injectionContent = `\n\n---\n\n## Workspace file: ${mdPath}\n\n${capped}\n\n---\n\n${injectionContent}`;
                ctl.debug(`Workspace .md injected: ${mdPath} (${mdContent.length} chars)`);
            } catch {
                // Missing file — expected for most candidates.
            }
        }
    } catch (e) {
        ctl.debug("Workspace .md auto-intake failed.", e);
    }

    currentContent = `${injectionContent}\n\n---\n\n${currentContent}`;
  }

  onPredictionStart();
  const statusBlock = injectPluginStatus("");

  if (currentContent !== userPrompt) {
      return `${currentContent}\n\n${statusBlock}`;
  }

  try {
    state.messageCount++;
    await savePersistedState(state);
  } catch (e) {
    ctl.debug("Failed to update message count or memory.", e);
  }
  
  return `${userMessage.getText()}\n\n${statusBlock}`;
}

async function prepareRetrievalResultsContextInjection(
  ctl: PromptPreprocessorController,
  originalUserPrompt: string,
  files: Array<FileHandle>,
): Promise<string> {
  const pluginConfig = ctl.getPluginConfig(pluginConfigSchematics);
  const retrievalLimit = pluginConfig.get("retrieval.limit");
  const retrievalAffinityThreshold = pluginConfig.get("retrieval.threshold");

  const statusSteps = new Map<FileHandle, PredictionProcessStatusController>();

  const rtRetrieve = getDict(
    ctl.getPluginConfig(pluginConfigSchematics).get("general.messageLanguage")
  ).runtime;

  const retrievingStatus = ctl.createStatus({
    status: "loading",
    text: rtRetrieve.statusLoadingEmbeddingModel,
  });
  
  const customModel = pluginConfig.get("retrieval.embeddingModel");
  const modelName = (customModel && customModel.trim() !== "") ? customModel : "nomic-ai/nomic-embed-text-v1.5-GGUF";

  const model = await ctl.client.embedding.model(modelName, {
    signal: ctl.abortSignal,
  });
  
  retrievingStatus.setState({
    status: "loading",
    text: rtRetrieve.statusRetrievingCitations,
  });
  const result = await ctl.client.files.retrieve(originalUserPrompt, files, {
    embeddingModel: model,
    limit: retrievalLimit,
    signal: ctl.abortSignal,
    onFileProcessList(filesToProcess) {
      for (const file of filesToProcess) {
        statusSteps.set(
          file,
          retrievingStatus.addSubStatus({
            status: "waiting",
            text: `Process ${file.name} for retrieval`,
          }),
        );
      }
    },
    onFileProcessingStart(file) {
      statusSteps
        .get(file)!
        .setState({ status: "loading", text: `Processing ${file.name} for retrieval` });
    },
    onFileProcessingEnd(file) {
      statusSteps
        .get(file)!
        .setState({ status: "done", text: `Processed ${file.name} for retrieval` });
    },
    onFileProcessingStepProgress(file, step, progressInStep) {
      const verb = step === "loading" ? "Loading" : step === "chunking" ? "Chunking" : "Embedding";
      statusSteps.get(file)!.setState({
        status: "loading",
        text: `${verb} ${file.name} for retrieval (${(progressInStep * 100).toFixed(1)}%)`,
      });
    },
  });

  result.entries = result.entries.filter(entry => entry.score > retrievalAffinityThreshold);

  let processedContent = "";
  const numRetrievals = result.entries.length;
  if (numRetrievals > 0) {
    retrievingStatus.setState({
      status: "done",
      text: rtRetrieve.statusRetrievedCitations(numRetrievals),
    });
    ctl.debug("Retrieval results", result);
    processedContent += rtRetrieve.citationPrefix;
    let citationNumber = 1;
    result.entries.forEach(result => {
      const completeText = result.content;
      processedContent += rtRetrieve.citationEntry(citationNumber, completeText);
      citationNumber++;
    });
    await ctl.addCitations(result);
    processedContent += rtRetrieve.citationSuffix(originalUserPrompt);
  } else {
    retrievingStatus.setState({
      status: "canceled",
      text: rtRetrieve.statusNoRelevantCitations,
    });
    ctl.debug("No relevant citations found for user query");
    processedContent =
      rtRetrieve.noRelevantCitationsNote + `\n\nUser Query:\n\n${originalUserPrompt}`;
  }
  ctl.debug("Processed content", processedContent);

  return processedContent;
}

async function prepareDocumentContextInjection(
  ctl: PromptPreprocessorController,
  input: ChatMessage,
): Promise<ChatMessage> {
  const documentInjectionSnippets: Map<FileHandle, string> = new Map();
  const files = input.consumeFiles(ctl.client, file => file.type !== "image");
  for (const file of files) {
    const { content } = await ctl.client.files.parseDocument(file, {
      signal: ctl.abortSignal,
    });

    ctl.debug(text`
      Strategy: inject-full-content. Injecting full content of file '${file}' into the
      context. Length: ${content.length}.
    `);
    documentInjectionSnippets.set(file, content);
  }

  let formattedFinalUserPrompt = "";

  const rtDoc = getDict(
    ctl.getPluginConfig(pluginConfigSchematics).get("general.messageLanguage")
  ).runtime;

  if (documentInjectionSnippets.size > 0) {
    formattedFinalUserPrompt += rtDoc.documentInjectionHeader;

    for (const [fileHandle, snippet] of documentInjectionSnippets) {
      formattedFinalUserPrompt += rtDoc.documentInjectionFileBlock(fileHandle.name, snippet);
    }

    formattedFinalUserPrompt += rtDoc.documentInjectionSuffix(input.getText());
  }

  input.replaceText(formattedFinalUserPrompt);
  return input;
}

async function measureContextWindow(ctx: Chat, model: LLMDynamicHandle) {
  let currentContextFormatted: string;
  
  try {
    const result = await Promise.resolve(model.applyPromptTemplate(ctx));
    currentContextFormatted = result;
  } catch (err) {
    console.warn("[PromptPreprocessor] applyPromptTemplate failed, using fallback:", err instanceof Error ? err.message : String(err));
    
    let combinedText = "";
    try {
      const chatObj = ctx as any;
      if (chatObj.messages && Array.isArray(chatObj.messages)) {
        for (const msg of chatObj.messages) {
          combinedText += (msg.getText ? msg.getText() : (msg.text || "")) + "\n";
        }
      } else if (chatObj.pullHistory) {
        const messages = await chatObj.pullHistory();
        if (Array.isArray(messages)) {
          for (const msg of messages) {
            combinedText += (msg.getText ? msg.getText() : (msg.text || "")) + "\n";
          }
        }
      }
    } catch (innerErr) {
      console.warn("[PromptPreprocessor] Failed to extract messages from Chat object:", innerErr instanceof Error ? innerErr.message : String(innerErr));
      combinedText = "";
    }
    currentContextFormatted = combinedText;
  }
  
  const totalTokensInContext = await model.countTokens(currentContextFormatted);
  const modelContextLength = await model.getContextLength();
  const modelRemainingContextLength = modelContextLength - totalTokensInContext;
  const contextOccupiedPercent = (totalTokensInContext / modelContextLength) * 100;
  return {
    totalTokensInContext,
    modelContextLength,
    modelRemainingContextLength,
    contextOccupiedPercent,
  };
}

async function chooseContextInjectionStrategy(
  ctl: PromptPreprocessorController,
  originalUserPrompt: string,
  files: Array<FileHandle>,
): Promise<DocumentContextInjectionStrategy> {
  const rtStrategy = getDict(
    ctl.getPluginConfig(pluginConfigSchematics).get("general.messageLanguage")
  ).runtime;

  const status = ctl.createStatus({
    status: "loading",
    text: rtStrategy.statusDecidingStrategy,
  });

  const model = await ctl.client.llm.model();
  const ctx = await ctl.pullHistory();

  let contextMeasurement: { totalTokensInContext: number; modelContextLength: number; modelRemainingContextLength: number; contextOccupiedPercent: number };
  try {
    contextMeasurement = await measureContextWindow(ctx, model);
  } catch (err) {
    console.warn("[PromptPreprocessor] measureContextWindow failed:", err instanceof Error ? err.message : String(err));
    status.setState({ status: "canceled", text: "Could not measure context; using retrieval" });
    return "retrieval";
  }

  // Estimate document token cost roughly
  let estimatedDocTokens = 0;
  for (const file of files) {
    try {
      const { content } = await ctl.client.files.parseDocument(file, { signal: ctl.abortSignal });
      estimatedDocTokens += Math.ceil(content.length / 4);
    } catch {
      estimatedDocTokens += 4000;
    }
  }

  if (estimatedDocTokens < contextMeasurement.modelRemainingContextLength * 0.5) {
    status.setState({ status: "done", text: rtStrategy.statusInjectFullContent || "Injecting full document content" });
    return "inject-full-content";
  }

  status.setState({ status: "done", text: rtStrategy.statusUseRetrieval || "Using retrieval" });
  return "retrieval";
}