import { type Chat, type GeneratorController, LMStudioClient } from "@lmstudio/sdk";
import { configSchematics, globalConfigSchematics } from "./config";
import { StreamRepetitionDetector, type DetectorConfig } from "./detector";

/**
 * Guarded generator.
 *
 * Wraps a local LM Studio model the same way the official remote-lmstudio
 * plugin does, but inspects every fragment (thinking + answer) on the fly. When
 * the stream degenerates into a repetition loop, the underlying prediction is
 * aborted immediately — already-generated text is preserved and committed as a
 * normal (stopped) turn, exactly like a user stop. Nothing is deleted.
 */
export async function generate(ctl: GeneratorController, chat: Chat) {
  const globalPluginConfig = ctl.getGlobalPluginConfig(globalConfigSchematics);
  const specifiedClientIdentifier = globalPluginConfig.get("clientIdentifier");
  const specifiedClientPasskey = globalPluginConfig.get("clientPasskey");

  await using client = new LMStudioClient({
    baseUrl: globalPluginConfig.get("baseUrl"),
    clientIdentifier: specifiedClientIdentifier === "" ? undefined : specifiedClientIdentifier,
    clientPasskey: specifiedClientPasskey === "" ? undefined : specifiedClientPasskey,
  });

  const pluginConfig = ctl.getPluginConfig(configSchematics);
  const modelKey = pluginConfig.get("modelKey");
  const temperature = pluginConfig.get("temperature");

  const detectorCfg: Partial<DetectorConfig> = {
    threshold: pluginConfig.get("threshold"),
    chantThreshold: pluginConfig.get("chantThreshold"),
  };
  const contentDetector = new StreamRepetitionDetector(detectorCfg);
  const reasoningDetector = new StreamRepetitionDetector(detectorCfg);
  const toolArgDetector = new StreamRepetitionDetector(detectorCfg);

  const monitorContent = pluginConfig.get("monitorContent");
  const monitorReasoning = pluginConfig.get("monitorReasoning");
  const notifyOnStop = pluginConfig.get("notifyOnStop");

  // Own abort controller so the guard can stop the prediction mid-stream.
  const guard = new AbortController();
  let stopReason: string | null = null;

  const combinedSignal = AbortSignal.any([ctl.abortSignal, guard.signal]);

  const model = await client.llm.model(modelKey, { signal: ctl.abortSignal });

  const checkAndStop = (
    detector: StreamRepetitionDetector,
    text: string,
    streamLabel: string,
  ): boolean => {
    const hit = detector.push(text);
    if (hit === null) return false;
    stopReason =
      `repetition-guard: stopped a repetition loop in the ${streamLabel} stream ` +
      `(${hit.kind}: unit ${JSON.stringify(hit.unit)} × ${hit.count}, threshold ${hit.threshold}). ` +
      `Kept everything generated before the loop.`;
    console.warn(stopReason);
    guard.abort();
    return true;
  };

  try {
    await model.respond(chat, {
      rawTools: {
        type: "toolArray",
        tools: ctl.getToolDefinitions(),
      },
      temperature,
      onPredictionFragment: fragment => {
        if (stopReason !== null) return;
        const isReasoning =
          (fragment as { reasoningType?: string | null }).reasoningType != null &&
          (fragment as { reasoningType?: string | null }).reasoningType !== "none";
        if (isReasoning ? monitorReasoning : monitorContent) {
          const detector = isReasoning ? reasoningDetector : contentDetector;
          if (checkAndStop(detector, fragment.content, isReasoning ? "thinking" : "answer")) {
            return;
          }
        }
        ctl.fragmentGenerated(fragment.content, fragment);
      },
      onToolCallRequestStart: () => ctl.toolCallGenerationStarted(),
      onToolCallRequestNameReceived: (_callId, name) => ctl.toolCallGenerationNameReceived(name),
      onToolCallRequestArgumentFragmentGenerated: (_callId, content) => {
        if (stopReason !== null) return;
        // Tool-call arguments can loop too (e.g. a huge repeated JSON blob).
        if (checkAndStop(toolArgDetector, content, "tool-call argument")) {
          ctl.toolCallGenerationFailed(new Error(stopReason ?? "repetition-guard"));
          return;
        }
        ctl.toolCallGenerationArgumentFragmentGenerated(content);
      },
      onToolCallRequestEnd: (_callId, info) => ctl.toolCallGenerationEnded(info.toolCallRequest),
      onToolCallRequestFailure: (_callId, error) => ctl.toolCallGenerationFailed(error),
      signal: combinedSignal,
    });
  } catch (error) {
    if (stopReason === null) {
      // Not the guard: user abort or a real failure — propagate as-is.
      throw error;
    }
    // Guard-triggered abort: end the turn cleanly with what was generated.
  }

  if (stopReason !== null && notifyOnStop) {
    ctl.fragmentGenerated(
      `\n\n> [repetition-guard] ✂️ Detected a repetition loop and stopped generation. ` +
        `Everything above was kept as-is. Try raising Temperature or lowering Repeat Penalty in the sidebar, ` +
        `or lower the threshold in plugin settings to cut earlier.`,
    );
  }
}

