Jyotisha/frontend/src/lib/rectification-agentic/v9/server-focus.ts

import {
  isPersistedFocusId,
  parseAgentChoiceCopy,
  serverOwnedChoiceCopy,
  type RectificationChoiceFrame,
} from "./choice-card";
import {
  askedProbeKeysFromReceipt,
  stampChoiceSchemaWithProbe,
  previousInferenceFromReceipt,
  withNakshatraBoundaryProbe,
} from "./inference-adapter";
import { spokenFollowupForUser, spokenCollectFallbackFollowup, collectQuestionDomain, type MethodFollowup } from "./method-followup";
import { refinementFromDecisionReceipt } from "./refinement-packet";
import {
  setV10ConversationFocus,
  RectificationToolServiceError,
  safeToolErrorCode,
  type AccountingClient,
  type ConversationFocus,
} from "./tool-service";

export { parsePersistedFollowupQuestionId } from "./method-followup";

export type PersistServerFocusStatus =
  | "created"
  | "already_open"
  | "duplicate_focus"
  | "probe_already_answered"
  | "zero_information_gain"
  | "invalid_choice_schema"
  | "skipped";

export type PersistServerFocusResult = Readonly<{
  status: PersistServerFocusStatus;
  focus: ConversationFocus | null;
  questionId: string | null;
  prompt: string | null;
}>;

export function stableFollowupQuestionId(followup: MethodFollowup): string {
  if (followup.semantic_key) return `probe:${followup.semantic_key}`.slice(0, 160);
  if (
    followup.probe_year
    && followup.domain
    && followup.intent !== "reverse_verify"
    && followup.intent !== "out_of_sample_check"
  ) {
    return `${followup.method_id}:${followup.domain}:${followup.probe_year}`.slice(0, 160);
  }
  if (followup.intent === "collect_method_evidence" && !followup.choice_frame) {
    return `collect:${collectQuestionDomain(followup.domain)}:${followup.intent}`.slice(0, 160);
  }
  return (followup.choice_frame?.question_id ?? `${followup.method_id}:${followup.ask_theme}`).slice(0, 160);
}

function schemaProbeId(schema: Readonly<Record<string, unknown>> | null | undefined): string | null {
  const probeId = schema?.probe_id;
  return typeof probeId === "string" && probeId.trim() ? probeId : null;
}

export function shouldSkipDiscriminatorFollowup(followup: MethodFollowup): PersistServerFocusStatus | null {
  if (
    followup.source === "event_probe"
    && (followup.information_gain ?? 0) <= 0
  ) {
    return "zero_information_gain";
  }
  if (
    followup.intent === "distinguish_candidates"
    && ((followup.candidate_ids?.length ?? 0) < 2 || (followup.expected_outcomes?.length ?? 0) < 2)
  ) {
    return "zero_information_gain";
  }
  return null;
}

export function expectedAnswerSchemaFor(
  frame: RectificationChoiceFrame,
  questionId: string,
  decisionReceipt: Readonly<Record<string, unknown>> | null | undefined,
  followup: MethodFollowup,
): Record<string, unknown> | null {
  const copy = serverOwnedChoiceCopy(frame);
  if (!copy) return null;
  const schema: Record<string, unknown> = {
    choice: {
      prompt: copy.prompt,
      option_a: copy.option_a,
      option_b: copy.option_b,
      option_c: copy.option_c,
      option_d: copy.option_d,
      options: copy.options,
    },
    semantic_key: followup.semantic_key ?? null,
    candidate_split_hash: followup.candidate_split_hash ?? null,
    choice_kind: frame.choice_kind ?? followup.choice_kind ?? "existence",
  };
  const receipt = decisionReceipt ?? null;
  const state = withNakshatraBoundaryProbe(
    previousInferenceFromReceipt(receipt),
    refinementFromDecisionReceipt(receipt).nakshatra_boundary,
  );
  if (decisionReceipt?.inference_state !== undefined && !state) return null;
  const verifyOnly = followup.intent === "reverse_verify" || followup.intent === "out_of_sample_check";
  const stamped = stampChoiceSchemaWithProbe(
    schema,
    verifyOnly ? null : state,
    questionId,
    {
      semantic_key: followup.semantic_key,
      candidate_split_hash: followup.candidate_split_hash,
    },
  );
  if (typeof followup.probe_year === "number" && followup.probe_year > 0) {
    stamped.probe_year = followup.probe_year;
  }
  if (verifyOnly) return stamped;
  return state && stamped.scoring !== false && !schemaProbeId(stamped) ? null : stamped;
}

