import { z } from "zod";
import { MAX_REDACTION_RULES, validateRedactionRule } from "./redaction.ts";
import { MAX_EXTRA_SCORES } from "../catalog/benchmark-feeds.ts";

/**
 * Input schema for `$AUTO_MODEL_ROUTER_HOME/config.yml`: a deep partial of
 * `RouterConfig` — every field is optional so the file carries only
 * overrides. Nested objects are strict so a misspelled key fails loudly
 * instead of being silently dropped.
 *
 * Array fields (`profiles`, `probeTiers`, `pin`, `allow`, `deny`,
 * `triggers`) REPLACE the default wholesale on merge; consequently a
 * `profiles` entry must be complete — there is nothing to merge it with.
 */

const tier = z.enum(["trivial", "simple", "moderate", "hard"]);
const qualityAxis = z.enum(["coding", "agentic", "intelligence"]);
const logLevel = z.enum(["silent", "error", "warn", "info", "debug"]);

const server = z.strictObject({
	host: z.string().min(1).optional(),
	port: z.number().int().min(0).max(65_535).optional(),
	apiKey: z.string().optional(),
	harnessId: z.string().optional(),
	subagentProfile: z.string().optional(),
	maxConcurrentTurns: z.number().int().positive().max(1_000).optional(),
});

const openrouter = z.strictObject({
	baseUrl: z.string().min(1).optional(),
	apiKey: z.string().optional(),
	referer: z.string().optional(),
	title: z.string().optional(),
	timeoutMs: z.number().positive().optional(),
	catalogTtlMs: z.number().positive().optional(),
	catalogRefreshMs: z.number().nonnegative().optional(),
	minCreditsUsd: z.number().nonnegative().optional(),
	usagePollMs: z.number().nonnegative().optional(),
});

const ollamaRate = z.strictObject({
	input: z.number().nonnegative(),
	cachedInput: z.number().nonnegative().optional(),
	output: z.number().nonnegative(),
});

const ollama = z.strictObject({
	enabled: z.boolean().optional(),
	baseUrl: z.string().min(1).optional(),
	apiKey: z.string().optional(),
	timeoutMs: z.number().positive().optional(),
	catalogTtlMs: z.number().positive().optional(),
	includeLocal: z.boolean().optional(),
	prices: z.record(z.string(), ollamaRate).optional(),
	twins: z.record(z.string(), z.string()).optional(),
	costBias: z.number().positive().optional(),
	biasUntilUsage: z.number().min(0).max(1).optional(),
	usagePollMs: z.number().nonnegative().optional(),
	quotaCooldownMs: z.number().nonnegative().optional(),
	rateLimitCooldownMs: z.number().nonnegative().optional(),
	planCreditsUsd: z.number().nonnegative().optional(),
});

/** Ids that would collide with OpenRouter's own vendor namespaces or the built-in upstreams. */
export const RESERVED_UPSTREAM_IDS: readonly string[] = ["openrouter", "ollama", "openai", "anthropic", "google", "meta-llama", "mistralai", "x-ai", "deepseek", "qwen", "amazon", "microsoft", "cohere", "perplexity", "nvidia", "moonshotai", "z-ai", "minimax"];

const upstreamModel = z.strictObject({
	id: z.string().min(1),
	name: z.string().optional(),
	input: z.number().nonnegative().optional(),
	output: z.number().nonnegative().optional(),
	cachedInput: z.number().nonnegative().optional(),
	cacheWrite: z.number().nonnegative().optional(),
	contextLength: z.number().int().positive().optional(),
	maxCompletionTokens: z.number().int().positive().optional(),
	supportsTools: z.boolean().optional(),
	supportsReasoning: z.boolean().optional(),
	supportsToolChoice: z.boolean().optional(),
	vision: z.boolean().optional(),
	quality: z.strictObject({ intelligence: z.number().optional(), coding: z.number().optional(), agentic: z.number().optional() }).optional(),
	twin: z.string().optional(),
});

