Jyotisha/frontend/src/lib/rectification-agentic/v9/collection-question-pool.ts

/**
 * S1 collection question pool: invite → anchored follow-ups from the
 * user's own years → yearless generic prompts. Ranking is by information
 * value, not a fixed domain rotation. Age-band years never enter prompts.
 */

import { USER_COLLECT_QUESTION } from "../user-copy.ts";

export const COLLECT_KIND_ORDER = [
  "education",
  "career",
  "relocation",
  "relationship",
  "family",
  "finance",
  "health_pressure",
] as const;

export type CollectKind = (typeof COLLECT_KIND_ORDER)[number];

export type CollectionPoolKind = "invite" | "anchor" | "generic";

export type CollectionEvidence = Readonly<{
  status: string;
  domain: string;
  datePrecision: string;
  occurredFrom: string | null;
  occurredTo: string | null;
  eventKind?: string | null;
  summary?: string | null;
}>;

export type CollectionTopic = Readonly<Record<string, unknown>>;

export type CollectionPoolItem = Readonly<{
  kind: CollectionPoolKind;
  value: number;
  prompt: string;
  key: string;
  domain: string;
  targetKind: string | null;
  year: number | null;
}>;

const KIND_EXAMPLES: Readonly<Record<CollectKind, string>> = {
  education: "上大学或毕业",
  career: "第一份工作",
  relocation: "搬到别的城市",
  relationship: "谈恋爱或结婚",
  family: "家里添丁或长辈住院",
  finance: "收入明显变过或大笔支出",
  health_pressure: "生病受伤",
};

const KIND_LABEL: Readonly<Record<CollectKind, string>> = {
  education: "上学",
  career: "工作",
  relocation: "搬家",
  relationship: "感情",
  family: "家里",
  finance: "钱的方面",
  health_pressure: "身体或压力",
};

const GENERIC_PROMPTS: Readonly<Record<CollectKind, string>> = {
  education: USER_COLLECT_QUESTION.education,
  career: USER_COLLECT_QUESTION.career,
  relocation: USER_COLLECT_QUESTION.relocation,
  relationship: USER_COLLECT_QUESTION.relationship,
  family: USER_COLLECT_QUESTION.family,
  finance: USER_COLLECT_QUESTION.finance,
  health_pressure: USER_COLLECT_QUESTION.health_pressure,
};

/** T2 lets the table start at 0.5; named rules keep the product examples. */
const ANCHOR_PROBABILITY = 0.5;

const INVITE_VALUE = 2;
const GENERIC_VALUE = 0.15;

export function normalizeCollectKind(domain: string | null | undefined): CollectKind | null {
  if (domain === "health" || domain === "health_pressure") return "health_pressure";
  if ((COLLECT_KIND_ORDER as readonly string[]).includes(domain ?? "")) {
    return domain as CollectKind;
  }
  return null;
}

function topicQuestionId(topic: CollectionTopic): string {
  return typeof topic.questionId === "string"
    ? topic.questionId
    : typeof topic.question_id === "string"
      ? topic.question_id
      : "";
}

function topicStatus(topic: CollectionTopic): string {
  return typeof topic.status === "string" ? topic.status : "";
}

function topicDomain(topic: CollectionTopic): string | null {
  return typeof topic.target_domain === "string"
    ? topic.target_domain
    : typeof topic.targetDomain === "string"
      ? topic.targetDomain
      : null;
}

function topicKind(topic: CollectionTopic): string {
  return typeof topic.target_kind === "string"
    ? topic.target_kind
    : typeof topic.targetKind === "string"
      ? topic.targetKind
      : "";
}

export function isInviteCollectTopic(topic: CollectionTopic): boolean {
  return topicQuestionId(topic).startsWith("collect:invite:")
    || topicKind(topic) === "invite_more";
}

export function askedCollectKeys(topics: readonly CollectionTopic[]): ReadonlySet<string> {
  const keys = new Set<string>();
  for (const topic of topics) {
    const status = topicStatus(topic);
    if (status === "active") continue;
    if (isInviteCollectTopic(topic)) continue;
    const questionId = topicQuestionId(topic).replace(/:(?:next|next2|next3)$/, "");
    if (questionId.startsWith("collect:")) keys.add(questionId);
    const kind = topicKind(topic);
    if (kind.startsWith("anchor:") || kind.startsWith("generic:")) {
      keys.add(`collect:${kind}`);
    }
  }
  return keys;
}

