/**
 * Question normalization and answer conformance.
 *
 * Both engines (native Jev and the LLM proxy) go through this module, so the tool's
 * input and output shapes are fixed by the plugin rather than by whichever backend
 * answered: questions are validated against the documented TypeSafe shapes, and every
 * answer is conformed into the same typed result (or reported as a per-question issue).
 */

import { JevError } from "./errors.ts";

export type JsonValue = string | number | boolean | null | JsonValue[] | { [key: string]: JsonValue };

export interface NoulQuestion {
  type: "noul";
  instructions: JsonValue;
  criteria?: { true?: string; false?: string };
}

export interface ChoiceQuestion {
  type: "choice";
  instructions: JsonValue;
  criteria: Record<string, string | null>;
}

export interface ScoreQuestion {
  type: "score";
  instructions: JsonValue;
  /** Ordered level descriptions: index 0 is the lowest level. */
  criteria: string[];
}

export type JevQuestion = NoulQuestion | ChoiceQuestion | ScoreQuestion;

export interface NoulAnswer {
  type: "noul";
  noul: number;
}

export interface ChoiceAnswer {
  type: "choice";
  choice: string;
  probabilities: Record<string, number>;
  /** Absent when the engine did not report one (the LLM proxy only relays what the model returned). */
  confidence?: number;
}

export interface ScoreAnswer {
  type: "score";
  score: number;
  legend: Record<string, string>;
  probabilities: Record<string, number>;
  confidence?: number;
}

export type JevAnswer = NoulAnswer | ChoiceAnswer | ScoreAnswer;

export interface NormalizedQuestions {
  ids: string[];
  questions: Record<string, JevQuestion>;
  /** Per-question backend/model overrides, stripped from the request payload sent to the service. */
  routes: Record<string, QuestionRoute>;
}

/** A single question can be answered by a different backend or model than the rest of the call. */
export interface QuestionRoute {
  backend?: string;
  model?: string;
}

export interface ConformResult {
  answers: Record<string, JevAnswer>;
  /** Fatal per-question problems: that question has no usable answer. */
  issues: Record<string, string>;
  /** Non-fatal per-question observations, e.g. probabilities that had to be renormalized. */
  notes: Record<string, string>;
}

const TYPE_ALIASES: Record<string, JevQuestion["type"]> = {
  noul: "noul",
  yes_no: "noul",
  yesno: "noul",
  boolean: "noul",
  bool: "noul",
  choice: "choice",
  classify: "choice",
  classification: "choice",
  select: "choice",
  score: "score",
  rating: "score",
  rate: "score",
  scale: "score",
  grade: "score",
};

/** Validate and canonicalize the tool's `questions` array into the request map TypeSafe expects. */
export function normalizeQuestions(raw: unknown): NormalizedQuestions {
  if (!Array.isArray(raw) || raw.length === 0) {
    throw new JevError("questions must be a non-empty array of question objects.");
  }
  const ids: string[] = [];
  const questions: Record<string, JevQuestion> = {};
  const routes: Record<string, QuestionRoute> = {};
  raw.forEach((entry, index) => {
    const where = `questions[${index}]`;
    if (entry === null || typeof entry !== "object" || Array.isArray(entry)) {
      throw new JevError(`${where} must be an object.`);
    }
    const input = entry as Record<string, unknown>;
    const id = readRequiredString(input.id, `${where}.id`);
    if (questions[id] !== undefined) throw new JevError(`${where}.id "${id}" is used more than once.`);
    const rawType = readRequiredString(input.type, `${where}.type`).toLowerCase().replace(/[\s-]+/g, "_");
    const type = TYPE_ALIASES[rawType];
    if (type === undefined) {
      throw new JevError(`${where}.type must be noul, choice, or score (got "${rawType}").`);
    }
    const instructions = readInstructions(
      input.instructions ?? input.question ?? input.prompt ?? input.instruction,
      `${where}.instructions`,
    );
    const route: QuestionRoute = {};
    const backendOverride = input.backend ?? input.provider;
    if (backendOverride !== undefined) route.backend = readRequiredString(backendOverride, `${where}.backend`);
    if (input.model !== undefined) route.model = readRequiredString(input.model, `${where}.model`);
    if (route.backend !== undefined || route.model !== undefined) routes[id] = route;
    questions[id] =
      type === "noul"
        ? { type, instructions, ...readNoulCriteria(input.criteria, `${where}.criteria`) }
        : type === "choice"
          ? { type, instructions, criteria: readChoiceOptions(input.criteria ?? input.options, `${where}.criteria`) }
          : { type, instructions, criteria: readScoreLevels(input.criteria ?? input.levels, `${where}.criteria`) };
    ids.push(id);
  });
  return { ids, questions, routes };
}

