import type { LLamaChatPromptOptions, LlamaGrammar } from "node-llama-cpp";
import { BOT_NAME } from "./prompt";

export const harnessNames = ["json", "speaker", "tagged"] as const;
export type HarnessName = (typeof harnessNames)[number];

type HarnessOptions = Pick<
  LLamaChatPromptOptions,
  "grammar" | "responsePrefix" | "customStopTriggers"
>;

export type OutputHarness = {
  name: HarnessName;
  promptSuffix: string;
  options: HarnessOptions;
  parse: (raw: string) => string | null;
  formatHistory: (text: string) => string;
};

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

export function createHarness(
  name: HarnessName,
  jsonGrammar: LlamaGrammar,
): OutputHarness {
  switch (name) {
    case "json":
      return {
        name,
        promptSuffix:
          '\n出力は {"name":"あるびのちゃん","text":"投稿本文"} という JSON オブジェクトだけにします。',
        options: { grammar: jsonGrammar },
        parse(raw) {
          try {
            const value: unknown = JSON.parse(raw.trim());
            if (
              typeof value === "object" &&
              value !== null &&
              "text" in value &&
              typeof value.text === "string"
            ) {
              return value.text.trim() || null;
            }
          } catch {
            // Reported as a harness failure by the caller.
          }
          return null;
        },
        formatHistory: (text) => JSON.stringify({ name: BOT_NAME, text }),
      };
    case "speaker": {
      const prefix = `${BOT_NAME}:\n`;
      return {
        name,
        promptSuffix:
          "\nアプリが書き手名を補うので、あなたは投稿本文だけを続けます。末尾に完了報告などを付けません。",
        options: {
          responsePrefix: prefix,
          // This was the separator used by the pre-JSON harness. Retaining it
          // also prevents the model from starting a fabricated next speaker.
          customStopTriggers: [
            "\n----------",
            `\n${BOT_NAME}:`,
            "\n<post",
            "\n</post>",
          ],
        },
        parse(raw) {
          const [firstTurn = ""] = stripPrefix(raw.trim(), prefix).split(
            `\n${BOT_NAME}:`,
          );
          const body = firstTurn
            .replace(/\n?<\/?post[^>]*>?[\s\S]*$/u, "")
            .trim();
          return body || null;
        },
        formatHistory: (text) => `${prefix}${text}`,
      };
    }
    case "tagged": {
      const prefix = `<post author="${BOT_NAME}">\n`;
      return {
        name,
        promptSuffix:
          "\n<post> はアプリが開始します。投稿本文だけを書き、直後に </post> で閉じます。",
        options: {
          responsePrefix: prefix,
          customStopTriggers: ["</post>"],
        },
        parse(raw) {
          const body = stripPrefix(raw.trim(), prefix)
            .replace(/<\/post>[\s\S]*$/u, "")
            .replace(/(?:<|&lt;)post(?:\s[^>]*?)?(?:>|&gt;)[\s\S]*$/u, "")
            .trim();
          return body || null;
        },
        formatHistory: (text) => `${prefix}${text}\n</post>`,
      };
    }
  }
}

export function isHarnessName(value: string): value is HarnessName {
  return harnessNames.some((name) => name === value);
}