/**
 * @copyright Sister Software
 * @license AGPL-3.0
 * @author Teffen Ellis, et al.
 *
 *   ONNX inference wrapper.
 *
 *   Loads a token-classification model exported by `packages/corpus-python/src/mailwoman_train/
 *   export_onnx.py` (BertForTokenClassification w/ inputs `input_ids` + `attention_mask`, output
 *   `logits` shape `[batch, sequence, num_labels]`).
 *
 *   Lazy-loads on first `infer()` call unless `warmup: true` is passed; the constructor itself is
 *   cheap and synchronous.
 */

import { promises as fs } from "node:fs"

import ort from "onnxruntime-node"

import { ANCHOR_FEATURE_DIM } from "./anchor-inference.ts"
import { COUNTRY_FEATURE_DIM } from "./country-inference.ts"
import { GAZETTEER_FEATURE_DIM, LOCALITY_SURFACE_FEATURE_DIM, STREET_TYPE_FEATURE_DIM } from "./gazetteer-inference.ts"

// Back-compat: the dims moved to gazetteer-inference.ts (browser-safe) so neural-web's runner can
// import them without touching this node-only module.
export { LOCALITY_SURFACE_FEATURE_DIM, STREET_TYPE_FEATURE_DIM } from "./gazetteer-inference.ts"

/**
 * Evidence-bundle zero-fallback widths (Option-A; must match the trained model's channel dims).
 */
/**
 * Channel width of the locality-surface evidence feature. Must match the trained model's input shape.
 */

export interface ONNXRunnerOpts {
	/**
	 * If true, load the model immediately in `create()`. Default false.
	 */
	warmup?: boolean
	/**
	 * Fixed sequence length the model expects. v0.1.0 / v0.2.0 quantization baked in 128 (the training-time max position)
	 * even though the fp32 export specified dynamic axes — re-quantize with a different shape to override. Inputs shorter
	 * than this are padded with id `0` and masked out via attention_mask=0; inputs longer are truncated.
	 */
	fixedSeqLen?: number
	/**
	 * ONNX Runtime execution providers to try, in priority order — e.g. `["cuda", "cpu"]` or `["webgpu", "cpu"]`.
	 * **Default `["cpu"]`** (unchanged behavior). GPU providers (`cuda`, `webgpu`) THROW at session-create when their
	 * runtime/driver is absent rather than soft-falling-back, so this is **guarded**: if the requested list fails to
	 * initialize, the runner retries on CPU alone. The cost of a failed GPU probe is a one-time sub-`100 ms` hit at load,
	 * so a GPU box lights up and a CPU box pays ~nothing. `cpu` is always appended if not present.
	 */
	executionProviders?: string[]
	/**
	 * Cap ONNX Runtime's INTRA-op thread pool — the threads a single operator splits its work across.
	 *
	 * Unset means ORT sizes the pool to the machine's core count. That is the right default for a server running one
	 * session over long sequences, and the wrong one for the shape this repo actually runs: short addresses, frequently
	 * several processes at once. Every CLI invocation is its own session, so N concurrent processes each claim every
	 * core, and the oversubscription surfaces as latency rather than error — measured 2026-08-03, eight concurrent
	 * `mailwoman geocode` calls took 8.75 s against a 10 s test timeout on an otherwise idle 16-core box, where one alone
	 * took 5.62 s.
	 *
	 * Set it when the caller knows it is one of many, or when sequences are short enough that thread coordination costs
	 * more than the parallelism returns.
	 */
	intraOpNumThreads?: number
}

/**
 * Default sequence length for v0.1.0 / v0.2.0 (BertConfig max_position_embeddings = 128).
 */
export const DEFAULT_FIXED_SEQ_LEN = 128

