/**
 * pi extension: TypeSafe Jev (and chat-model stand-ins) as a single decision tool.
 *
 * Registers:
 *   - `decide` tool   — state + typed questions (noul/choice/score) -> typed answers
 *   - `/decide`       — status | add | set | unset | question | models | setup | help  (alias: `/jev`)
 *
 * Wiring only: the dispatch lives in lib/dispatch.ts and the command surface in lib/commands.ts,
 * so both stay testable without a harness. Every call picks a backend, and both the call and each
 * individual question may override it, so one request can combine several models.
 */

import type { ExtensionAPI, ExtensionCommandContext } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import {
  BACKEND_IDS,
  describeBackends,
  resolveConfig,
  selectBackend,
  type BackendId,
} from "./lib/config.ts";
import { createCommandHandler, type CommandContext } from "./lib/commands.ts";
import { buildGroups, runGroup, normalizeState, type GroupOutcome } from "./lib/dispatch.ts";
import { JevError, messageOf } from "./lib/errors.ts";
import { formatAnswers, type AnswerReport, type GroupMeta } from "./lib/format.ts";
import { resolveAgentDirWithHarness } from "./lib/harness.ts";
import { normalizeQuestions, type JevAnswer, type JevQuestion, type JsonValue } from "./lib/questions.ts";
import { stringEnum } from "./lib/schema.ts";
import { toPiUsage, type TokenCounts } from "./lib/usage.ts";

const QUESTION_TYPES = ["noul", "choice", "score"] as const;

const QUESTION_SCHEMA = Type.Object({
  id: Type.String({ description: "Answer key for this question; the answer comes back under the same id." }),
  type: stringEnum(QUESTION_TYPES),
  instructions: Type.Union([Type.String(), Type.Object({}, { additionalProperties: true }), Type.Array(Type.Any())], {
    description:
      "One narrow, coherent judgment. A string for plain questions, or structured JSON when definitions, contrasts, or examples clarify it.",
  }),
  criteria: Type.Optional(
    Type.Union(
      [
        Type.Record(Type.String(), Type.Union([Type.String(), Type.Null()])),
        Type.Array(Type.Union([Type.String(), Type.Object({}, { additionalProperties: true })])),
      ],
      {
        description:
          "noul: {true?, false?} descriptions. choice: {option: rubric|null} (or an array of options), at least two. score: ordered level descriptions, lowest first, at least two.",
      },
    ),
  ),
  backend: Type.Optional(stringEnum(BACKEND_IDS)),
  model: Type.Optional(Type.String({ description: "Route only this question to a specific model id." })),
});

const TOOL_PARAMETERS = Type.Object({
  state: Type.Union([Type.String(), Type.Object({}, { additionalProperties: true }), Type.Array(Type.Any())], {
    description:
      "Everything the judgments need: source text, records, policies, current facts. A string, or structured JSON (a JSON-encoded string is also accepted).",
  }),
  questions: Type.Array(QUESTION_SCHEMA, {
    description:
      "One entry per independent judgment. type: noul = P(yes) in [0,1]; choice = one option from criteria plus a full distribution; score = graded position across ordered levels. Ask independent questions together — they are answered in parallel.",
  }),
  backend: Type.Optional(
    stringEnum(
      BACKEND_IDS,
      "Backend for every question that does not set its own. Defaults to the configured backend. Requests to different backends run in parallel.",
    ),
  ),
  model: Type.Optional(
    Type.String({ description: "Model id for every question that does not set its own; overrides the backend's configured model." }),
  ),
});

interface DeciderToolDetails {
  status: "asking" | "answered";
  groups?: GroupMeta[];
  questions?: Record<string, JevQuestion>;
  answers?: Record<string, JevAnswer>;
  issues?: Record<string, string>;
  notes?: Record<string, string>;
  costKnown?: boolean;
  message?: string;
}