export type PersistedOpenQuestion = Readonly<{
  question_id: string | null;
  prompt: string | null;
  status: PersistServerFocusStatus;
  kind: "choice" | "collect_spoken";
  focus_id?: string | null;
  probe_id?: string | null;
  intent?: string;
  domain?: string | null;
  unrenderable?: true;
  reason?: string;
}>;

export function openQuestionFromPersistedFocus(result: PersistServerFocusResult): PersistedOpenQuestion | null {
  if (
    (result.status !== "created" && result.status !== "already_open")
    || !result.focus
    || !isPersistedFocusId(result.focus.id)
    || result.focus.questionId !== result.questionId
  ) return null;
  const schema = result.focus.expectedAnswerSchema;
  if (isCollectFocusSchema(schema) && schema) {
    const prompt = (typeof result.prompt === "string" && result.prompt.trim()
      ? result.prompt.trim()
      : typeof schema.prompt === "string" ? schema.prompt.trim() : "");
    if (!prompt) return null;
    return {
      question_id: result.questionId,
      prompt,
      status: result.status,
      kind: "collect_spoken",
      focus_id: result.focus.id,
      probe_id: typeof schema.probe_id === "string" ? schema.probe_id : null,
      intent: result.focus.intent,
      domain: result.focus.targetDomain ?? null,
    };
  }
  if (!result.prompt || !parseAgentChoiceCopy(schema)) {
    return {
      question_id: result.questionId,
      prompt: null,
      status: result.status,
      kind: "choice",
      unrenderable: true,
      reason: "invalid_choice_schema",
    };
  }
  return {
    question_id: result.questionId,
    prompt: result.prompt,
    status: result.status,
    kind: "choice",
  };
}

export function isRenderableChoiceOpenQuestion(
  open: PersistedOpenQuestion | null | undefined,
): open is PersistedOpenQuestion & { kind: "choice" } {
  return Boolean(open && open.kind === "choice" && open.unrenderable !== true);
}

export const COLLECT_FOCUS_SCHEMA_KEY = "collect";
export const COLLECT_FOCUS_RETRY_SUFFIX = "next";

const PERSISTABLE_FOCUS_DOMAINS = new Set([
  "education",
  "career",
  "relationship",
  "relocation",
  "finance",
  "health",
  "family",
  "other",
]);

export function persistableFocusDomain(domain: string | null | undefined): string | null {
  if (!domain || domain === "unknown" || domain === "active_focus") return null;
  if (domain === "health_pressure") return "health";
  if (domain === "occupation") return "other";
  if (PERSISTABLE_FOCUS_DOMAINS.has(domain)) return domain;
  return domain;
}

function isFocusIdempotencyConflict(error: unknown): boolean {
  const code = error instanceof RectificationToolServiceError
    ? error.code
    : safeToolErrorCode(error);
  return code === "focus_idempotency_conflict" || code.includes("focus_idempotency_conflict");
}

export async function linkFocusAskedTurn(input: {
  accounting: AccountingClient;
  userId: string;
  caseId: string;
  focus: ConversationFocus;
  askedTurnId?: string | null;
}): Promise<ConversationFocus> {
  if (!input.askedTurnId || input.focus.askedTurnId) return input.focus;
  try {
    const result = await setV10ConversationFocus(input.accounting, input.userId, input.caseId, {
      questionId: input.focus.questionId,
      intent: input.focus.intent,
      targetEvidenceId: input.focus.targetEvidenceId,
      targetDomain: input.focus.targetDomain,
      targetKind: input.focus.targetKind,
      expectedAnswerSchema: input.focus.expectedAnswerSchema,
      askedTurnId: input.askedTurnId,
    });
    return result.focus;
  } catch {
    return input.focus;
  }
}