/**
 * Intra-op thread cap applied by `NeuralAddressClassifier.loadFromWeights`, overridable per-process via
 * `MAILWOMAN_INTRA_OP_THREADS`.
 *
 * THERE IS NO VALUE THAT IS RIGHT FOR BOTH REGIMES, which is why this is a knob with a compromise default rather than a
 * tuned constant. Measured on a 16-core box:
 *
 * - ONE process, 120 warm parses: 1 thread 18.3 ms/parse, 2 threads 12.5, 4 threads 9.2, ORT's all-cores default 9.3.
 *   More threads win; the parallelism is doing real work.
 * - FOUR concurrent processes, full geocode: 1 thread 32 req/s each, 2 threads 45, 4 threads 33. Fewer threads win,
 *   because N processes each sizing a pool to the machine oversubscribe it N-fold.
 *
 * Two is the compromise: it costs a single process ~35% latency against its own optimum, and buys a four-process server
 * ~36% throughput against the single-process optimum applied blindly. A server that knows its own worker count should
 * set `MAILWOMAN_INTRA_OP_THREADS` to roughly cores/workers instead of accepting this.
 *
 * Re-derive both curves before changing it. They are properties of the model and the box, and the single-process one
 * alone will point at the wrong answer.
 */
export const DEFAULT_INTRA_OP_THREADS = 2

export interface InferResult {
	/**
	 * Logits per token per label, indexed as `logits[tokenIdx][labelIdx]`.
	 */
	logits: number[][]
	/**
	 * Number of label classes (the inner-dim of the logits tensor).
	 */
	numLabels: number
	/**
	 * Pooled locale-head posterior (`locale_logits` output, LOCALE_COUNTRIES order), when the model exports it (v1.1.0+,
	 * #511 Tier A). Absent on older bundles — consumers must treat undefined as "no address-system detection available".
	 */
	localeLogits?: number[]
	/**
	 * #727 stage-2: per-span type scores from the semi-Markov span head (`span_scores` output, v3.x+). Indexed
	 * `spanScores[tokenIdx][lengthIdx][segmentTypeIdx]` — the segment starting at `tokenIdx`, of length `lengthIdx + 1`
	 * tokens, typed `SEGMENT_TYPES[segmentTypeIdx]` (that axis ships in the weights bundle's `semi-crf-transitions.json`,
	 * never hardcoded — the PLACETYPE_ORDER class).
	 *
	 * Absent on every pre-v3 bundle, so consumers MUST treat undefined as "no span decode available" and fall back to the
	 * BIO path. Fetching it costs ~0.75 ms (CPU, S=128); a runtime that never reads it pays nothing (ORT prunes the
	 * unfetched branch) — measured in `docs/articles/evals/2026-07-15-v301-phase2-export.md`.
	 */
	spanScores?: number[][][]
	/**
	 * Max span length (the `L` axis of {@link spanScores}). Absent iff `spanScores` is.
	 */
	maxSpan?: number
}

export class ONNXRunner {
	private session: ort.InferenceSession | null = null
	private loadPromise: Promise<ort.InferenceSession> | null = null
	public readonly fixedSeqLen: number

	private readonly executionProviders: string[]
	private readonly intraOpNumThreads: number | undefined
	private readonly modelPath: string
	private readonly modelBytes: Uint8Array | null

	private constructor(modelPath: string, modelBytes: Uint8Array | null, opts: ONNXRunnerOpts) {
		this.modelPath = modelPath
		this.modelBytes = modelBytes
		this.fixedSeqLen = opts.fixedSeqLen ?? DEFAULT_FIXED_SEQ_LEN
		const requested = opts.executionProviders ?? ["cpu"]
		// CPU is the universal final fallback — append it so a GPU-only list still has somewhere to land.
		this.executionProviders = requested.includes("cpu") ? requested : [...requested, "cpu"]
		this.intraOpNumThreads = opts.intraOpNumThreads
	}

	/**
	 * Load by path. Reads the model lazily unless `warmup` is true.
	 */
	static async create(modelPath: string, opts: ONNXRunnerOpts = {}): Promise<ONNXRunner> {
		const runner = new ONNXRunner(modelPath, null, opts)

		if (opts.warmup) {
			await runner.ensureSession()
		}

		return runner
	}

	/**
	 * Load from an already-read byte buffer.
	 */
	static async fromBytes(modelBytes: Uint8Array, opts: ONNXRunnerOpts = {}): Promise<ONNXRunner> {
		const runner = new ONNXRunner("(bytes)", modelBytes, opts)

		if (opts.warmup) {
			await runner.ensureSession()
		}

		return runner
	}

