/**
 * @copyright Sister Software
 * @license AGPL-3.0
 * @author Teffen Ellis, et al.
 *
 * Splits corpus rows into train, val and test by holding out whole places: a random row split leaks by
 * neighborhood, because the model can memorize a street it saw in train. Rows inside a held-out place
 * go to val or test 50/50 by a hash of `source_id`; every other row goes to train.
 *
 * A holdout change takes effect at the next corpus rebuild. Each versioned corpus keeps the
 * `SPLIT_MANIFEST.json` it was built with, because adding a holdout to a built corpus would leak rows
 * the model already trained on.
 */

import { delimitedSource, preferCompressed } from "@mailwoman/core/fs/delimited"
import { openWriteStream } from "@mailwoman/core/fs/streams"
import { writeLocalJSONFile, writeLocalTextFile, makeDirectories, removePath } from "@mailwoman/core/fs/writers"
import { createHash } from "@mailwoman/core/hash"
import { spawnProcess } from "@mailwoman/core/process"
import { childEnv } from "@mailwoman/core/scripting/utils"
import { type PathBuilderLike, resolvePathBuilder } from "path-ts"
import { JSONSpliterator } from "spliterator"

import type { CanonicalRow, LabeledRow } from "#types"

/**
 * The name of one corpus split.
 */
export type SplitName = "train" | "val" | "test"

/**
 * The component values that identify the places one country holds out.
 *
 * Sources represent a place in different components: `usgov-nad` and `tiger` emit `region`,
 * but BAN emits none, so a French holdout must also match on postcode.
 * A row is held out when any declared matcher fires.
 * A policy with no matchers holds out no entry.
 */
export interface HoldoutPolicy {
	regions?: readonly string[]
	/**
	 * Prefixes matched against `row.components.postcode`; use one only where it identifies
	 * a single place by itself, such as the first two digits of a French postcode.
	 */
	postcodePrefixes?: readonly string[]
	localities?: readonly string[]
}

/**
 * One country's holdout as a bare region list or a {@link HoldoutPolicy}.
 *
 * Older committed `SPLIT_MANIFEST.json` files use a bare array of region values.
 */
export type CountryHoldout = readonly string[] | HoldoutPolicy

/**
 * Options for {@link splitRows}.
 */
export interface SplitOptions {
	/**
	 * The holdout policy keyed by ISO 3166-1 alpha-2 country code.
	 */
	holdouts?: Record<string, CountryHoldout>
}

/**
 * The `source_id` lists for each split.
 */
export interface SplitManifest {
	train: string[]
	val: string[]
	test: string[]
	holdouts: Record<string, CountryHoldout>
	/**
	 * The corpus version, read from the first row that has one.
	 */
	corpus_version: string
	counts: { train: number; val: number; test: number; total: number }
}

/**
 * Returns the default holdouts — small, peripheral places in each country.
 * A change here takes effect at the next base corpus rebuild.
 */