const upstream = z.strictObject({
	id: z
		.string()
		.regex(/^[a-z0-9][a-z0-9-]{0,31}$/, "lowercase letters, digits and dashes")
		.refine((id) => !RESERVED_UPSTREAM_IDS.includes(id), { message: "this id is a vendor namespace on OpenRouter or a built-in upstream; use e.g. openai-direct, azure-eu, anthropic-direct, vllm" }),
	kind: z.enum(["openai", "azure", "anthropic"]),
	enabled: z.boolean().optional(),
	baseUrl: z.string().min(1),
	apiKey: z.string().optional(),
	auth: z.enum(["api-key", "oauth-bearer"]).optional(),
	apiVersion: z.string().optional(),
	headers: z.record(z.string(), z.string()).optional(),
	timeoutMs: z.number().positive().optional(),
	rateLimitCooldownMs: z.number().nonnegative().optional(),
	quotaCooldownMs: z.number().nonnegative().optional(),
	costBias: z.number().positive().max(1).optional(),
	models: z.array(upstreamModel),
});

const axisScore = z.number().min(0).max(100).optional();

/**
 * One `benchmarks.extraScores` row, as a config FILE may write it. Strict here
 * on purpose: a typo in a file the operator edits is an error they want told,
 * exactly like every other key. The same rows arriving over `reconfigure` from a
 * front door skip this schema and are sanitised at use instead (`suppliedScores`),
 * where a bad row is dropped with a warning rather than failing the whole patch.
 */
const extraScore = z.strictObject({
	key: z.string().min(1),
	// Only ever a tie-break between two rows sharing a key; absent is normal.
	creator: z.string().default(""),
	coding: axisScore,
	intelligence: axisScore,
	agentic: axisScore,
	// Config supplies provenance; it may not claim a fetched feed or the local lane.
	source: z.enum(["neutral", "vendor"]),
});

const benchmarks = z.strictObject({
	enabled: z.boolean().optional(),
	artificialAnalysisApiKey: z.string().optional(),
	benchlm: z.boolean().optional(),
	refreshMs: z.number().nonnegative().optional(),
	timeoutMs: z.number().positive().optional(),
	useLocalScores: z.boolean().optional(),
	extraScores: z.array(extraScore).max(MAX_EXTRA_SCORES).optional(),
});

const tierConfig = z.strictObject({
	minQuality: z.number().min(0).max(100).optional(),
	maxInputPerMtok: z.number().nonnegative().optional(),
	maxOutputPerMtok: z.number().nonnegative().optional(),
	qualityExponent: z.number().nonnegative().optional(),
	qualityNormalization: z.boolean().optional(),
	capabilityFloorUsd: z.number().positive().optional(),
	pin: z.array(z.string()).optional(),
});

const taskConfig = z.strictObject({
	axis: qualityAxis.optional(),
	minQuality: z.number().min(0).max(100).optional(),
	requireImage: z.boolean().optional(),
	prefer: z.array(z.string()).optional(),
});

const filters = z.strictObject({
	allow: z.array(z.string()).optional(),
	deny: z.array(z.string()).optional(),
	includeFree: z.boolean().optional(),
	requireToolSupport: z.boolean().optional(),
	agenticAxisForToolTurns: z.boolean().optional(),
	minAgenticForToolTurns: z.number().min(0).max(100).optional(),
	minTrust: z.number().min(0).max(1).optional(),
	feedbackWeight: z.number().nonnegative().optional(),
	feedbackByTask: z.boolean().optional(),
	minTrustSamples: z.number().int().nonnegative().optional(),
	trustScopedByHarness: z.boolean().optional(),
	trustWindowDays: z.number().nonnegative().optional(),
	contextHeadroom: z.number().positive().optional(),
	reasoningCompletionFloor: z.number().int().nonnegative().optional(),
	latencyWeight: z.number().nonnegative().optional(),
	latencyReferenceMs: z.number().positive().optional(),
	latencyReferenceTokensPerSec: z.number().positive().optional(),
	latencyMinSamples: z.number().int().nonnegative().optional(),
	cacheReliabilityMinSamples: z.number().int().nonnegative().optional(),
	maxExpectedWaitMs: z.number().positive().optional(),
	latencyWeightContinuation: z.number().nonnegative().optional(),
	escalationCostWeight: z.number().min(0).max(1).optional(),
});

