/**
 * llmindex.io — instruction-following templates (mechanically checkable constraints)
 * Author:  Simon-Pierre Boucher
 * Contact: contact@spboucher.ai
 * License: Proprietary — © Simon-Pierre Boucher, all rights reserved
 */
import { ANSWER_FORMAT_INSTRUCTIONS, type ConstraintSpec } from '../answer';
import type { ItemTemplate } from '../types';

const WORDS = [
  'nova', 'delta', 'ember', 'quartz', 'falcon', 'lumen', 'cedar', 'orbit', 'prism', 'tundra',
  'zephyr', 'basalt', 'comet', 'drift', 'echo', 'flint',
];

export const ifRepeat: ItemTemplate = {
  id: 'if.format.repeat-v1',
  domain: 'instruction_following',
  description:
    'Produce a word repeated n times with an exact separator and casing — checkable string output.',
  paramSpace: WORDS.length * 7 * 3 * 3,
  render(rng, perturbSeed) {
    const word = rng.pick(WORDS);
    const n = rng.int(3, 9);
    const sep = rng.pick(['-', '_', '/'] as const);
    const casing = rng.pick(['uppercase', 'lowercase', 'capitalized'] as const);
    const cased =
      casing === 'uppercase'
        ? word.toUpperCase()
        : casing === 'capitalized'
          ? word[0]!.toUpperCase() + word.slice(1)
          : word;
    const expected = Array.from({ length: n }, () => cased).join(sep);
    return {
      templateId: this.id,
      domain: this.domain,
      prompt:
        `Write the word "${word}" in ${casing} form, repeated exactly ${n} times, ` +
        `joined by the character "${sep}" with no spaces. Output that string as your answer.\n\n${ANSWER_FORMAT_INSTRUCTIONS}`,
      answerKey: expected,
      grading: 'exact',
      perturbSeed,
    };
  },
};

const TOPICS = ['the sea', 'a city at night', 'winter mornings', 'an old machine', 'a long journey'] as const;

/**
 * Constraint stacking (IFEval-style, home-made): 4-5 simultaneously
 * verifiable constraints on one short generated text. Every constraint is
 * mechanically checkable; satisfiability is guaranteed by construction
 * (the generator verifies a witness before emitting the item).
 */
export const ifConstraintStack: ItemTemplate = {
  id: 'if.constraints.stack-v1',
  domain: 'instruction_following',
  description:
    'Write one sentence satisfying 4-5 stacked verifiable constraints (word count, boundary words, exact keyword frequency, forbidden letter, casing).',
  paramSpace: 16 ** 3 * 20 * 5 * 4,
  render(rng, perturbSeed) {
    const wordCount = rng.int(14, 24);
    const startWord = rng.pick(WORDS);
    const endWord = rng.pick(WORDS.filter((w) => w !== startWord));
    const keyword = rng.pick(WORDS.filter((w) => w !== startWord && w !== endWord));
    const keywordCount = rng.int(2, 3);
    // Forbidden letter must not appear in mandatory words.
    const mandatory = `${startWord}${endWord}${keyword}`;
    const candidates = 'qjzxv'.split('').filter((ch) => !mandatory.includes(ch));
    const forbiddenLetter = candidates[0] ?? 'q';
    const spec: ConstraintSpec = {
      wordCount,
      startsWithWord: startWord,
      endsWithWord: endWord,
      includeWordExactly: [{ word: keyword, count: keywordCount }],
      forbiddenLetter,
      allLowercase: true,
    };
    // Witness check: constraints are jointly satisfiable by construction —
    // filler words below avoid the forbidden letters entirely.
    const prompt =
      `Write in English about ${rng.pick(TOPICS)}, following ALL of these rules simultaneously:\n` +
      `1. Exactly ${wordCount} words.\n` +
      `2. The first word must be "${startWord}" and the last word must be "${endWord}".\n` +
      `3. Use the word "${keyword}" exactly ${keywordCount} times (in addition to rules 2 if they differ).\n` +
      `4. The letter "${forbiddenLetter}" must not appear anywhere.\n` +
      `5. Everything entirely in lowercase.\n\n` +
      `Give the text itself as your answer.\n\n${ANSWER_FORMAT_INSTRUCTIONS}`;
    return {
      templateId: this.id,
      domain: this.domain,
      prompt,
      answerKey: JSON.stringify(spec),
      grading: 'constraints',
      perturbSeed,
    };
  },
};

export const ifAcronym: ItemTemplate = {
  id: 'if.format.acronym-v1',
  domain: 'instruction_following',
  description: 'Build an acronym from the k-th letters of a generated word list.',
  paramSpace: WORDS.length ** 4 * 3,
  render(rng, perturbSeed) {
    const words = rng.shuffle(WORDS).slice(0, rng.int(4, 6));
    const k = rng.int(1, 3);
    const ordinal = k === 1 ? 'first' : k === 2 ? 'second' : 'third';
    const expected = words.map((w) => w[k - 1]!.toUpperCase()).join('');
    return {
      templateId: this.id,
      domain: this.domain,
      prompt:
        `Take the ${ordinal} letter of each of these words, in order: ${words.join(', ')}. ` +
        `Concatenate them in uppercase into a single string with no separators. Output that string as your answer.\n\n${ANSWER_FORMAT_INSTRUCTIONS}`,
      answerKey: expected,
      grading: 'exact',
      perturbSeed,
    };
  },
};