/**
 * Conform raw answers (native JSON or parsed model output) into the canonical answer shapes.
 * Answers that cannot be conformed are reported as issues instead of being silently repaired.
 */
export function conformAnswers(questions: Record<string, JevQuestion>, raw: unknown): ConformResult {
  const answers: Record<string, JevAnswer> = {};
  const issues: Record<string, string> = {};
  const notes: Record<string, string> = {};
  if (raw === null || typeof raw !== "object" || Array.isArray(raw)) {
    for (const id of Object.keys(questions)) issues[id] = "engine returned no answer object";
    return { answers, issues, notes };
  }
  const source = raw as Record<string, unknown>;
  for (const [id, question] of Object.entries(questions)) {
    const entry = source[id];
    if (entry === undefined || entry === null) {
      issues[id] = "engine returned no answer for this question";
      continue;
    }
    if (question.type === "noul") {
      const value = readNumberField(entry, ["noul", "probability", "p", "yes", "value"]);
      if (value === undefined || value < 0 || value > 1) {
        issues[id] = `noul answer must be a number in [0,1], got ${describeValue(entry)}`;
        continue;
      }
      answers[id] = { type: "noul", noul: value };
      continue;
    }
    if (question.type === "choice") {
      const options = Object.keys(question.criteria);
      const rawChoice = readStringField(entry, ["choice", "selected", "option", "label", "answer", "value"]);
      const choice = rawChoice === undefined ? undefined : matchOption(options, rawChoice);
      if (choice === undefined) {
        issues[id] =
          rawChoice === undefined
            ? `choice answer is missing a "choice" field`
            : `choice "${rawChoice}" is not one of the criteria options: ${options.join(", ")}`;
        continue;
      }
      const distribution = readDistribution(options, readProbabilityField(entry));
      if (distribution === undefined) {
        issues[id] = `choice probabilities must cover every option (${options.join(", ")}) with numbers >= 0`;
        continue;
      }
      if (distribution.sum <= 0) {
        issues[id] = "choice probabilities sum to 0";
        continue;
      }
      if (Math.abs(distribution.sum - 1) > 1e-6) {
        notes[id] = `probabilities renormalized (sum was ${round(distribution.sum, 6)})`;
      }
      answers[id] = {
        type: "choice",
        choice,
        probabilities: scale(distribution.values, distribution.sum),
        ...readConfidence(entry, notes, id),
      };
      continue;
    }
    const levels = question.criteria;
    const value = readNumberField(entry, ["score", "value", "rating", "level"]);
    if (value === undefined || value < 0 || value > levels.length - 1) {
      issues[id] = `score answer must be a number between 0 and ${levels.length - 1}, got ${describeValue(entry)}`;
      continue;
    }
    const keys = levels.map((_, index) => String(index));
    const distribution = readDistribution(keys, readProbabilityField(entry));
    if (distribution === undefined) {
      issues[id] = `score probabilities must cover every level index (0..${levels.length - 1}) with numbers >= 0`;
      continue;
    }
    if (distribution.sum <= 0) {
      issues[id] = "score probabilities sum to 0";
      continue;
    }
    if (Math.abs(distribution.sum - 1) > 1e-6) {
      notes[id] = `probabilities renormalized (sum was ${round(distribution.sum, 6)})`;
    }
    const legend: Record<string, string> = {};
    levels.forEach((level, index) => {
      legend[String(index)] = level;
    });
    answers[id] = {
      type: "score",
      score: value,
      legend,
      probabilities: scale(distribution.values, distribution.sum),
      ...readConfidence(entry, notes, id),
    };
  }
  return { answers, issues, notes };
}