export function inviteDeclined(topics: readonly CollectionTopic[]): boolean {
  return topics.some((topic) => {
    if (!isInviteCollectTopic(topic)) return false;
    const status = topicStatus(topic);
    return status === "declined" || status === "skipped";
  });
}

function isConfirmedDated(item: CollectionEvidence): boolean {
  return item.status === "confirmed"
    && item.datePrecision !== "unknown"
    && Boolean(item.occurredFrom || item.occurredTo);
}

export function evidenceYear(item: CollectionEvidence): number | null {
  const raw = item.occurredFrom || item.occurredTo;
  if (!raw || raw.length < 4 || !/^\d{4}/.test(raw)) return null;
  const year = Number(raw.slice(0, 4));
  return year >= 1900 && year <= 2100 ? year : null;
}

function evidenceMonth(item: CollectionEvidence): number | null {
  const raw = item.occurredFrom || item.occurredTo || "";
  if (raw.length < 7 || raw[4] !== "-") return null;
  const month = Number(raw.slice(5, 7));
  return month >= 1 && month <= 12 ? month : null;
}

export function coveredCollectKinds(evidence: readonly CollectionEvidence[]): ReadonlySet<CollectKind> {
  const covered = new Set<CollectKind>();
  for (const item of evidence) {
    if (!isConfirmedDated(item)) continue;
    const kind = normalizeCollectKind(item.domain);
    if (kind) covered.add(kind);
  }
  return covered;
}

function declinedKinds(topics: readonly CollectionTopic[]): ReadonlySet<CollectKind> {
  const declined = new Set<CollectKind>();
  for (const topic of topics) {
    const status = topicStatus(topic);
    if (status !== "declined" && status !== "skipped") continue;
    const questionId = topicQuestionId(topic);
    if (questionId.startsWith("collect:invite:")) continue;
    if (questionId.startsWith("collect:other:")) continue;
    const kind = normalizeCollectKind(topicDomain(topic))
      ?? (questionId.startsWith("collect:generic:")
        ? normalizeCollectKind(questionId.split(":")[2] ?? "")
        : null);
    if (kind) declined.add(kind);
  }
  return declined;
}

function uncoveredKinds(
  evidence: readonly CollectionEvidence[],
  declined: ReadonlySet<CollectKind>,
): CollectKind[] {
  const covered = coveredCollectKinds(evidence);
  return COLLECT_KIND_ORDER.filter((kind) => !covered.has(kind) && !declined.has(kind));
}

export function examplePhrases(
  evidence: readonly CollectionEvidence[],
  declined: ReadonlySet<CollectKind>,
  limit = 4,
): string[] {
  return uncoveredKinds(evidence, declined)
    .map((kind) => KIND_EXAMPLES[kind])
    .slice(0, limit);
}

function invitePrompt(evidence: readonly CollectionEvidence[], declined: ReadonlySet<CollectKind>): string {
  const examples = examplePhrases(evidence, declined, 4);
  if (examples.length >= 2) {
    return `还有吗？比如${examples.join("、")}。`;
  }
  if (examples.length === 1) {
    return `还有吗？比如${examples[0]}、家里添丁或长辈住院。`;
  }
  return "还有吗？比如第一份工作、搬到别的城市。";
}

type AnchorRule = Readonly<{
  keyKind: string;
  target: CollectKind;
  probability: number;
  prompt: (year: number, month: number | null) => string;
  applies: (item: CollectionEvidence, year: number) => boolean;
}>;

function kindLooksLike(item: CollectionEvidence, tokens: readonly string[]): boolean {
  const blob = `${item.eventKind ?? ""} ${item.summary ?? ""}`.toLowerCase();
  return tokens.some((token) => blob.includes(token));
}

