/**
 * Copyright (c) 2026, Salesforce, Inc.,
 * All rights reserved.
 * For full license text, see the LICENSE.txt file
 */

import { z } from "zod";
import { stripControlChars, stripLineSeparators } from "../lib/control-chars.js";
import { DOTTED_GRAPHQL_NAME_RE, GRAPHQL_NAME_RE } from "../lib/graphql-name.js";

/**
 * Some MCP clients/models JSON-stringify complex tool arguments (e.g. send
 * `"25"` for a number or `'{"x":1}'` for an object). This coerces such a
 * string back to its parsed value so the real value validates — lossless,
 * and scoped to strings that LOOK like JSON (object/array/number) so `$var`
 * placeholders and enum strings (e.g. scope "MINE") are left untouched.
 * Same "tolerate predictable client encoding quirks" philosophy as the
 * detailOrderBy array-collapse shim.
 */
export const coerceJsonArg = (v: unknown): unknown => {
	const once = coerceJsonOnce(v);
	// Unwrapping a double-quoted string (e.g. '"25"' → "25", '"$first"' → "$first")
	// can yield another coercible string; coerce one more level so '"25"' → 25.
	if (typeof once === "string" && once !== v) return coerceJsonOnce(once);
	return once;
};

const coerceJsonOnce = (v: unknown): unknown => {
	if (typeof v !== "string") return v;
	const t = v.trim();
	// object/array, bare number, OR a JSON-quoted string literal
	if (!(t.startsWith("{") || t.startsWith("[") || /^-?\d+(\.\d+)?$/.test(t) || t.startsWith('"')))
		return v;
	try {
		return JSON.parse(t);
	} catch {
		return v;
	}
};

/** Wrap a schema so a JSON-stringified value is coerced before validation. */
export const jsonCoercible = (schema: z.ZodTypeAny) => z.preprocess(coerceJsonArg, schema);

/**
 * Wrap an enum (or any string schema) so Unicode control/format chars are
 * STRIPPED from a string value before validation (W-23336443). Motivation: a
 * `z.enum` rejection is reflected VERBATIM by the MCP SDK's input validation,
 * which runs UPSTREAM of `runTool` — so the adapter's `neutralizeControlChars`
 * never executes, and a poisoned value like `"describe_object‮"` reaches
 * the host raw inside `received '…'`. `JSON.stringify` (which the SDK uses)
 * escapes only C0, so DEL and the entire Cf class (bidi overrides, zero-width)
 * survive. Stripping here closes that channel two ways: a control-char-poisoned
 * but otherwise-valid value strips to the valid enum member (accepted, no
 * message), and a genuinely-invalid value rejects with a message free of raw
 * control chars.
 *
 * This upstream path also strips U+2028/U+2029 (via {@link stripLineSeparators}),
 * matching what every envelope sink does in `tool-adapter.ts`: those separators
 * are NOT Cc/Cf so `stripControlChars` correctly ignores them, but left raw in
 * host-visible text they trip a Claude.AI 408 (MCP TS SDK #2155) — and this
 * rejection message reaches the host BEFORE the envelope's own line-separator
 * strip can run, so it must strip them itself.
 *
 * Why `z.preprocess` and not `.refine`: preprocess can TRANSFORM (strip) the
 * value, and — verified against the SDK's zod-to-json-schema converter — the
 * wrapper PRESERVES the published JSON-Schema `enum` and `description`, so the
 * advertised allowed-value list the LLM sees is unchanged. `.refine` can only
 * reject, not strip.
 *
 * CAVEAT: applying this to a `z.discriminatedUnion` discriminator member is
 * pointless, though not harmful. It is NOT harmful because zod's discriminator-
 * map builder recurses through a `z.preprocess` wrapper (`getDiscriminator`
 * reads `ZodEffects.innerType()`), so construction does NOT throw and clean
 * values still route correctly (verified against zod 3.25.76 on AGGREGATE_INPUT).
 * It is POINTLESS because discrimination reads the RAW `ctx.data[discriminator]`
 * to pick a branch BEFORE that branch's preprocess ever runs — so the strip
 * cannot influence branch selection. A discriminator that matches no branch
 * raises `invalid_union_discriminator`, whose issue lists only the EXPECTED
 * options and does NOT echo the received value — so, unlike a plain
 * `invalid_enum_value`, there is no verbatim-reflection channel there to close.
 * The aggregate `function` discriminators are therefore left un-wrapped per
 * W-23336443; only plain enums (invalid_enum_value DOES echo `received`) need it.
 */
export const enumStripControlChars = <T extends z.ZodTypeAny>(schema: T) =>
	z.preprocess(
		(v) => (typeof v === "string" ? stripLineSeparators(stripControlChars(v)) : v),
		schema,
	);

/** A string that looks like a JSON object/array literal — used so the advertised
 *  schema accepts a stringified object/array; coerceJsonArg then parses it. A
 *  `{`/`[`-prefixed but invalid-JSON string is accepted as a literal (rare; the
 *  model reliably sends valid JSON) rather than erroring. */
