// Chat titles.
//
// LM Studio leaves naming a chat to the plugin that runs the prediction loop, so while this
// plugin is enabled, chats were never named. After each reply, if the chat still needs a name,
// the model is asked for a short title. The answer is constrained to a small JSON object, which
// keeps reasoning models from spending their budget thinking first. If that fails, the title is
// the start of the user's first message.

import { Chat, type PredictionLoopHandlerController } from "@lmstudio/sdk";
import { stripReasoningMarkers } from "./reasoning";

const TITLE_SCHEMA = {
  type: "object",
  properties: { title: { type: "string", minLength: 1, maxLength: 60 } },
  required: ["title"],
  additionalProperties: false,
};

const MAX_TITLE = 60;

// The user's own words in a stored message, without text other plugins wrapped around them:
// rag-v1's injected documents ("User Query:") and throughLine's state blocks.
export function ownWords(stored: string): string {
  let t = stored;
  const q = t.lastIndexOf("User Query:");
  if (q >= 0) t = t.slice(q + "User Query:".length);
  return t
    .replace(/\[throughLine[^\]]*\][\s\S]*?\[end throughLine\]/g, "")
    .replace(/\[throughLine:[^\]]*\]/g, "")
    .trim();
}

export function firstUserText(history: Chat): string {
  for (const m of history.getMessagesArray()) {
    if (m.getRole() !== "user") continue;
    const t = ownWords(m.getText());
    if (t && !t.startsWith("/")) return t;
  }
  return "";
}

export function cleanTitle(raw: string): string {
  let t = raw.split("\n")[0].trim();
  t = t.replace(/^(title\s*:\s*)/i, "").replace(/^["'“”‘’`*_#\s]+|["'“”‘’`*_\s]+$/g, "").replace(/[.\s]+$/, "");
  if (t.length > MAX_TITLE) t = t.slice(0, MAX_TITLE).replace(/\s+\S*$/, "");
  return t.trim();
}

// The first few words of the user's message, used when the model gives no usable title.
export function fallbackTitle(text: string): string {
  const words = text.replace(/\s+/g, " ").trim().split(" ").slice(0, 6).join(" ");
  const t = cleanTitle(words.replace(/[?!.,;:]+$/, ""));
  return t ? t.charAt(0).toUpperCase() + t.slice(1) : "New chat";
}

export async function nameChatIfNeeded(ctl: PredictionLoopHandlerController): Promise<void> {
  if (ctl.abortSignal.aborted) return;
  try {
    if (!(await ctl.needsNaming())) return;
  } catch {
    return;
  }
  const first = firstUserText(await ctl.pullHistory());
  if (!first) return;

  let title = "";
  try {
    const source = await ctl.tokenSource();
    const prediction = source.respond(
      Chat.from([
        {
          role: "user",
          content:
            "Give this conversation a short title of two to six words, based on its first message. " +
            "Answer only with the JSON object.\n\nFirst message:\n" +
            first.slice(0, 1200),
        },
      ]),
      {
        maxTokens: 60,
        temperature: 0.3,
        structured: { type: "json", jsonSchema: TITLE_SCHEMA },
        contextOverflowPolicy: "stopAtLimit",
      },
    );
    ctl.onAborted(() => void prediction.cancel());
    const result = await prediction;
    const text = stripReasoningMarkers(result.nonReasoningContent ?? result.content);
    title = cleanTitle(String((JSON.parse(text) as { title?: unknown }).title ?? ""));
  } catch (error) {
    ctl.debug("chat title: model title unavailable, using the first message", error);
  }
  if (!title) title = fallbackTitle(first);
  try {
    await ctl.suggestName(title);
  } catch (error) {
    ctl.debug("chat title: could not set the name", error);
  }
}

// Chat titles.
//
// LM Studio leaves naming a chat to the plugin that runs the prediction loop, so while this
// plugin is enabled, chats were never named. After each reply, if the chat still needs a name,
// the model is asked for a short title. The answer is constrained to a small JSON object, which
// keeps reasoning models from spending their budget thinking first. If that fails, the title is
// the start of the user's first message.

import { Chat, type PredictionLoopHandlerController } from "@lmstudio/sdk";
import { stripReasoningMarkers } from "./reasoning";

const TITLE_SCHEMA = {
  type: "object",
  properties: { title: { type: "string", minLength: 1, maxLength: 60 } },
  required: ["title"],
  additionalProperties: false,
};

const MAX_TITLE = 60;

// The user's own words in a stored message, without text other plugins wrapped around them:
// rag-v1's injected documents ("User Query:") and throughLine's state blocks.
export function ownWords(stored: string): string {
  let t = stored;
  const q = t.lastIndexOf("User Query:");
  if (q >= 0) t = t.slice(q + "User Query:".length);
  return t
    .replace(/\[throughLine[^\]]*\][\s\S]*?\[end throughLine\]/g, "")
    .replace(/\[throughLine:[^\]]*\]/g, "")
    .trim();
}

export function firstUserText(history: Chat): string {
  for (const m of history.getMessagesArray()) {
    if (m.getRole() !== "user") continue;
    const t = ownWords(m.getText());
    if (t && !t.startsWith("/")) return t;
  }
  return "";
}

export function cleanTitle(raw: string): string {
  let t = raw.split("\n")[0].trim();
  t = t.replace(/^(title\s*:\s*)/i, "").replace(/^["'“”‘’`*_#\s]+|["'“”‘’`*_\s]+$/g, "").replace(/[.\s]+$/, "");
  if (t.length > MAX_TITLE) t = t.slice(0, MAX_TITLE).replace(/\s+\S*$/, "");
  return t.trim();
}

// The first few words of the user's message, used when the model gives no usable title.
export function fallbackTitle(text: string): string {
  const words = text.replace(/\s+/g, " ").trim().split(" ").slice(0, 6).join(" ");
  const t = cleanTitle(words.replace(/[?!.,;:]+$/, ""));
  return t ? t.charAt(0).toUpperCase() + t.slice(1) : "New chat";
}

export async function nameChatIfNeeded(ctl: PredictionLoopHandlerController): Promise<void> {
  if (ctl.abortSignal.aborted) return;
  try {
    if (!(await ctl.needsNaming())) return;
  } catch {
    return;
  }
  const first = firstUserText(await ctl.pullHistory());
  if (!first) return;

  let title = "";
  try {
    const source = await ctl.tokenSource();
    const prediction = source.respond(
      Chat.from([
        {
          role: "user",
          content:
            "Give this conversation a short title of two to six words, based on its first message. " +
            "Answer only with the JSON object.\n\nFirst message:\n" +
            first.slice(0, 1200),
        },
      ]),
      {
        maxTokens: 60,
        temperature: 0.3,
        structured: { type: "json", jsonSchema: TITLE_SCHEMA },
        contextOverflowPolicy: "stopAtLimit",
      },
    );
    ctl.onAborted(() => void prediction.cancel());
    const result = await prediction;
    const text = stripReasoningMarkers(result.nonReasoningContent ?? result.content);
    title = cleanTitle(String((JSON.parse(text) as { title?: unknown }).title ?? ""));
  } catch (error) {
    ctl.debug("chat title: model title unavailable, using the first message", error);
  }
  if (!title) title = fallbackTitle(first);
  try {
    await ctl.suggestName(title);
  } catch (error) {
    ctl.debug("chat title: could not set the name", error);
  }
}