const ANCHOR_RULES: readonly AnchorRule[] = [
  {
    keyKind: "education_completion",
    target: "career",
    probability: 0.8,
    prompt: (year) => `${year} 年毕业后第一份工作大概哪年开始？`,
    applies: (item, year) => {
      if (normalizeCollectKind(item.domain) !== "education") return false;
      if (item.eventKind === "education_completion") return true;
      return kindLooksLike(item, ["毕业", "completion"]) && evidenceYear(item) === year;
    },
  },
  {
    keyKind: "education_start",
    target: "relocation",
    probability: 0.6,
    prompt: (year) => `${year} 年上大学是搬到别的城市住吗？`,
    applies: (item, year) => {
      if (normalizeCollectKind(item.domain) !== "education") return false;
      if (item.eventKind === "education_start") return true;
      return kindLooksLike(item, ["入学", "大学", "start"]) && evidenceYear(item) === year;
    },
  },
  {
    keyKind: "relationship_follow",
    target: "relationship",
    probability: ANCHOR_PROBABILITY,
    prompt: () => "后来是分开了还是结婚了，大概哪年？",
    applies: (item) => {
      if (normalizeCollectKind(item.domain) !== "relationship") return false;
      return item.eventKind === "relationship_start"
        || kindLooksLike(item, ["在一起", "交往", "恋爱", "start"]);
    },
  },
  {
    keyKind: "after_event",
    target: "career",
    probability: ANCHOR_PROBABILITY,
    prompt: (year, month) => (
      month
        ? `${year} 年 ${month} 月那次之后，工作或住处有没有变？`
        : `${year} 年那次之后，工作或住处有没有变？`
    ),
    applies: () => true,
  },
];

function gapWeight(target: CollectKind, covered: ReadonlySet<CollectKind>): number {
  return covered.has(target) ? 0.5 : 1;
}

function educationCompletionYear(evidence: readonly CollectionEvidence[]): number | null {
  const edu = evidence
    .filter((item) => isConfirmedDated(item) && normalizeCollectKind(item.domain) === "education")
    .map((item) => ({ item, year: evidenceYear(item) }))
    .filter((row): row is { item: CollectionEvidence; year: number } => row.year != null)
    .sort((left, right) => left.year - right.year);
  if (edu.length === 0) return null;
  const marked = edu.find((row) => (
    row.item.eventKind === "education_completion" || kindLooksLike(row.item, ["毕业", "completion"])
  ));
  return (marked ?? edu[edu.length - 1]).year;
}

export function anchoredFollowups(
  evidence: readonly CollectionEvidence[],
  declined: ReadonlySet<CollectKind>,
  asked: ReadonlySet<string>,
): CollectionPoolItem[] {
  const covered = coveredCollectKinds(evidence);
  const items: CollectionPoolItem[] = [];
  const seen = new Set<string>();
  const completionYear = educationCompletionYear(evidence);
  for (const item of evidence) {
    if (!isConfirmedDated(item)) continue;
    const year = evidenceYear(item);
    if (year == null) continue;
    const month = evidenceMonth(item);
    const sourceKind = normalizeCollectKind(item.domain);
    if (!sourceKind) continue;
    let specific = 0;
    for (const rule of ANCHOR_RULES) {
      if (rule.keyKind === "after_event" && specific > 0) continue;
      if (!rule.applies(item, year)) continue;
      if (rule.keyKind === "education_completion" && completionYear != null && year !== completionYear) {
        continue;
      }
      if (declined.has(rule.target) && rule.target !== sourceKind) continue;
      const key = `collect:anchor:${rule.keyKind}:${year}`;
      if (asked.has(key) || seen.has(key)) continue;
      seen.add(key);
      if (rule.keyKind !== "after_event") specific += 1;
      items.push({
        kind: "anchor",
        value: rule.probability * gapWeight(rule.target, covered),
        prompt: rule.prompt(year, month),
        key,
        domain: rule.target,
        targetKind: rule.keyKind,
        year,
      });
    }
  }
  return items.sort((left, right) => right.value - left.value || left.key.localeCompare(right.key));
}

function genericItems(
  evidence: readonly CollectionEvidence[],
  declined: ReadonlySet<CollectKind>,
  asked: ReadonlySet<string>,
): CollectionPoolItem[] {
  const covered = coveredCollectKinds(evidence);
  return uncoveredKinds(evidence, declined)
    .filter((kind) => !asked.has(`collect:generic:${kind}`))
    .map((kind) => ({
      kind: "generic" as const,
      value: GENERIC_VALUE * gapWeight(kind, covered),
      prompt: GENERIC_PROMPTS[kind],
      key: `collect:generic:${kind}`,
      domain: kind,
      targetKind: null,
      year: null,
    }))
    .sort((left, right) => right.value - left.value || left.key.localeCompare(right.key));
}