export function collectFocusSchema(followup: MethodFollowup): Record<string, unknown> | null {
  const prompt = spokenFollowupForUser({ ...followup, choice_frame: null });
  if (!prompt) return null;
  const schema: Record<string, unknown> = {
    prompt,
    [COLLECT_FOCUS_SCHEMA_KEY]: true,
  };
  if (followup.semantic_key) schema.semantic_key = followup.semantic_key;
  return schema;
}

export function serverOwnedExpectedAnswerSchema(
  followup: MethodFollowup,
  decisionReceipt?: Readonly<Record<string, unknown>> | null,
): Record<string, unknown> | null {
  const frame = followup.choice_frame;
  if (frame) {
    const schema = expectedAnswerSchemaFor(
      frame,
      stableFollowupQuestionId(followup),
      decisionReceipt,
      followup,
    );
    if (schema?.choice) return schema;
    return collectFocusSchema(spokenCollectFallbackFollowup(followup));
  }
  return collectFocusSchema(followup);
}

export function isCollectFocusSchema(schema: Readonly<Record<string, unknown>> | null | undefined): boolean {
  return schema?.[COLLECT_FOCUS_SCHEMA_KEY] === true && typeof schema.prompt === "string";
}

async function persistCollectFocus(input: {
  accounting: AccountingClient;
  userId: string;
  caseId: string;
  activeFocus: ConversationFocus | null;
  followup: MethodFollowup;
  askedTurnId?: string | null;
}): Promise<PersistServerFocusResult> {
  const schema = collectFocusSchema(input.followup);
  if (!schema) {
    return {
      status: "skipped",
      focus: input.activeFocus,
      questionId: null,
      prompt: null,
    };
  }
  const questionId = stableFollowupQuestionId(input.followup);
  const prompt = typeof schema.prompt === "string" ? schema.prompt : null;
  const active = input.activeFocus;
  if (
    active
    && active.questionId === questionId
    && active.questionId !== "active_focus:active_focus"
    && isCollectFocusSchema(active.expectedAnswerSchema)
  ) {
    const focus = await linkFocusAskedTurn({
      accounting: input.accounting,
      userId: input.userId,
      caseId: input.caseId,
      focus: active,
      askedTurnId: input.askedTurnId,
    });
    return { status: "already_open", focus, questionId: focus.questionId, prompt };
  }
  const insertFocus = async (id: string): Promise<PersistServerFocusResult> => {
    const result = await setV10ConversationFocus(input.accounting, input.userId, input.caseId, {
      questionId: id,
      intent: input.followup.intent,
      targetEvidenceId: null,
      targetDomain: persistableFocusDomain(input.followup.domain)
        ?? (input.followup.intent === "collect_method_evidence"
          ? collectQuestionDomain(input.followup.domain)
          : null),
      targetKind: null,
      expectedAnswerSchema: schema,
      askedTurnId: input.askedTurnId ?? null,
    });
    return {
      status: result.idempotent ? "already_open" : "created",
      focus: result.focus,
      questionId: result.focus.questionId,
      prompt,
    };
  };
  try {
    return await insertFocus(questionId);
  } catch (error) {
    if (!isFocusIdempotencyConflict(error)) {
      return {
        status: "skipped",
        focus: input.activeFocus,
        questionId: null,
        prompt: null,
      };
    }
    const retryId = `${questionId}:${COLLECT_FOCUS_RETRY_SUFFIX}`.slice(0, 160);
    if (retryId === questionId) {
      return {
        status: "duplicate_focus",
        focus: input.activeFocus,
        questionId,
        prompt,
      };
    }
    try {
      return await insertFocus(retryId);
    } catch (retryError) {
      if (isFocusIdempotencyConflict(retryError)) {
        return {
          status: "duplicate_focus",
          focus: input.activeFocus,
          questionId: retryId,
          prompt,
        };
      }
      return {
        status: "skipped",
        focus: input.activeFocus,
        questionId: null,
        prompt: null,
      };
    }
  }
}

async function persistSpokenChoiceFallback(input: {
  accounting: AccountingClient;
  userId: string;
  caseId: string;
  activeFocus: ConversationFocus | null;
  followup: MethodFollowup;
  askedTurnId?: string | null;
}): Promise<PersistServerFocusResult> {
  return persistCollectFocus({
    accounting: input.accounting,
    userId: input.userId,
    caseId: input.caseId,
    activeFocus: input.activeFocus,
    followup: spokenCollectFallbackFollowup(input.followup),
    askedTurnId: input.askedTurnId,
  });
}

