import type { Disturbance, Stats, WorldState } from "./types.ts";
// Tokenling world model: prompt headers, event lines and tool results.
// Mirrors train/world.py exactly; test/parity.mjs checks the two agree.

export const STAGES = ["baby", "kid", "teen", "adult"];
export const PERSONALITIES = [
  "bouncy",
  "shy",
  "curious",
  "sleepy",
  "dramatic",
  "foodie",
  "silly",
  "sweet",
  "grumpy",
  "brave",
  "artsy",
  "dreamy",
  "chatty",
  "clingy",
  "sassy",
  "polite",
];
export const FOODS = ["strawberry", "cookie", "rice ball", "carrot", "pudding", "apple"];
export const GIFTS = ["flower", "pebble", "star", "leaf", "shell", "heart", "drawing", "feather"];
export const DREAMS = [
  ...FOODS,
  ...GIFTS,
  "ball",
  "moon",
  "sun",
  "cloud",
  "rain",
  "frog",
  "cat",
  "ghost",
  "music",
  "sparkle",
];
export const EMOTES = [
  "happy",
  "love",
  "sad",
  "sleepy",
  "wave",
  "dance",
  "spin",
  "blush",
  "grumpy",
  "wow",
  "giggle",
  "cry",
];
export const WANTS = ["snack", "play", "bath", "sleep", "pats", "medicine"];
export const WALKS = ["left", "right"];
export const DOODLES = [
  "heart",
  "star",
  "strawberry",
  "cloud",
  "moon",
  "sun",
  "flower",
  "cookie",
  "frog",
  "rain",
  "music",
  "cat",
  "ghost",
  "sparkle",
];
export const WEATHER_DAY = ["sunny", "cloudy", "rainy", "windy", "snowy", "foggy"];
export const WEATHER_NIGHT = ["clear", "cloudy", "rainy", "windy", "snowy", "starry"];
export const TOYS = ["ball", "yarn", "tiny drum", "rubber duck", "blocks"];

export const EFFECT_TOOLS = new Set([
  "emote",
  "want",
  "remember",
  "gift",
  "sleep",
  "dream",
  "diary",
  "walk",
  "yes",
  "no",
]);
// [yes]/[no] answer an offer event line and may appear nowhere else.
export const ANSWER_TOOLS = new Set(["yes", "no"]);
export const LOOKUP_TOOLS = new Set(["clock", "recall", "look", "yesterday"]);
export const TAG_RE = /\[(\w+)(?::([^\]\n]{1,80}))?\](\{[^}\n]*\})?/g;

export function partOfDay(h: number) {
  if (h >= 5 && h < 12) return "morning";
  if (h >= 12 && h < 17) return "afternoon";
  if (h >= 17 && h < 21) return "evening";
  return "night";
}

export function levelWords(s: Stats) {
  const pick = (v: number, words: string[]) =>
    words[Math.min(words.length - 1, Math.floor((Math.round(v) * words.length) / 101))];
  return {
    tummy: pick(s.tummy, ["starving", "hungry", "peckish", "full", "stuffed"]),
    fun: pick(s.fun, ["sad", "bored", "okay", "happy", "giddy"]),
    energy: pick(s.energy, ["exhausted", "sleepy", "okay", "awake", "hyper"]),
    clean: pick(s.clean, ["stinky", "messy", "clean", "sparkly"]),
  };
}

const pad = (n: number) => String(n).padStart(2, "0");

// st: {pet, owner, stage, day, traits, hh, mm, weather, stats, memories, today, asleep, sick}
export function header(kind: "chat" | "diary", st: WorldState) {
  const w = levelWords(st.stats);
  const lines = [
    `## tokenling ${kind}`,
    `pet: ${st.pet} · ${st.stage} · day ${st.day} · ${st.traits.join(", ")}`,
    `owner: ${st.owner}`,
  ];
  if (kind === "chat") {
    let now = `now: ${pad(st.hh)}:${pad(st.mm)} ${partOfDay(st.hh)}, ${st.weather}`;
    if (st.asleep) now += `, ${st.pet} is asleep`;
    if (st.sick) now += `, ${st.pet} is sick`;
    lines.push(now);
    lines.push(`needs: tummy ${w.tummy}, fun ${w.fun}, energy ${w.energy}, ${w.clean}`);
  } else {
    lines.push(`weather: ${st.weather}`);
  }
  if (st.memories.length) lines.push("remembers: " + st.memories.slice(-3).join("; "));
  const today = st.today.filter((t) => kind !== "chat" || !t.startsWith(st.owner + " said: "));
  if (today.length) lines.push("today: " + today.slice(-6).join("; "));
  lines.push("---");
  return lines.join("\n") + "\n";
}