	private async ensureSession(): Promise<ort.InferenceSession> {
		if (this.session) return this.session

		if (!this.loadPromise) {
			this.loadPromise = (async () => {
				const bytes = this.modelBytes ?? new Uint8Array(await fs.readFile(this.modelPath))
				this.session = await this.createSession(bytes)

				return this.session
			})()
		}

		return this.loadPromise
	}

	/**
	 * Create the session on the configured execution providers, guarded: GPU providers (`cuda`/`webgpu`) throw at
	 * create-time when their runtime/driver is missing, so on failure we retry on CPU alone. A box with the GPU runtime
	 * uses it; a box without one transparently lands on CPU.
	 */
	private async createSession(bytes: Uint8Array): Promise<ort.InferenceSession> {
		try {
			return await ort.InferenceSession.create(bytes, {
				executionProviders: this.executionProviders,
				graphOptimizationLevel: "all",
				...(this.intraOpNumThreads ? { intraOpNumThreads: this.intraOpNumThreads } : {}),
			})
		} catch (error) {
			if (this.executionProviders.length === 1 && this.executionProviders[0] === "cpu") throw error

			// A requested GPU provider failed to initialize — fall back to CPU so inference still loads.
			console.warn(
				`[ONNXRunner] execution providers [${this.executionProviders.join(", ")}] failed to initialize ` +
					// oxlint-disable-next-line mailwoman/prefer-spliterator -- In-memory error message; only its first line is logged.
					`(${(error as Error).message.split("\n")[0]}); falling back to CPU.`
			)

			return ort.InferenceSession.create(bytes, {
				executionProviders: ["cpu"],
				graphOptimizationLevel: "all",
				...(this.intraOpNumThreads ? { intraOpNumThreads: this.intraOpNumThreads } : {}),
			})
		}
	}