const classifier = z.strictObject({
	ambiguityThreshold: z.number().min(0).max(1).optional(),
	model: z.string().min(1).optional(),
	learnedModelPath: z.string().optional(),
	maxCostFraction: z.number().min(0).max(1).optional(),
	maxCostUsd: z.number().nonnegative().optional(),
	timeoutMs: z.number().positive().optional(),
	cacheSize: z.number().int().nonnegative().optional(),
	toolAxis: qualityAxis.optional(),
	chatAxis: qualityAxis.optional(),
	agenticLoopDepth: z.number().int().nonnegative().optional(),
	mechanicalRetryFactor: z.number().min(0).max(1).optional(),
	readOnlyToolWeight: z.number().nonnegative().optional(),
	reasoningWeights: z
		.strictObject({
			medium: z.number().nonnegative().optional(),
			high: z.number().nonnegative().optional(),
			xhigh: z.number().nonnegative().optional(),
			max: z.number().nonnegative().optional(),
		})
		.optional(),
});

const escalation = z.strictObject({
	enabled: z.boolean().optional(),
	probeTokens: z.number().int().positive().optional(),
	maxHoldMs: z.number().positive().optional(),
	maxAttempts: z.number().int().positive().optional(),
	probeTiers: z.array(tier).optional(),
	triggers: z.array(z.string()).optional(),
	escalateOnLengthStop: z.boolean().optional(),
});

const hysteresis = z.strictObject({
	holdTurns: z.number().int().nonnegative().optional(),
	holdTurnsAfterEscalation: z.number().int().nonnegative().optional(),
	switchMargin: z.number().positive().optional(),
	switchHorizonTurns: z.number().int().positive().optional(),
	confirmUpgradesBelowConfidence: z.number().min(0).max(1).optional(),
	cacheWarmTtlMs: z.number().nonnegative().optional(),
	maxDowngradePerTurn: z.number().int().nonnegative().optional(),
	breakHoldOnMechanical: z.boolean().optional(),
});

const exploration = z.strictObject({
	enabled: z.boolean().optional(),
	// Spelled out per tier rather than z.record so an unknown tier name is a
	// config error instead of a silently ignored key.
	rates: z
		.strictObject({
			trivial: z.number().min(0).max(1).optional(),
			simple: z.number().min(0).max(1).optional(),
			moderate: z.number().min(0).max(1).optional(),
			hard: z.number().min(0).max(1).optional(),
		})
		.optional(),
	stickyPolicy: z.enum(["never", "cold-cache", "always"]).optional(),
	holdTurns: z
		.strictObject({
			enabled: z.boolean().optional(),
			// Non-empty and positive: an empty set or a 0 would silently disable
			// the experiment while reading as enabled.
			values: z.array(z.number().int().positive()).min(1).optional(),
		})
		.optional(),
});
const cache = z.strictObject({
	injectBreakpoints: z.boolean().optional(),
	maxBreakpoints: z.number().int().positive().optional(),
	minPromptTokens: z.number().int().nonnegative().optional(),
	milestoneTokens: z.number().int().positive().optional(),
});

const context = z.strictObject({
	enabled: z.boolean().optional(),
	baseUrl: z.string().optional(),
	token: z.string().optional(),
	defaultScope: z.string().optional(),
	timeoutMs: z.number().int().positive().optional(),
	maxStalenessMs: z.number().int().nonnegative().optional(),
	maxBlockChars: z.number().int().positive().optional(),
	memoryLimit: z.number().int().positive().optional(),
	docsLimit: z.number().int().nonnegative().optional(),
	sessionLimit: z.number().int().nonnegative().optional(),
	briefChars: z.number().int().nonnegative().optional(),
	layers: z.boolean().optional(),
	recordTurns: z.boolean().optional(),
	injectWithoutTools: z.boolean().optional(),
	maxQueue: z.number().int().positive().optional(),
});

const compaction = z.strictObject({
	enabled: z.boolean().optional(),
	budgetTokens: z.number().int().positive().optional(),
	floorRatio: z.number().positive().max(1).optional(),
	replanGrowthRatio: z.number().min(1).optional(),
	fitToWindow: z.boolean().optional(),
	protectRecentTurns: z.number().int().positive().optional(),
	maxToolResultBytes: z.number().int().positive().optional(),
	keepHeadBytes: z.number().int().nonnegative().optional(),
	keepTailBytes: z.number().int().nonnegative().optional(),
	elideSupersededReads: z.boolean().optional(),
	collapseDuplicateResults: z.boolean().optional(),
	digestToolResults: z.boolean().optional(),
	digestMaxPerTurn: z.number().int().nonnegative().optional(),
});