export function defaultHoldouts(): Record<string, CountryHoldout> {
	return {
		US: ["Vermont", "VT", "Wyoming", "WY", "North Dakota", "ND"],
		FR: {
			regions: ["Corse", "Lozère", "Lozere", "Creuse"],
			// BAN street rows have no region, so these postcode departments hold out the
			// same three places: Corse 20, Creuse 23 and Lozère 48.
			postcodePrefixes: ["20", "23", "48"],
		},
		DE: {
			// Only `wof-admin` has a DE region, so these Länder hold out admin rows only.
			regions: ["Saarland", "SL", "Mecklenburg-Vorpommern", "MV"],
			// DE street rows come from two OpenAddresses members, Berlin and Sachsen,
			// so the regions above match no street row.
			// Postcode area `02` (Upper Lusatia in eastern Sachsen) is the street-level holdout.
			// German postcodes have five digits, so this holdout cannot collide.
			postcodePrefixes: ["02"],
		},
		GB: {
			// GB street rows have no region, so the holdout matches postcode areas —
			// Truro, Llandudno and Halifax.
			// Each prefix has two letters and none is a prefix of another
			// (`L` for Liverpool would also match `LL`).
			postcodePrefixes: ["TR", "LL", "HX"],
			// Cornwall matches the admin rows for the same place `TR` covers.
			regions: ["Cornwall"],
		},
		// The eight countries below enter training through OpenStreetMap rows, most of which carry no
		// region and many no postcode, so each holdout matches its place three ways: the region an admin
		// row carries, the postcode prefix, and the city as the OSM rows spell it. The spellings are the
		// ones the 2026-10-02 extracts write.
		CN: {
			// Gansu by its province and capital, and Qinghai `81` and Tibet `85` by postcode.
			regions: ["甘肃省", "甘肃", "Gansu"],
			postcodePrefixes: ["81", "85"],
			localities: ["兰州市", "兰州", "Lanzhou"],
		},
		IN: {
			// Kerala, whose PIN codes run 67 to 69.
			regions: ["Kerala"],
			postcodePrefixes: ["67", "68", "69"],
			localities: ["Kochi", "Kollam", "Thiruvananthapuram", "Kozhikode"],
		},
		ID: {
			// Bali, whose postcodes begin `80`.
			regions: ["Bali"],
			postcodePrefixes: ["80"],
			localities: ["Denpasar", "Bali"],
		},
		NG: {
			// Enugu State.
			// OSM Nigeria rows carry almost no postcode or region, so the cities carry it.
			regions: ["Enugu"],
			localities: ["Enugu", "Nsukka"],
		},
		PK: {
			// Islamabad Capital Territory, in both scripts the extract writes.
			regions: ["Islamabad Capital Territory"],
			postcodePrefixes: ["44"],
			localities: ["Islamabad", "islamabad", "اسلام آباد"],
		},
		BD: {
			// Chattogram, Khulna and Sylhet, whose postcodes begin 4, 9 and 3 respectively.
			postcodePrefixes: ["40", "41", "42", "43", "90", "91", "92", "31"],
			localities: ["Chittagong", "Chattogram", "Khulna", "Sylhet"],
		},
		VN: {
			regions: ["Thành phố Cần Thơ", "Cần Thơ", "Thừa Thiên Huế"],
			localities: ["Cần Thơ", "Thành phố Cần Thơ", "Huế", "Thành phố Huế"],
		},
		RU: {
			// Primorsky Krai, whose postcodes run 690 to 692.
			regions: ["Приморский край"],
			postcodePrefixes: ["690", "691", "692"],
			localities: ["Владивосток", "Уссурийск", "Находка"],
		},
		// Each French overseas jurisdiction is one INSEE department with one postcode range,
		// so the BAN rows hold out by commune, spelled as BAN's `nom_commune` writes it.
		// Saint-Barthélemy and Saint-Martin are each one commune under one postcode,
		// so neither has a place to hold out.
		GP: { localities: ["Le Moule"] },
		MQ: { localities: ["Le Robert", "Ducos"] },
		GF: { localities: ["Kourou"] },
		RE: { localities: ["Saint-Louis", "Saint-Leu"] },
		PM: { localities: ["Miquelon-Langlade"] },
		YT: { localities: ["Dzaoudzi", "Pamandzi"] },
		// Tumaraa holds 1,522 of PF's 6,755 rows with a street, more than a tenth,
		// because Pirae's 958 rows carry no street name and Papeete holds the remaining 5,233.
		PF: { localities: ["Tumaraa"] },
		NC: { localities: ["Koné", "Pouembout", "Kaala-Gomen"] },
		LV: {
			// Latgale's south-east: the Krāslava, Ludza, Balvi, Preiļi and Līvāni
			// municipalities, whose postcodes begin LV-56, LV-57, LV-45 and LV-53,
			// and the state city Rēzekne, which lies in no municipality.
			// The address register writes `nov.`; other sources spell the municipality out.
			regions: [
				"Krāslavas nov.",
				"Ludzas nov.",
				"Balvu nov.",
				"Preiļu nov.",
				"Līvānu nov.",
				"Krāslavas novads",
				"Ludzas novads",
				"Balvu novads",
				"Preiļu novads",
				"Līvānu novads",
			],
			postcodePrefixes: ["LV-45", "LV-53", "LV-56", "LV-57"],
			localities: ["Rēzekne"],
		},
		RO: {
			// Alba 51, Hunedoara 33, Mureș 54 and Sibiu 55: a Romanian postcode's first two digits name its county.
			regions: ["Alba", "Hunedoara", "Mureș", "Sibiu"],
			postcodePrefixes: ["51", "33", "54", "55"],
		},
	}
}

/**
 * Returns the matchers for one country, reading a bare array as a region list.
 */
export function holdoutPolicyFor(holdout: CountryHoldout | undefined): HoldoutPolicy {
	if (!holdout) return {}

	return Array.isArray(holdout) ? { regions: holdout } : (holdout as HoldoutPolicy)
}

type SplitInputRow = Pick<CanonicalRow, "source_id" | "country" | "corpus_version" | "components">

/**
 * Returns the split for one row.
 *
 * Both `splitRows` and the streaming `buildCorpus` loop call this, so every
 * caller assigns a row to the same split.
 */
export function splitForRow(
	row: Pick<SplitInputRow, "source_id" | "country" | "components">,
	holdouts: Record<string, CountryHoldout> = defaultHoldouts()
): SplitName {
	const policy = holdoutPolicyFor(holdouts[row.country])
	const { region, postcode, locality } = row.components

	const isHeldOut =
		(region !== undefined && (policy.regions?.includes(region) ?? false)) ||
		(postcode !== undefined && (policy.postcodePrefixes?.some((prefix) => postcode.startsWith(prefix)) ?? false)) ||
		(locality !== undefined && (policy.localities?.includes(locality) ?? false))

	if (!isHeldOut) return "train"

	return hashBucket(row.source_id, 2) === 0 ? "val" : "test"
}

/**
 * Builds a `SplitManifest` in memory from labeled or canonical rows.
 *
 * Tests and small fixtures use this, while `buildCorpus` uses `splitForRow`
 * and `writeSplitManifestsFromLabeledFiles` so it never holds every row's split in memory.
 */
