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"),
    checkInterval: pluginConfig.get("checkInterval"),
    maxTextLength: pluginConfig.get("maxTextLength"),
  };
  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]);

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

  let model;
  let activeModelKey = modelKey && String(modelKey).trim();

  if (!activeModelKey) {
    try {
      const loadedModels = await client.llm.listLoaded();
      if (loadedModels && loadedModels.length > 0) {
        // Берем первую попавшуюся модель, которая сейчас загружена в памяти
        activeModelKey = loadedModels[0].modelKey;
      }
    } catch (e) {
      console.error("Failed to fetch automatically loaded model from LM Studio:", e);
    }
  }

  if (activeModelKey) {
    model = await client.llm.model(activeModelKey, { signal: ctl.abortSignal });
  } else {
    throw new Error("No model specified in config and no active models found running in LM Studio.");
  }
  
  //stopReason = null; - NOT here!

  const checkAndStop = (
    detector: StreamRepetitionDetector,
    text: string,
    streamLabel: string,
  ): boolean => {
    // ИСПРАВЛЕНИЕ: Читаем конфигурацию по канонам LM Studio SDK через метод .get()
    // Сначала получаем объект конфига с помощью getPluginConfig
    const pluginConfig = ctl.getPluginConfig(configSchematics);
    
    // Вытаскиваем актуальное значение ползунка по его имени или шортхэнду
    const currentMultiplier = pluginConfig.get("repetitionThreshold") ?? pluginConfig.get("repetitionThreshold") ?? 1.5;

    // Передаем этот живой коэффициент прямо в метод push
    const hit = detector.push(text, currentMultiplier);
    if (hit === null) return false;

    const formatSize = (chars: number) => {
      return chars >= 1024 ? `${(chars / 1024).toFixed(1)} KB` : `${chars} chars`;
    };

    stopReason =
      `repetition-guard: stopped a repetition loop in the answer stream.` +
      `Detector: ${hit.kind}. The text's suffix matches the buffer's prefix for ${hit.matchLen} chars. ` +
      `Repeating pattern ${JSON.stringify(hit.unit)} length: ${hit.threshold} (truncated to 40). Approx repeats ≥ ${hit.count} (threshold ${hit.threshold}).` +
      `Overlap Length: ${formatSize(hit.matchLen)}. Loop distance: ${formatSize(hit.distance)}. Everything generated before the loop was kept in chat.`;

    console.warn(stopReason);
    guard.abort();
    return true;
  };

  stopReason = null; // <-- ОБЯЗАТЕЛЬНЫЙ СБРОС ДЛЯ КАЖДОГО НОВОГО ХОДА!
  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. \n` +
        ` Try raising Temperature in the sidebar, or lower the threshold in plugin settings to cut earlier.` +
        ` Increase Repeat Penalty above 1 by 0.1`,
    );
  }
}

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"),
    checkInterval: pluginConfig.get("checkInterval"),
    maxTextLength: pluginConfig.get("maxTextLength"),
  };
  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]);

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

  let model;
  let activeModelKey = modelKey && String(modelKey).trim();

  if (!activeModelKey) {
    try {
      const loadedModels = await client.llm.listLoaded();
      if (loadedModels && loadedModels.length > 0) {
        // Берем первую попавшуюся модель, которая сейчас загружена в памяти
        activeModelKey = loadedModels[0].modelKey;
      }
    } catch (e) {
      console.error("Failed to fetch automatically loaded model from LM Studio:", e);
    }
  }

  if (activeModelKey) {
    model = await client.llm.model(activeModelKey, { signal: ctl.abortSignal });
  } else {
    throw new Error("No model specified in config and no active models found running in LM Studio.");
  }
  
  //stopReason = null; - NOT here!

  const checkAndStop = (
    detector: StreamRepetitionDetector,
    text: string,
    streamLabel: string,
  ): boolean => {
    // ИСПРАВЛЕНИЕ: Читаем конфигурацию по канонам LM Studio SDK через метод .get()
    // Сначала получаем объект конфига с помощью getPluginConfig
    const pluginConfig = ctl.getPluginConfig(configSchematics);
    
    // Вытаскиваем актуальное значение ползунка по его имени или шортхэнду
    const currentMultiplier = pluginConfig.get("repetitionThreshold") ?? pluginConfig.get("repetitionThreshold") ?? 1.5;

    // Передаем этот живой коэффициент прямо в метод push
    const hit = detector.push(text, currentMultiplier);
    if (hit === null) return false;

    const formatSize = (chars: number) => {
      return chars >= 1024 ? `${(chars / 1024).toFixed(1)} KB` : `${chars} chars`;
    };

    stopReason =
      `repetition-guard: stopped a repetition loop in the answer stream.` +
      `Detector: ${hit.kind}. The text's suffix matches the buffer's prefix for ${hit.matchLen} chars. ` +
      `Repeating pattern ${JSON.stringify(hit.unit)} length: ${hit.threshold} (truncated to 40). Approx repeats ≥ ${hit.count} (threshold ${hit.threshold}).` +
      `Overlap Length: ${formatSize(hit.matchLen)}. Loop distance: ${formatSize(hit.distance)}. Everything generated before the loop was kept in chat.`;

    console.warn(stopReason);
    guard.abort();
    return true;
  };

  stopReason = null; // <-- ОБЯЗАТЕЛЬНЫЙ СБРОС ДЛЯ КАЖДОГО НОВОГО ХОДА!
  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. \n` +
        ` Try raising Temperature in the sidebar, or lower the threshold in plugin settings to cut earlier.` +
        ` Increase Repeat Penalty above 1 by 0.1`,
    );
  }
}