const budget = z.strictObject({
	perTurnUsd: z.number().nonnegative().optional(),
	perConversationUsd: z.number().nonnegative().optional(),
	perDayUsd: z.number().nonnegative().optional(),
	perMonthUsd: z.number().nonnegative().optional(),
	onExceeded: z.enum(["downgrade", "reject"]).optional(),
});

const fallbackBlend = z.strictObject({
	inputPerMtok: z.number().nonnegative().optional(),
	outputPerMtok: z.number().nonnegative().optional(),
});

const ledger = z.strictObject({
	path: z.string().optional(),
	blendWindowDays: z.number().positive().optional(),
	blendMinSamples: z.number().int().nonnegative().optional(),
	fallbackBlend: fallbackBlend.optional(),
	conversationTtlMs: z.number().positive().optional(),
	// `null` and 0 both mean "keep everything"; null is the default.
	retentionDays: z.number().int().nonnegative().nullable().optional(),
});

/**
 * Redaction rules are compiled at load, and a pattern that cannot be run
 * safely is rejected HERE, with the reason, rather than at the first turn:
 * a rule an operator believes is removing something must never be a rule the
 * router quietly skipped.
 */
const redactionRule = z.strictObject({
	name: z.string().min(1),
	pattern: z.string().min(1),
	replacement: z.string().optional(),
}).superRefine((rule, ctx) => {
	const reason = validateRedactionRule(rule);
	if (reason !== null) ctx.addIssue({ code: "custom", message: reason });
});

const redaction = z.strictObject({
	enabled: z.boolean().optional(),
	rules: z.array(redactionRule).max(MAX_REDACTION_RULES).optional(),
	scanTools: z.boolean().optional(),
});

// Complete entries: arrays replace wholesale, so a partial profile would
// produce an invalid resolved config.
const profile = z.strictObject({
	id: z.string().min(1),
	name: z.string().min(1),
	maxTier: tier,
	minTier: tier,
	contextWindow: z.number().int().positive(),
	maxTokens: z.number().int().positive(),
	budget: budget.optional(),
});

export const configInputSchema = z.strictObject({
	server: server.optional(),
	openrouter: openrouter.optional(),
	benchmarks: benchmarks.optional(),
	tiers: z
		.strictObject({
			trivial: tierConfig.optional(),
			simple: tierConfig.optional(),
			moderate: tierConfig.optional(),
			hard: tierConfig.optional(),
		})
		.optional(),
	tasks: z
		.strictObject({
			coding: taskConfig.optional(),
			vision: taskConfig.optional(),
			documentation: taskConfig.optional(),
			data: taskConfig.optional(),
			chat: taskConfig.optional(),
		})
		.optional(),
	ollama: ollama.optional(),
	upstreams: z.array(upstream).optional(),
	filters: filters.optional(),
	classifier: classifier.optional(),
	escalation: escalation.optional(),
	hysteresis: hysteresis.optional(),
	exploration: exploration.optional(),
	cache: cache.optional(),
	context: context.optional(),
	compaction: compaction.optional(),
	budget: budget.optional(),
	profiles: z.array(profile).optional(),
	report: z.strictObject({ baselines: z.array(z.string()).optional(), dailySummary: z.boolean().optional() }).optional(),
	anthropic: z.object({ models: z.record(z.string(), z.string()).optional() }).strict().optional(),
	harnessSwitch: z
		.strictObject({
			enabled: z.boolean().optional(),
			models: z.record(tier, z.string().regex(/^[^/]+\/.+$/, "provider/id")).optional(),
			minConfidence: z.number().min(0).max(1).optional(),
		})
		.optional(),
	digest: z
		.strictObject({
			enabled: z.boolean().optional(),
			minBytes: z.number().int().nonnegative().optional(),
			maxBytes: z.number().int().positive().optional(),
			tools: z.array(z.string()).optional(),
			fromTier: tier.optional(),
			tier: tier.optional(),
			model: z.string().optional(),
			maxOutputTokens: z.number().int().positive().optional(),
			maxCostUsd: z.number().nonnegative().optional(),
			timeoutMs: z.number().int().positive().optional(),
			toolAliases: z.record(z.string(), z.string()).optional(),
		})
		.optional(),
	ledger: ledger.optional(),
	redaction: redaction.optional(),
	adaptiveTierFloors: z.boolean().optional(),
	adaptivePriceCeilings: z.boolean().optional(),
	logLevel: logLevel.optional(),
});

export type ConfigInput = z.infer<typeof configInputSchema>;