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

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 remoteClient = 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 model = await remoteClient.llm.model(modelKey, { signal: ctl.abortSignal });
  const toolDefinitions = ctl.getToolDefinitions();

  await model.respond(chat, {
    rawTools: toolDefinitions.length === 0 ? undefined : {
      type: "toolArray",
      tools: toolDefinitions,
    },
    temperature,
    onPredictionFragment: fragment => ctl.fragmentGenerated(fragment.content, fragment),
    onToolCallRequestStart: () => ctl.toolCallGenerationStarted(),
    onToolCallRequestNameReceived: name => ctl.toolCallGenerationNameReceived(name),
    onToolCallRequestArgumentFragmentGenerated: (content) => ctl.toolCallGenerationArgumentFragmentGenerated(content),
    onToolCallRequestEnd: info => ctl.toolCallGenerationEnded(info.toolCallRequest),
    onToolCallRequestFailure: error => ctl.toolCallGenerationFailed(error),
    signal: ctl.abortSignal,
  });
}