const STOP = new Set([
  "the",
  "and",
  "you",
  "your",
  "are",
  "was",
  "has",
  "have",
  "for",
  "with",
  "what",
  "who",
  "does",
  "did",
  "about",
  "that",
  "this",
  "like",
  "likes",
  "she",
  "her",
  "his",
  "him",
  "they",
  "them",
  "named",
  "name",
  "is",
]);

export function words(text: string) {
  const out = new Set<string>();
  for (let wd of text.toLowerCase().match(/[a-z0-9']+/g) || []) {
    wd = wd.replace(/^'+|'+$/g, "");
    if (wd.length < 3 || STOP.has(wd)) continue;
    if (wd.endsWith("s") && wd.length > 3) wd = wd.slice(0, -1);
    out.add(wd);
  }
  return out;
}

export function recall(memories: string[], query: string) {
  const q = words(query);
  const hits = [...memories].reverse().filter((m) => [...words(m)].some((w) => q.has(w)));
  return hits.length ? hits.slice(0, 2).join("; ") : "nothing";
}

export function look(st: WorldState) {
  const shelf = st.shelf.length ? st.shelf.slice(-4).join(", ") : "empty";
  let out = `shelf: ${shelf} · toy: ${st.toy} · window: ${st.weather}`;
  if (st.poop) out += ` · poop x${st.poop}`;
  return out;
}

export function clock(st: WorldState) {
  return `${pad(st.hh)}:${pad(st.mm)} ${partOfDay(st.hh)}`;
}

export function toolResult(st: WorldState, name: string, arg?: string) {
  if (name === "clock") return clock(st);
  if (name === "look") return look(st);
  if (name === "recall") return recall(st.memories, arg || "");
  if (name === "yesterday") return st.yesterday || "no diary yet";
  return null;
}

export function a(food: string) {
  return ("aeiou".includes(food[0]) ? "an " : "a ") + food;
}

export function dreamThing(item: string) {
  return ["music", "rain"].includes(item)
    ? item
    : ["moon", "sun"].includes(item)
      ? "the " + item
      : a(item);
}

export function wakeChance(energy: number, kind: Disturbance, n = 1) {
  return Math.min(
    0.9,
    Math.max(0, { talk: 0.03, touch: 0.1, light: 0.2 }[kind] + 0.007 * energy + 0.12 * (n - 1)),
  );
}

// App event lines, exactly as train/world.py event_patterns() accepts them.
// Offers ask the pet to decide; its next line answers with [yes] or [no].
export const events = {
  offerFood: (o: string, p: string, food: string) => `* ${o} offered ${p} ${a(food)} *`,
  offerBath: (o: string, p: string) => `* ${o} offered ${p} a bath *`,
  offerMedicine: (o: string, p: string) => `* ${o} offered ${p} medicine *`,
  offerCatch: (o: string, p: string) => `* ${o} asked ${p} to play catch *`,
  caught: (o: string, p: string, n: number) =>
    `* ${o} played catch with ${p} (${n} ${n === 1 ? "star" : "stars"}) *`,
  pat: (o: string, p: string) => `* ${o} patted ${p} *`,
  lights: (o: string, on: boolean) => `* ${o} turned ${on ? "on" : "off"} the lights *`,
  woke: (p: string) => `* ${p} woke up *`,
  pooped: (p: string) => `* ${p} pooped *`,
  cleaned: (o: string) => `* ${o} cleaned up the poop *`,
  back: (o: string, span: string) => `* ${o} came back after ${span} *`,
  grew: (p: string, stage: string) => `* ${p} grew into a ${stage}! *`,
};