export async function persistServerOwnedFocus(input: {
  accounting: AccountingClient;
  userId: string;
  caseId: string;
  activeFocus: ConversationFocus | null;
  decisionReceipt: Readonly<Record<string, unknown>> | null | undefined;
  followup: MethodFollowup | null;
  askedTurnId?: string | null;
}): Promise<PersistServerFocusResult> {
  const followup = input.followup;
  const frame = followup?.choice_frame ?? null;
  if (!followup) {
    return {
      status: "skipped",
      focus: input.activeFocus,
      questionId: null,
      prompt: null,
    };
  }
  if (!frame) {
    if (followup.intent === "distinguish_candidates") {
      return persistSpokenChoiceFallback({
        accounting: input.accounting,
        userId: input.userId,
        caseId: input.caseId,
        activeFocus: input.activeFocus,
        followup,
        askedTurnId: input.askedTurnId,
      });
    }
    if (followup.intent === "collect_method_evidence") {
      return persistCollectFocus({
        accounting: input.accounting,
        userId: input.userId,
        caseId: input.caseId,
        activeFocus: input.activeFocus,
        followup,
        askedTurnId: input.askedTurnId,
      });
    }
    return {
      status: "skipped",
      focus: input.activeFocus,
      questionId: null,
      prompt: null,
    };
  }
  const skip = shouldSkipDiscriminatorFollowup(followup);
  if (skip) {
    return { status: skip, focus: input.activeFocus, questionId: null, prompt: null };
  }
  const questionId = stableFollowupQuestionId(followup);
  const answeredKeys = new Set(askedProbeKeysFromReceipt(input.decisionReceipt));
  if (followup.semantic_key && answeredKeys.has(followup.semantic_key)) {
    return {
      status: "probe_already_answered",
      focus: input.activeFocus,
      questionId,
      prompt: null,
    };
  }
  const copy = serverOwnedChoiceCopy(frame);
  if (!copy) {
    if (followup.intent === "distinguish_candidates") {
      return persistSpokenChoiceFallback({
        accounting: input.accounting,
        userId: input.userId,
        caseId: input.caseId,
        activeFocus: input.activeFocus,
        followup,
        askedTurnId: input.askedTurnId,
      });
    }
    return { status: "invalid_choice_schema", focus: input.activeFocus, questionId, prompt: null };
  }
  const schema = expectedAnswerSchemaFor(frame, questionId, input.decisionReceipt, followup);
  if (!schema?.choice) {
    return { status: "invalid_choice_schema", focus: input.activeFocus, questionId, prompt: null };
  }
  const prompt = copy?.prompt ?? null;
  const active = input.activeFocus;
  if (
    active
    && (
      active.questionId === questionId
      || (schemaProbeId(schema) && schemaProbeId(active.expectedAnswerSchema) === schemaProbeId(schema))
    )
  ) {
    const focus = await linkFocusAskedTurn({
      accounting: input.accounting,
      userId: input.userId,
      caseId: input.caseId,
      focus: active,
      askedTurnId: input.askedTurnId,
    });
    return { status: "already_open", focus, questionId: focus.questionId, prompt };
  }
  try {
    const result = await setV10ConversationFocus(input.accounting, input.userId, input.caseId, {
      questionId,
      intent: followup.intent,
      targetEvidenceId: null,
      targetDomain: followup.domain,
      targetKind: null,
      expectedAnswerSchema: schema,
      askedTurnId: input.askedTurnId ?? null,
    });
    return {
      status: result.idempotent ? "already_open" : "created",
      focus: result.focus,
      questionId: result.focus.questionId,
      prompt,
    };
  } catch (error) {
    const code = error instanceof RectificationToolServiceError
      ? error.code
      : safeToolErrorCode(error);
    if (code === "focus_idempotency_conflict" || code.includes("focus_idempotency_conflict")) {
      return {
        status: "duplicate_focus",
        focus: input.activeFocus,
        questionId,
        prompt,
      };
    }
    throw error;
  }
}