function readConfidence(
  entry: unknown,
  notes: Record<string, string>,
  id: string,
): { confidence?: number } {
  const value = readNumberField(entry, ["confidence"]);
  if (value === undefined) return {};
  if (value < 0 || value > 1) {
    notes[id] = `${notes[id] ? `${notes[id]}; ` : ""}dropped confidence ${value} (outside [0,1])`;
    return {};
  }
  return { confidence: value };
}

function readRequiredString(raw: unknown, where: string): string {
  if (typeof raw !== "string" || raw.trim() === "") throw new JevError(`${where} must be a non-empty string.`);
  return raw.trim();
}

function readInstructions(raw: unknown, where: string): JsonValue {
  if (typeof raw === "string") {
    const trimmed = raw.trim();
    if (trimmed === "") throw new JevError(`${where} must not be empty.`);
    return trimmed;
  }
  if (Array.isArray(raw)) {
    if (raw.length === 0) throw new JevError(`${where} must not be an empty array.`);
    return raw as JsonValue;
  }
  if (raw !== null && typeof raw === "object" && Object.keys(raw as Record<string, unknown>).length > 0) {
    return raw as JsonValue;
  }
  throw new JevError(`${where} must be a non-empty string, object, or array.`);
}

function readNoulCriteria(raw: unknown, where: string): { criteria?: { true?: string; false?: string } } {
  if (raw === undefined || raw === null) return {};
  if (typeof raw !== "object" || Array.isArray(raw)) {
    throw new JevError(`${where} must be an object with optional "true" and "false" descriptions.`);
  }
  const input = raw as Record<string, unknown>;
  const criteria: { true?: string; false?: string } = {};
  const yes = input.true ?? input.yes;
  const no = input.false ?? input.no;
  if (typeof yes === "string" && yes.trim() !== "") criteria.true = yes.trim();
  if (typeof no === "string" && no.trim() !== "") criteria.false = no.trim();
  return Object.keys(criteria).length === 0 ? {} : { criteria };
}

function readChoiceOptions(raw: unknown, where: string): Record<string, string | null> {
  const options: Record<string, string | null> = {};
  if (Array.isArray(raw)) {
    for (const entry of raw) {
      if (typeof entry === "string" && entry.trim() !== "") {
        options[entry.trim()] ??= null;
        continue;
      }
      if (entry !== null && typeof entry === "object" && !Array.isArray(entry)) {
        const label = readStringField(entry, ["label", "option", "name", "value", "key"]);
        if (label === undefined) throw new JevError(`${where} entries need a label/option/name field.`);
        const description = readStringField(entry, ["description", "rubric", "detail"]);
        options[label] = description ?? null;
        continue;
      }
      throw new JevError(`${where} must contain option strings or objects.`);
    }
  } else if (raw !== null && typeof raw === "object") {
    for (const [key, value] of Object.entries(raw as Record<string, unknown>)) {
      if (key.trim() === "") throw new JevError(`${where} has an empty option key.`);
      if (value === null || value === undefined) {
        options[key] = null;
        continue;
      }
      if (typeof value !== "string") {
        throw new JevError(`${where}.${key} must be a description string or null.`);
      }
      options[key] = value.trim() === "" ? null : value.trim();
    }
  } else {
    throw new JevError(`${where} must be an object of option -> description, or an array of options.`);
  }
  if (Object.keys(options).length < 2) {
    throw new JevError(`${where} needs at least two options; a single-option choice cannot discriminate.`);
  }
  return options;
}