	/**
	 * Run inference on a single token id sequence.
	 *
	 * Pads to `fixedSeqLen` (default 128) with id 0 + mask 0; truncates if longer. Output is sliced back to the actual
	 * input length.
	 *
	 * @param tokenIDs The id sequence produced by the tokenizer (no special tokens added).
	 * @param anchor Optional postcode-anchor channel (#239/#240). When supplied (only for anchor models — exported with
	 *   the `anchor_features`/`anchor_confidence` inputs), per-piece features `(seqLen × dim)` + confidence `(seqLen,)`
	 *   are fed, zero-padded to `fixedSeqLen`. Omit for plain models, whose ONNX has no anchor inputs.
	 */
	async infer(
		tokenIDs: number[],
		anchor?: { features: ReadonlyArray<ReadonlyArray<number>>; confidence: ReadonlyArray<number> },
		gazetteer?: { features: ReadonlyArray<ReadonlyArray<number>>; confidence: ReadonlyArray<number> },
		country?: { features: ReadonlyArray<ReadonlyArray<number>>; confidence: ReadonlyArray<number> },
		evidence?: {
			streetType?: { features: ReadonlyArray<ReadonlyArray<number>>; confidence: ReadonlyArray<number> }
			localitySurface?: { features: ReadonlyArray<ReadonlyArray<number>>; confidence: ReadonlyArray<number> }
		}
	): Promise<InferResult> {
		const session = await this.ensureSession()
		const seqLen = Math.min(tokenIDs.length, this.fixedSeqLen)
		const padded = new BigInt64Array(this.fixedSeqLen)
		const mask = new BigInt64Array(this.fixedSeqLen)

		for (let i = 0; i < seqLen; i++) {
			padded[i] = BigInt(tokenIDs[i]!)
			mask[i] = 1n
		}

		const feeds: Record<string, ort.Tensor> = {
			input_ids: new ort.Tensor("int64", padded, [1, this.fixedSeqLen]),
			attention_mask: new ort.Tensor("int64", mask, [1, this.fixedSeqLen]),
		}

		if (anchor) {
			const dim = anchor.features[0]?.length ?? 0
			const af = new Float32Array(this.fixedSeqLen * dim)
			const ac = new Float32Array(this.fixedSeqLen)

			for (let i = 0; i < seqLen; i++) {
				ac[i] = anchor.confidence[i] ?? 0
				const row = anchor.features[i]

				if (row) {
					for (let d = 0; d < dim; d++) {
						af[i * dim + d] = row[d] ?? 0
					}
				}
			}

			feeds.anchor_features = new ort.Tensor("float32", af, [1, this.fixedSeqLen, dim])
			feeds.anchor_confidence = new ort.Tensor("float32", ac, [1, this.fixedSeqLen])
		} else if (session.inputNames.includes("anchor_features")) {
			// Anchor-trained model (its ONNX declares the anchor inputs as mandatory) but no anchor data
			// was supplied: feed zeros. That's the `confidence = 0` identity — the model's anchor-off
			// behavior. Without it the session throws on the missing required inputs.
			feeds.anchor_features = new ort.Tensor("float32", new Float32Array(this.fixedSeqLen * ANCHOR_FEATURE_DIM), [
				1,
				this.fixedSeqLen,
				ANCHOR_FEATURE_DIM,
			])

			feeds.anchor_confidence = new ort.Tensor("float32", new Float32Array(this.fixedSeqLen), [1, this.fixedSeqLen])
		}

		// Gazetteer-anchor channel (#464): same feed contract as the postcode anchor. Feature width is
		// read from the supplied rows (the lexicon's slot count); a gazetteer-trained model with no clue
		// data supplied gets the confidence=0 identity (the model's gazetteer-off behavior).
		if (gazetteer && session.inputNames.includes("gazetteer_features")) {
			const dim = gazetteer.features[0]?.length ?? 0
			const gf = new Float32Array(this.fixedSeqLen * dim)
			const gc = new Float32Array(this.fixedSeqLen)

			for (let i = 0; i < seqLen; i++) {
				gc[i] = gazetteer.confidence[i] ?? 0
				const row = gazetteer.features[i]

				if (row) {
					for (let d = 0; d < dim; d++) {
						gf[i * dim + d] = row[d] ?? 0
					}
				}
			}

			feeds.gazetteer_features = new ort.Tensor("float32", gf, [1, this.fixedSeqLen, dim])
			feeds.gazetteer_confidence = new ort.Tensor("float32", gc, [1, this.fixedSeqLen])
		} else if (session.inputNames.includes("gazetteer_features")) {
			feeds.gazetteer_features = new ort.Tensor("float32", new Float32Array(this.fixedSeqLen * GAZETTEER_FEATURE_DIM), [
				1,
				this.fixedSeqLen,
				GAZETTEER_FEATURE_DIM,
			])

			feeds.gazetteer_confidence = new ort.Tensor("float32", new Float32Array(this.fixedSeqLen), [1, this.fixedSeqLen])
		}

		// Country-lexicon channel (#1104): same feed contract as the gazetteer. Feature width read from the supplied
		// rows (COUNTRY_FEATURE_DIM); a country-trained model with no lexicon supplied gets the confidence=0 identity.
		if (country && session.inputNames.includes("country_features")) {
			const dim = country.features[0]?.length ?? 0
			const cf = new Float32Array(this.fixedSeqLen * dim)
			const cc = new Float32Array(this.fixedSeqLen)

			for (let i = 0; i < seqLen; i++) {
				cc[i] = country.confidence[i] ?? 0
				const row = country.features[i]

				if (row) {
					for (let d = 0; d < dim; d++) {
						cf[i * dim + d] = row[d] ?? 0
					}
				}
			}

			feeds.country_features = new ort.Tensor("float32", cf, [1, this.fixedSeqLen, dim])
			feeds.country_confidence = new ort.Tensor("float32", cc, [1, this.fixedSeqLen])
		} else if (session.inputNames.includes("country_features")) {
			feeds.country_features = new ort.Tensor("float32", new Float32Array(this.fixedSeqLen * COUNTRY_FEATURE_DIM), [
				1,
				this.fixedSeqLen,
				COUNTRY_FEATURE_DIM,
			])

			feeds.country_confidence = new ort.Tensor("float32", new Float32Array(this.fixedSeqLen), [1, this.fixedSeqLen])
		}

		// Evidence-bundle channels (Option-A, Phase 2): same feed contract as every soft channel —
		// present-conditional on the graph's declared inputs, confidence=0 identity zero-fallback for a
		// bundle-trained model run without lexicons. Inert against every pre-bundle model by construction.
		const evidenceFeeds = [
			{ prefix: "street_type", dim: STREET_TYPE_FEATURE_DIM, data: evidence?.streetType },
			{ prefix: "locality_surface", dim: LOCALITY_SURFACE_FEATURE_DIM, data: evidence?.localitySurface },
		] as const

		for (const { prefix, dim: fallbackDim, data } of evidenceFeeds) {
			if (!session.inputNames.includes(`${prefix}_features`)) continue

			if (data) {
				const dim = data.features[0]?.length ?? fallbackDim
				const ef = new Float32Array(this.fixedSeqLen * dim)
				const ec = new Float32Array(this.fixedSeqLen)

				for (let i = 0; i < seqLen; i++) {
					ec[i] = data.confidence[i] ?? 0
					const row = data.features[i]

					if (row) {
						for (let d = 0; d < dim; d++) {
							ef[i * dim + d] = row[d] ?? 0
						}
					}
				}

				feeds[`${prefix}_features`] = new ort.Tensor("float32", ef, [1, this.fixedSeqLen, dim])
				feeds[`${prefix}_confidence`] = new ort.Tensor("float32", ec, [1, this.fixedSeqLen])
			} else {
				feeds[`${prefix}_features`] = new ort.Tensor("float32", new Float32Array(this.fixedSeqLen * fallbackDim), [
					1,
					this.fixedSeqLen,
					fallbackDim,
				])

				feeds[`${prefix}_confidence`] = new ort.Tensor("float32", new Float32Array(this.fixedSeqLen), [
					1,
					this.fixedSeqLen,
				])
			}
		}

		const output = await session.run(feeds)
		const logitsTensor = output.logits

		if (!logitsTensor) throw new Error("ONNX model did not return a `logits` output")
		const data = logitsTensor.data as Float32Array
		// dims are [batch, sequence, labels].
		const numLabels = (logitsTensor.dims as readonly [number, number, number])[2]

		const logits: number[][] = []

		for (let t = 0; t < seqLen; t++) {
			const row: number[] = new Array(numLabels)
			const base = t * numLabels

			for (let l = 0; l < numLabels; l++) {
				row[l] = data[base + l]!
			}

			logits.push(row)
		}

		// Locale head (#511 Tier A): present on v1.1.0+ exports, absent (and optional) before.
		const localeTensor = output.locale_logits
		const localeLogits = localeTensor ? Array.from(localeTensor.data as Float32Array) : undefined

		// Span head (#727 stage-2): present on v3.x+ exports. Same optional contract as the locale head
		// — a pre-v3 bundle simply has no `span_scores` output and the BIO path is unaffected.
		const spanTensor = output.span_scores
		let spanScores: number[][][] | undefined
		let maxSpan: number | undefined

		if (spanTensor) {
			const spanData = spanTensor.data as Float32Array
			// dims are [batch, sequence, span, type].
			const spanDims = spanTensor.dims as readonly [number, number, number, number]
			const spanLen = spanDims[2]
			const numTypes = spanDims[3]
			maxSpan = spanLen
			spanScores = []

			// Only the first `seqLen` token rows are real; the rest is the fixed-length pad tail.
			for (let t = 0; t < seqLen; t++) {
				const perLength: number[][] = new Array(spanLen)

				for (let l = 0; l < spanLen; l++) {
					const row: number[] = new Array(numTypes)
					const base = (t * spanLen + l) * numTypes

					for (let ty = 0; ty < numTypes; ty++) {
						row[ty] = spanData[base + ty]!
					}

					perLength[l] = row
				}

				spanScores.push(perLength)
			}
		}

		return {
			logits,
			numLabels,
			...(localeLogits ? { localeLogits } : {}),
			...(spanScores ? { spanScores, maxSpan } : {}),
		}
	}

	/**
	 * The model's declared input names (loads the session if not already loaded). Used by the ProductionScorer (#718)
	 * back-compat path: when a model-card has no `requires` block, the required soft-feature channels are INFERRED from
	 * the graph — a model exporting `anchor_features` / `gazetteer_features` declared those channels mandatory at train
	 * time.
	 */
	async inputNames(): Promise<readonly string[]> {
		const session = await this.ensureSession()

		return session.inputNames
	}
}