import { type Chat, type GeneratorController, LMStudioClient } from "@lmstudio/sdk";
import { configSchematics, globalConfigSchematics } from "./config";
import { StreamRepetitionDetector, type DetectorConfig } from "./detector";

/**
 * Guarded generator.
 *
 * Wraps a local LM Studio model the same way the official remote-lmstudio
 * plugin does, but inspects every fragment (thinking + answer) on the fly. When
 * the stream degenerates into a repetition loop, the underlying prediction is
 * aborted immediately — already-generated text is preserved and committed as a
 * normal (stopped) turn, exactly like a user stop. Nothing is deleted.
 */
export async function generate(ctl: GeneratorController, chat: Chat) {
  const globalPluginConfig = ctl.getGlobalPluginConfig(globalConfigSchematics);
  const specifiedClientIdentifier = globalPluginConfig.get("clientIdentifier");
  const specifiedClientPasskey = globalPluginConfig.get("clientPasskey");

  await using client = new LMStudioClient({
    baseUrl: globalPluginConfig.get("baseUrl"),
    clientIdentifier: specifiedClientIdentifier === "" ? undefined : specifiedClientIdentifier,
    clientPasskey: specifiedClientPasskey === "" ? undefined : specifiedClientPasskey,
  });

  const pluginConfig = ctl.getPluginConfig(configSchematics);
  const modelKey = pluginConfig.get("modelKey");
  const temperature = pluginConfig.get("temperature");

  const detectorCfg: Partial<DetectorConfig> = {
    threshold: pluginConfig.get("threshold"),
    chantThreshold: pluginConfig.get("chantThreshold"),
  };
  const contentDetector = new StreamRepetitionDetector(detectorCfg);
  const reasoningDetector = new StreamRepetitionDetector(detectorCfg);
  const toolArgDetector = new StreamRepetitionDetector(detectorCfg);

  const monitorContent = pluginConfig.get("monitorContent");
  const monitorReasoning = pluginConfig.get("monitorReasoning");
  const notifyOnStop = pluginConfig.get("notifyOnStop");

  // Own abort controller so the guard can stop the prediction mid-stream.
  const guard = new AbortController();
  let stopReason: string | null = null;

  const combinedSignal = AbortSignal.any([ctl.abortSignal, guard.signal]);

  const model = await client.llm.model(modelKey, { signal: ctl.abortSignal });

  const checkAndStop = (
    detector: StreamRepetitionDetector,
    text: string,
    streamLabel: string,
  ): boolean => {
    const hit = detector.push(text);
    if (hit === null) return false;
    stopReason =
      `repetition-guard: stopped a repetition loop in the ${streamLabel} stream ` +
      `(${hit.kind}: unit ${JSON.stringify(hit.unit)} × ${hit.count}, threshold ${hit.threshold}). ` +
      `Kept everything generated before the loop.`;
    console.warn(stopReason);
    guard.abort();
    return true;
  };

  try {
    await model.respond(chat, {
      rawTools: {
        type: "toolArray",
        tools: ctl.getToolDefinitions(),
      },
      temperature,
      onPredictionFragment: fragment => {
        if (stopReason !== null) return;
        const isReasoning =
          (fragment as { reasoningType?: string | null }).reasoningType != null &&
          (fragment as { reasoningType?: string | null }).reasoningType !== "none";
        if (isReasoning ? monitorReasoning : monitorContent) {
          const detector = isReasoning ? reasoningDetector : contentDetector;
          if (checkAndStop(detector, fragment.content, isReasoning ? "thinking" : "answer")) {
            return;
          }
        }
        ctl.fragmentGenerated(fragment.content, fragment);
      },
      onToolCallRequestStart: () => ctl.toolCallGenerationStarted(),
      onToolCallRequestNameReceived: (_callId, name) => ctl.toolCallGenerationNameReceived(name),
      onToolCallRequestArgumentFragmentGenerated: (_callId, content) => {
        if (stopReason !== null) return;
        // Tool-call arguments can loop too (e.g. a huge repeated JSON blob).
        if (checkAndStop(toolArgDetector, content, "tool-call argument")) {
          ctl.toolCallGenerationFailed(new Error(stopReason ?? "repetition-guard"));
          return;
        }
        ctl.toolCallGenerationArgumentFragmentGenerated(content);
      },
      onToolCallRequestEnd: (_callId, info) => ctl.toolCallGenerationEnded(info.toolCallRequest),
      onToolCallRequestFailure: (_callId, error) => ctl.toolCallGenerationFailed(error),
      signal: combinedSignal,
    });
  } catch (error) {
    if (stopReason === null) {
      // Not the guard: user abort or a real failure — propagate as-is.
      throw error;
    }
    // Guard-triggered abort: end the turn cleanly with what was generated.
  }

  if (stopReason !== null && notifyOnStop) {
    ctl.fragmentGenerated(
      `\n\n> [repetition-guard] ✂️ Detected a repetition loop and stopped generation. ` +
        `Everything above was kept as-is. Try raising Temperature or lowering Repeat Penalty in the sidebar, ` +
        `or lower the threshold in plugin settings to cut earlier.`,
    );
  }
}