function readScoreLevels(raw: unknown, where: string): string[] {
  let levels: string[];
  if (Array.isArray(raw)) {
    levels = raw.map((entry) => readLevelText(entry, where));
  } else if (raw !== null && typeof raw === "object") {
    const entries = Object.entries(raw as Record<string, unknown>);
    const numericKeys = entries.length > 0 && entries.every(([key]) => /^\d+$/.test(key.trim()));
    levels = entries.map(([key, value]) => readLevelText(numericKeys ? value : key, where));
  } else {
    throw new JevError(`${where} must be an ordered array of level descriptions.`);
  }
  if (levels.length < 2) {
    throw new JevError(`${where} needs at least two levels.`);
  }
  return levels;
}

function readLevelText(entry: unknown, where: string): string {
  if (typeof entry === "string" && entry.trim() !== "") return entry.trim();
  if (entry !== null && typeof entry === "object" && !Array.isArray(entry)) {
    const text = readStringField(entry, ["description", "label", "text", "name", "value"]);
    if (text !== undefined) return text;
  }
  throw new JevError(`${where} level descriptions must be non-empty strings.`);
}

function readNumberField(source: unknown, keys: string[]): number | undefined {
  if (typeof source === "number" && Number.isFinite(source)) return source;
  if (source === null || typeof source !== "object" || Array.isArray(source)) return undefined;
  const record = source as Record<string, unknown>;
  for (const key of keys) {
    const value = record[key];
    if (typeof value === "number" && Number.isFinite(value)) return value;
  }
  return undefined;
}

function readStringField(source: unknown, keys: string[]): string | undefined {
  if (source === null || typeof source !== "object" || Array.isArray(source)) return undefined;
  const record = source as Record<string, unknown>;
  for (const key of keys) {
    const value = record[key];
    if (typeof value === "string" && value.trim() !== "") return value.trim();
  }
  return undefined;
}

function readProbabilityField(entry: unknown): unknown {
  if (entry === null || typeof entry !== "object" || Array.isArray(entry)) return undefined;
  const record = entry as Record<string, unknown>;
  return record.probabilities ?? record.probability ?? record.distribution ?? record.scores;
}

function matchOption(options: string[], value: string): string | undefined {
  const exact = options.find((option) => option === value);
  if (exact !== undefined) return exact;
  const needle = value.trim().toLowerCase();
  return options.find((option) => option.trim().toLowerCase() === needle);
}

/** Read `{ key: number }` (or an array aligned with `keys`) and report the raw sum. */
function readDistribution(keys: string[], raw: unknown): { values: Record<string, number>; sum: number } | undefined {
  if (Array.isArray(raw)) {
    if (raw.length !== keys.length) return undefined;
    const values: Record<string, number> = {};
    let sum = 0;
    keys.forEach((key, index) => {
      const value = raw[index];
      if (typeof value !== "number" || !Number.isFinite(value) || value < 0) return;
      values[key] = value;
      sum += value;
    });
    return Object.keys(values).length === keys.length ? { values, sum } : undefined;
  }
  if (raw === null || typeof raw !== "object") return undefined;
  const source = raw as Record<string, unknown>;
  const values: Record<string, number> = {};
  let sum = 0;
  for (const key of keys) {
    const value = source[key];
    if (typeof value !== "number" || !Number.isFinite(value) || value < 0) return undefined;
    values[key] = value;
    sum += value;
  }
  return { values, sum };
}

function scale(values: Record<string, number>, sum: number): Record<string, number> {
  const scaled: Record<string, number> = {};
  for (const [key, value] of Object.entries(values)) scaled[key] = round(value / sum, 6);
  return scaled;
}

function round(value: number, digits: number): number {
  const factor = 10 ** digits;
  return Math.round(value * factor) / factor;
}

function describeValue(value: unknown): string {
  if (typeof value === "string") return JSON.stringify(value);
  if (typeof value === "number" || typeof value === "boolean") return String(value);
  return value === null ? "null" : JSON.stringify(value).slice(0, 200);
}