import type { LlamaModel } from "node-llama-cpp";
import { LlmSession } from "./llm";
import { xmlEscape } from "./prompt";
import { isPoliteEnough } from "./util";

export const politenessModes = ["auto", "always", "never"] as const;
export type PolitenessMode = (typeof politenessModes)[number];

export const isPolitenessMode = (value: string): value is PolitenessMode =>
  politenessModes.some((mode) => mode === value);

const rephrasePrompt = `あなたは日本語の文体変換器です。
<original> 内の文章を丁寧な「です・ます」調に直してください。
- 発想、意味、語彙、比喩、批評の強さ、文の数を保ちます。
- 主に文末と、それに必要な最小限の助詞・活用だけを変えます。
- 内容の追加、削除、要約、説明、感想、自己紹介をしません。
- 変換後の本文だけを出力します。`;

const stripPrefix = (raw: string, prefix: string) =>
  (raw.startsWith(prefix) ? raw.slice(prefix.length) : raw).trim();

export async function ensurePolite(
  model: LlamaModel,
  text: string,
  mode: PolitenessMode = "auto",
  stream = true,
) {
  if (mode === "never" || (mode === "auto" && isPoliteEnough(text))) {
    return { text, changed: false, attempts: 0 } as const;
  }

  await using session = new LlmSession(model, rephrasePrompt);
  await session.init();
  const prefix = "変換後:\n";
  const raw = await session.prompt(`<original>${xmlEscape(text)}</original>`, {
    responsePrefix: prefix,
    customStopTriggers: [
      "\n\n",
      "\n\n説明",
      "\n説明",
      "\n\n変更点",
      "\n変更点",
      "\n\n* **",
    ],
    maxTokens: 256,
    temperature: 0.1,
    minP: 0.05,
    repeatPenalty: { lastTokens: 128, penalty: 1.05 },
    ...(stream ? {} : { onResponseChunk() {} }),
  });
  if (stream) process.stderr.write("\n");
  const converted = stripPrefix(raw, prefix) || text;

  return { text: converted, changed: converted !== text, attempts: 1 } as const;
}