export const jsonLiteralString = () => z.string().regex(/^\s*[{[]/);
/** A string of digits — a stringified positive integer for `first`. */
export const intLiteralString = () => z.string().regex(/^\d+$/);
/** A JSON-quoted string literal, e.g. "\"$first\"" or "\"25\"" — some models
 *  double-encode args. coerceJsonArg unwraps it before validation. */
export const quotedString = () => z.string().regex(/^".*"$/s);

/**
 * Zod string validator that enforces the GraphQL Name production. Returned
 * schema is `.describe()`-tagged with the supplied description so it appears
 * in the MCP tool's advertised JSON Schema.
 */
export const graphqlName = (description: string): z.ZodString =>
	z.string().regex(GRAPHQL_NAME_RE, "must be a valid GraphQL Name").describe(description);

/**
 * Zod validator for a dotted field path (`Id`, `Owner.Name`, …). Each
 * `.`-separated segment must be a valid GraphQL Name. Applied to every
 * caller-supplied field-selection input (`returnFields`, `fields`,
 * `parentFields`, child-relationship `fields`) so a selection-set breakout is
 * rejected at the MCP boundary (W-22735537); the builders carry the matching
 * `assertDottedGraphqlName` guard for direct (CLI / eval) callers.
 */
export const dottedGraphqlName = (description: string): z.ZodString =>
	z
		.string()
		.regex(DOTTED_GRAPHQL_NAME_RE, "must be a valid field path (dot-separated GraphQL Names)")
		.describe(description);

/**
 * `scope` argument validator: a bare enum token (e.g. `MINE`, `EVERYTHING`) or a
 * `$varName` placeholder. `scope` renders into an argument position, so an
 * unconstrained string is a selection-set / argument breakout (W-22735537).
 */
const SCOPE_RE = /^\$?[A-Za-z_][A-Za-z0-9_]*$/;
export const scopeArg = (description: string): z.ZodString =>
	z.string().regex(SCOPE_RE, "must be a Scope enum token or $varName").describe(description);

/**
 * Org alias / username charset + length, matching `auth.ts`'s
 * `assertValidOrgAlias`. Defense-in-depth at the MCP boundary so the
 * advertised JSON Schema teaches the LLM the constraint upfront, while
 * `assertValidOrgAlias` still gates the shell-out path.
 */
const ORG_ALIAS_RE = /^[A-Za-z0-9_][A-Za-z0-9_.+@-]{0,252}$/;
export const orgAlias = (description: string): z.ZodString =>
	z.string().regex(ORG_ALIAS_RE, "must be a valid org alias or username").describe(description);

/**
 * `<Object>_OrderBy` shape. A FACTORY (not a shared instance): each call returns
 * a fresh schema so the MCP SDK's zod-to-json-schema converter INLINES it at every
 * use site instead of collapsing reuse into cross-`$ref`s (which intermittently
 * dropped union branches in the advertised schema).
 */
export const orderByObject = () => z.record(z.unknown());

/**
 * A bare top-level `$varName` string standing in for an entire `filter` /
 * `orderBy` / `first` argument. A FACTORY (not a shared instance) so the advertised
 * JSON Schema inlines it at every use site rather than emitting cross-`$ref`s.
 */
export const varPlaceholder = () =>
	z.string().regex(/^\$[A-Za-z_]\w*$/, "must be a $variable placeholder, e.g. $filter");

/**
 * Build a `childRelationships[]` element schema. The `orderBy` schema is a
 * parameter because tools differ on whether they advertise the array shape
 * (`sf_gql_list`: union; `sf_gql_detail`: singleton with array-collapse
 * preprocess). Everything else — `relationshipName`, `fields`, `first`,
 * `filter` — is identical across callers.
 *
 * `relationshipName` flows into a rendered field-path position in the
 * resulting query, so it carries the same GraphQL Name constraint as
 * top-level `object` — same failure mode if it's invalid.
 */
export const childRelationshipSchema = (orderBy: z.ZodTypeAny) =>
	z.object({
		relationshipName: graphqlName(
			'Child relationship API name, e.g. "Contacts", "Opportunities". Must be a valid GraphQL Name.',
		),
		fields: z.array(
			dottedGraphqlName(
				'Scalar field API names on the child record; dot-paths like "Owner.Name" allowed.',
			),
		),
		first: jsonCoercible(
			z.union([z.number().int().positive(), varPlaceholder(), intLiteralString(), quotedString()]),
		).optional(),
		filter: jsonCoercible(
			z.union([z.record(z.unknown()), varPlaceholder(), jsonLiteralString(), quotedString()]),
		).optional(),
		orderBy: jsonCoercible(
			z.union([orderBy, varPlaceholder(), jsonLiteralString(), quotedString()]),
		).optional(),
	});