export function splitRows(rows: Iterable<SplitInputRow>, opts: SplitOptions = {}): SplitManifest {
	const holdouts = opts.holdouts ?? defaultHoldouts()
	const train: string[] = []
	const val: string[] = []
	const test: string[] = []
	let corpus_version = ""

	for (const row of rows) {
		if (!corpus_version && row.corpus_version) {
			corpus_version = row.corpus_version
		}

		const split = splitForRow(row, holdouts)

		if (split === "train") {
			train.push(row.source_id)
		} else if (split === "val") {
			val.push(row.source_id)
		} else {
			test.push(row.source_id)
		}
	}

	const total = train.length + val.length + test.length

	return {
		train,
		val,
		test,
		holdouts,
		corpus_version,
		counts: { train: train.length, val: val.length, test: test.length, total },
	}
}

/**
 * Returns a deterministic bucket in `0..n-1` for an id, from the first four bytes
 * of the SHA-256 digest read as a big-endian uint32.
 */
export function hashBucket(id: string, n: number): number {
	const digest = createHash("sha256").update(id).digest()
	const u = digest[0]! * 0x01_00_00_00 + digest[1]! * 0x01_00_00 + digest[2]! * 0x01_00 + digest[3]!

	return u % n
}

/**
 * Writes a `SplitManifest` to `<outputDir>/{train,val,test}.txt` (sorted `source_id` values, one per line)
 * and `SPLIT_MANIFEST.json` (the corpus version, holdouts and counts).
 */
export async function writeSplitManifests(manifest: SplitManifest, outputDir: PathBuilderLike): Promise<void> {
	await makeDirectories(outputDir)

	for (const name of ["train", "val", "test"] as const) {
		const sorted = [...manifest[name]].toSorted()

		await writeLocalTextFile(sorted, outputDir, `${name}.txt`)
	}

	const summary = {
		corpus_version: manifest.corpus_version,
		holdouts: manifest.holdouts,
		counts: manifest.counts,
	}

	await writeLocalJSONFile(summary, outputDir, "SPLIT_MANIFEST.json")
}

/**
 * The `LabeledRow` fields that the split functions read.
 */
export type SplitInputLabeledRow = Pick<LabeledRow, "source_id" | "country" | "corpus_version" | "components">

/**
 * Writes the same files as `writeSplitManifests` by streaming one labeled JSONL file per split.
 *
 * `buildCorpus` calls this after its align loop partitions rows with `splitForRow`.
 * The caller passes the counts to avoid rescanning the files.
 * `sort(1)` spills to disk, so memory stays constant.
 */
export async function writeSplitManifestsFromLabeledFiles(opts: {
	labeledPaths: Record<SplitName, PathBuilderLike>
	outputDir: PathBuilderLike
	corpusVersion: string
	counts: Record<SplitName, number>
	holdouts?: Record<string, CountryHoldout>
}): Promise<SplitManifest["counts"]> {
	await makeDirectories(opts.outputDir)
	const holdouts = opts.holdouts ?? defaultHoldouts()

	for (const split of ["train", "val", "test"] as const) {
		const labeledPath = opts.labeledPaths[split]
		const outPath = resolvePathBuilder(opts.outputDir, `${split}.txt`)
		await streamSortedSourceIDs(labeledPath, outPath)
	}

	const total = opts.counts.train + opts.counts.val + opts.counts.test

	const summary = {
		corpus_version: opts.corpusVersion,
		holdouts,
		counts: { ...opts.counts, total },
	}

	await writeLocalJSONFile(summary, opts.outputDir, "SPLIT_MANIFEST.json")

	return summary.counts
}

/**
 * Writes the sorted `source_id` values from a labeled JSONL file to `outPath`;
 * an empty input still produces an empty output file.
 */
async function streamSortedSourceIDs(labeledJsonlPath: PathBuilderLike, outPath: PathBuilderLike): Promise<void> {
	const unsortedPath = `${outPath}.unsorted`
	const out = openWriteStream(unsortedPath, { encoding: "utf8" })

	const outClosed = new Promise<void>((resolve, reject) => {
		out.on("close", () => resolve())
		out.on("error", reject)
	})

	// A malformed row throws out of the loop.
	// The `finally` block still closes the write stream.
	try {
		for await (const obj of JSONSpliterator.fromAsync<{ source_id?: string }>(
			delimitedSource(await preferCompressed(labeledJsonlPath))
		)) {
			if (typeof obj.source_id === "string") {
				out.write(`${obj.source_id}\n`)
			}
		}
	} finally {
		out.end()
	}

	await outClosed

	await new Promise<void>((resolve, reject) => {
		// `LC_ALL=C` sorts by byte, so the order is the same on every host.
		const proc = spawnProcess("sort", [unsortedPath, "-o", outPath], { env: childEnv({ LC_ALL: "C" }) })
		proc.on("error", reject)

		proc.on("exit", (code) => {
			if (code === 0) {
				resolve()
			} else {
				reject(new Error(`sort exited with code ${code}`))
			}
		})
	})

	await removePath(unsortedPath).catch(() => {})
}