import { LayaRuntimeError } from "./errors";
import type { LayaQuestion, ModelManifest, QuestionType, TokenUsage } from "./shared/types";
export interface TokenEncoding { ids: number[]; attention_mask?: number[] }
export interface TokenizerLike { encode(text: string, options?: { add_special_tokens?: boolean }): TokenEncoding | Promise<TokenEncoding> }
export interface PreparedQuestion { id: string; inputIds: number[]; attentionMask: number[]; markerPositions: number[]; markerMask: number[]; questionType: number; type: QuestionType; optionLabels: string[]; criteria: unknown; usage: TokenUsage }
const QUESTION_TYPES = { choice: 0, score: 1, noul: 2 } as const;
export function serializeState(value: unknown): string { return typeof value === "string" ? value : serializeJson(value); }
function serializeJson(value: unknown): string {
  if (value === null) return "null";
  if (typeof value === "string") return JSON.stringify(value);
  if (typeof value === "number") return Number.isFinite(value) ? String(value) : "null";
  if (typeof value === "boolean") return value ? "true" : "false";
  if (Array.isArray(value)) return `[${value.map(serializeJson).join(", ")}]`;
  if (typeof value === "object") return `{${Object.entries(value as Record<string, unknown>).map(([k, v]) => `${JSON.stringify(k)}: ${serializeJson(v)}`).join(", ")}}`;
  return "null";
}
const renderCriterion = (value: unknown) => typeof value === "string" ? value : serializeJson(value);
export function renderOptions(question: LayaQuestion): { labels: string[]; text: string[] } {
  if (question.type === "choice") {
    const entries: Array<[string, unknown]> = Array.isArray(question.criteria) ? question.criteria.map((value) => [value, null]) : Object.entries(question.criteria);
    return { labels: entries.map(([key]) => key), text: entries.map(([key, value]) => value === null || value === "" ? key : `${key}: ${renderCriterion(value)}`) };
  }
  if (question.type === "score") return { labels: question.criteria.map((_, i) => String(i)), text: question.criteria.map((value, i) => `level ${i}: ${renderCriterion(value)}`) };
  const criteria = question.criteria ?? {};
  return { labels: ["false", "true"], text: [`false: ${criteria.false == null || criteria.false === "" ? "no, the statement does not hold" : renderCriterion(criteria.false)}`, `true: ${criteria.true == null || criteria.true === "" ? "yes, the statement holds" : renderCriterion(criteria.true)}`] };
}
export function validateQuestion(question: LayaQuestion): void {
  if (!question.instructions.trim()) throw new LayaRuntimeError("Question instructions cannot be empty", "INVALID_QUESTION");
  if (question.type === "choice" && Object.keys(question.criteria).length < 1) throw new LayaRuntimeError("Choice questions require at least one option", "INVALID_QUESTION");
  if (question.type === "score" && question.criteria.length < 2) throw new LayaRuntimeError("Score questions require at least two levels", "INVALID_QUESTION");
}
async function encode(tokenizer: TokenizerLike, text: string): Promise<number[]> { return (await tokenizer.encode(text, { add_special_tokens: false })).ids; }
export async function prepareQuestion(tokenizer: TokenizerLike, manifest: ModelManifest, state: unknown, id: string, question: LayaQuestion): Promise<PreparedQuestion> {
  validateQuestion(question);
  const scrub = (text: string) => text.split(manifest.specialTokens.mask).join(" ");
  const options = renderOptions(question);
  let headIds = await encode(tokenizer, `${question.type} question: ${scrub(question.instructions)}`);
  let optionIds = await Promise.all(options.text.map(async (text) => [manifest.specialTokens.maskId, ...(await encode(tokenizer, ` ${scrub(text)}`)).slice(0, 48)]));
  let budget = manifest.headTokens - optionIds.reduce((sum, item) => sum + item.length, 0);
  if (budget < 16) { const per = Math.max(4, Math.floor((manifest.headTokens - 16) / Math.max(1, optionIds.length))); optionIds = optionIds.map((item) => item.slice(0, per)); budget = manifest.headTokens - optionIds.reduce((sum, item) => sum + item.length, 0); }
  headIds = headIds.slice(0, Math.max(8, budget));
  const inputIds = [manifest.specialTokens.clsId, ...headIds, manifest.specialTokens.sepId];
  const markerPositions: number[] = [];
  for (const item of optionIds) { markerPositions.push(inputIds.length); inputIds.push(...item); }
  inputIds.push(manifest.specialTokens.sepId);
  const room = Math.max(0, manifest.contextLength - inputIds.length - 1);
  const fullState = await encode(tokenizer, scrub(serializeState(state)));
  inputIds.push(...fullState.slice(0, room), manifest.specialTokens.sepId);
  if (markerPositions.length !== options.labels.length || markerPositions.some((position) => position >= manifest.contextLength)) throw new LayaRuntimeError(`Options exceed head budget ${manifest.headTokens}`, "CONTEXT_OVERFLOW");
  const clipped = inputIds.slice(0, manifest.contextLength);
  return { id, inputIds: clipped, attentionMask: clipped.map(() => 1), markerPositions, markerMask: markerPositions.map(() => 1), questionType: QUESTION_TYPES[question.type], type: question.type, optionLabels: options.labels, criteria: question.criteria, usage: { inputTokens: clipped.length, truncated: fullState.length > room, optionBudgetExceeded: false } };
}
export function padPrepared(batch: PreparedQuestion[], padId: number) {
  if (!batch.length) throw new LayaRuntimeError("At least one question is required", "INVALID_QUESTION");
  const sequence = Math.max(...batch.map((item) => item.inputIds.length)); const options = Math.max(...batch.map((item) => item.markerPositions.length));
  return { inputIds: batch.flatMap((item) => [...item.inputIds, ...Array(sequence - item.inputIds.length).fill(padId)]), attentionMask: batch.flatMap((item) => [...item.attentionMask, ...Array(sequence - item.attentionMask.length).fill(0)]), markerPositions: batch.flatMap((item) => [...item.markerPositions, ...Array(options - item.markerPositions.length).fill(0)]), markerMask: batch.flatMap((item) => [...item.markerMask, ...Array(options - item.markerMask.length).fill(0)]), questionTypes: batch.map((item) => item.questionType), dimensions: { batch: batch.length, sequence, options } };
}
export function tempBucket(type: QuestionType, count: number): string { return `${type}:${count <= 2 ? "2" : count <= 5 ? "3-5" : count <= 10 ? "6-10" : "11+"}`; }
export function clampTemperature(value: unknown): number { const parsed = Number(value); return Number.isFinite(parsed) ? Math.min(5, Math.max(0.5, parsed)) : 1; }
export function confidenceFromProbabilities(values: number[]): number { if (values.length < 2) return 1; const entropy = -values.reduce((sum, value) => sum + value * Math.log(Math.max(value, 1e-12)), 0); return Math.max(0, Math.min(1, 1 - entropy / Math.log(values.length))); }