export default async function deciderExtension(pi: ExtensionAPI) {
  // Resolved once at load time so every tool call and command share one config location.
  const agentDir = await resolveAgentDirWithHarness();

  pi.registerTool({
    name: "decide",
    label: "Decide (Jev / decision backends)",
    description:
      "Ask a decision backend (TypeSafe Jev, OpenRouter's Decisions API, or a chat-model proxy) for calibrated typed " +
      "judgments over state you provide: a yes/no probability (noul), one option out of a defined set with a full " +
      "probability distribution (choice), or a graded score across ordered levels (score). Use it for routing, " +
      "ranking, extraction-verification, moderation, and any decision where a probability beats prose. Code keeps " +
      "control: answers come back with probabilities and confidence you can threshold. Each question may name its own " +
      "backend/model, so one call can combine several. " +
      "Configured by <agentDir>/decider.json — run /decide setup, or set TYPESAFE_API_KEY (typesafe), " +
      "OPENROUTER_API_KEY (openrouter), or DECIDER_LLM_API_KEY plus a chat model id (llm). " +
      "Load the `decide` skill for question-design and shape guidance, including the gate and composite patterns.",
    promptSnippet: "Ask Jev or another decision backend for calibrated typed judgments (noul/choice/score) over supplied state",
    promptGuidelines: [
      "Use decide when a decision needs a calibrated probability, a pick among named options, or a graded score instead of a general-purpose model's prose.",
      "Batch independent decide questions into one call; keep each question narrow and put all needed evidence in state.",
      "Route individual decide questions with their own backend/model only when a second opinion or a cheaper model is genuinely useful; unset questions use the call-level backend.",
      "Give every decide choice question a no-match option, and threshold the returned probabilities/confidence in code rather than in the prompt.",
    ],
    parameters: TOOL_PARAMETERS,
    async execute(_toolCallId, params, signal, onUpdate, _ctx) {
      const cfg = resolveConfig({ agentDir });
      const { ids, questions, routes } = normalizeQuestions(params.questions);
      const state = normalizeState(params.state);
      const { groups, failures } = buildGroups(cfg, ids, routes, params.backend, params.model);

      onUpdate?.({
        content: [
          {
            type: "text",
            text: `Asking ${ids.length} question(s) across ${groups.length} backend batch(es): ${groups
              .map((group) => `${group.backend.id}/${group.model} (${group.ids.length})`)
              .join(", ")}`,
          },
        ],
        details: { status: "asking" } satisfies DeciderToolDetails,
      });

      const settled = await Promise.all(
        groups.map(async (group) => {
          try {
            return { group, outcome: await runGroup(group, state, questions, signal), error: undefined };
          } catch (error) {
            return { group, outcome: undefined, error };
          }
        }),
      );

      const answers: Record<string, JevAnswer> = {};
      const issues: Record<string, string> = { ...failures };
      const notes: Record<string, string> = {};
      const metas: GroupMeta[] = [];
      const tokens: TokenCounts = { input: 0, output: 0 };
      let costUsd = 0;
      let costKnown = Object.keys(failures).length === 0;

      for (const { group, outcome, error } of settled) {
        if (outcome === undefined) {
          for (const id of group.ids) issues[id] = `${group.backend.id} backend failed: ${messageOf(error)}`;
          costKnown = false;
          continue;
        }
        Object.assign(answers, outcome.conform.answers);
        Object.assign(issues, outcome.conform.issues);
        Object.assign(notes, outcome.conform.notes);
        metas.push(outcome.meta);
        tokens.input = (tokens.input ?? 0) + (outcome.tokens.input ?? 0);
        tokens.output = (tokens.output ?? 0) + (outcome.tokens.output ?? 0);
        if (outcome.costUsd === undefined) costKnown = false;
        else costUsd += outcome.costUsd;
      }

      if (Object.keys(answers).length === 0) {
        const detail = Object.entries(issues)
          .map(([id, issue]) => `${id}: ${issue}`)
          .join("; ");
        throw new JevError(
          `No question could be answered: ${detail}`,
          groups.length === 0 ? { hint: describeBackends(cfg) } : {},
        );
      }

      const report: AnswerReport = { ids, questions, answers, issues, notes };
      const text = formatAnswers(report, metas);
      const failed = Object.keys(issues).length;
      return {
        content: [
          { type: "text", text },
          ...(failed > 0 ? [{ type: "text" as const, text: `${failed} question(s) failed; see the lines marked ERROR.` }] : []),
        ],
        details: {
          status: "answered",
          groups: metas,
          questions,
          answers,
          issues,
          notes,
          costKnown,
        } satisfies DeciderToolDetails,
        usage: toPiUsage(tokens, costKnown ? costUsd : undefined),
      };
    },
  });

  // The command surface lives in lib/commands.ts so it is testable without a harness.
  const handle = createCommandHandler(agentDir);

  pi.registerCommand("decide", {
    description: "Decision backends: status, add, set, unset, question, models, setup",
    getArgumentCompletions: (prefix) =>
      ["status", "add", "set", "unset", "question", "models", "setup", "help"]
        .filter((name) => name.startsWith(prefix))
        .map((name) => ({ value: name, label: name })),
    handler: handle,
  });
  pi.registerCommand("jev", {
    description: "Alias of /decide",
    handler: handle,
  });
}