export function collectionQuestionPool(
  evidence: readonly CollectionEvidence[],
  declinedTopics: readonly CollectionTopic[] = [],
  inviteHistory: readonly CollectionTopic[] = declinedTopics,
): CollectionPoolItem[] {
  const topics = [...declinedTopics, ...inviteHistory];
  const asked = askedCollectKeys(topics);
  const declined = declinedKinds(topics);
  const pool: CollectionPoolItem[] = [];
  if (!inviteDeclined(topics)) {
    pool.push({
      kind: "invite",
      value: INVITE_VALUE,
      prompt: invitePrompt(evidence, declined),
      key: "collect:invite:more",
      domain: "other",
      targetKind: null,
      year: null,
    });
  }
  const anchors = anchoredFollowups(evidence, declined, asked);
  pool.push(...anchors);
  if (!pool.some((item) => item.kind === "invite") && anchors.length === 0) {
    pool.push(...genericItems(evidence, declined, asked));
  }
  return [...pool].sort((left, right) => right.value - left.value || left.key.localeCompare(right.key));
}

function eventPhrase(item: CollectionEvidence): string {
  const year = evidenceYear(item);
  const month = evidenceMonth(item);
  const when = year == null
    ? ""
    : month
      ? `${year} 年 ${month} 月`
      : `${year} 年`;
  if (item.eventKind === "education_completion" || kindLooksLike(item, ["毕业", "completion"])) {
    return `${when}毕业`.trim();
  }
  if (item.eventKind === "education_start" || kindLooksLike(item, ["入学", "大学", "start"])) {
    return `${when}上大学`.trim();
  }
  const kind = normalizeCollectKind(item.domain);
  const label = kind ? KIND_LABEL[kind] : "这件事";
  const summary = item.summary?.trim();
  if (summary && !/^[a-z_]+$/i.test(summary)) return `${when}${summary}`.trim();
  return `${when}${label}`.trim();
}

export function preciseGapNarration(
  evidence: readonly CollectionEvidence[],
  declinedTopics: readonly CollectionTopic[] = [],
): string {
  const dated = evidence.filter(isConfirmedDated);
  const phrases = dated.map(eventPhrase).filter(Boolean);
  const covered = coveredCollectKinds(evidence);
  const declined = declinedKinds(declinedTopics);
  const examples = examplePhrases(evidence, declined, 4);
  const exampleText = examples.length >= 2
    ? examples.slice(0, 3).join("、")
    : "第一份工作哪年开始、哪年搬到别的城市";
  const recorded = phrases.length > 0
    ? `现在记下的是${phrases.join("和")}。`
    : "现在还没有记下带年月的经历。";
  const sameKind = covered.size === 1
    ? `都是${KIND_LABEL[[...covered][0]]}的事。`
    : covered.size > 1
      ? ""
      : "";
  const needOther = covered.size === 1
    ? `再来一件不是${KIND_LABEL[[...covered][0]]}的、记得大概年月的事就能开始筛，比如${exampleText}。`
    : `再来一件记得大概年月的事就能开始筛，比如${exampleText}。`;
  return `${recorded}${sameKind}${needOther}`;
}

export function moreCollectHint(
  evidence: readonly CollectionEvidence[],
  declinedTopics: readonly CollectionTopic[] = [],
): string {
  const remaining = collectionQuestionPool(evidence, declinedTopics)
    .filter((item) => item.kind !== "invite");
  const top = remaining[0];
  if (top?.kind === "anchor") {
    const trimmed = top.prompt.replace(/[？?]$/, "");
    return `如果还记得${trimmed.replace(/大概/g, "")}，范围还能再收一截。`;
  }
  const examples = examplePhrases(evidence, declinedKinds(declinedTopics), 3);
  const exampleText = examples.length >= 2
    ? examples.slice(0, 2).join("、")
    : "第一份工作哪年开始、哪年搬到别的城市";
  return `如果还记得${exampleText}，范围还能再收一截。`;
}

export const COLLECT_FLOW_BANNED_PHRASES = [
  "任何领域",
  "领域不限",
  "方法覆盖",
  "训练门",
  "holdout",
  "做不了",
  "材料不够",
  "还差",
  "探针",
  "分盘",
  "领域",
] as const;