/**
 * @copyright Sister Software
 * @license AGPL-3.0
 * @author Teffen Ellis, et al.
 *
 *   Pre-compute corpus-wide token + bigram label distributions for the corpus linter.
 *
 *   Reads one or more Parquet shards, builds per-(token, label) and per-(bigram, label-bigram)
 *   histograms, and serializes them as JSON. The output file is consumed by `lint-corpus-shard.ts`
 *   as the baseline against which a new shard is compared.
 *
 *   Stats are cheap to compute (~5–30s per 100K rows) but expensive enough that we cache them between
 *   linter invocations. Re-run this script whenever the corpus changes substantially (a new
 *   mainline shard added, a source-pool re-weighted, etc.).
 *
 *   Output schema:
 *
 *   ```ts
 *   interface CorpusStats {
 *     row_count: number
 *     shard_paths: string[]
 *     tokens: { [token: string]: { [label: string]: number } }
 *     bigrams: { [token_bigram: string]: { [label_bigram: string]: number } }
 *     // token_bigram = "tok1tok2" (US sep), label_bigram = "lab1lab2"
 *     // For memory: only keep bigrams with count >= MIN_BIGRAM_COUNT (2).
 *   }
 * ```
 *
 *   Usage: node scripts/build-corpus-stats.ts\
 *   --shards <glob-pattern-or-dir>\
 *   --output <stats.json>
 *
 *   For a quick local-corpus baseline (limited but useful for linter testing): node
 *   scripts/build-corpus-stats.ts\
 *   --shards $MAILWOMAN_DATA_ROOT/corpus/versioned/v0.4.0/corpus-v0.4.0/train/\
 *   --output /tmp/corpus-stats-local.json
 */

import { readdirSync, statSync, writeFileSync } from "node:fs"
import { join } from "node:path"

import { ParquetReader } from "../parquet-wrapper/index.ts"

const SEP = ""
const MIN_BIGRAM_COUNT = 2

export interface CorpusStatsOptions {
	shardsArg: string
	outputPath: string
	limitPerShard?: number
}

function discoverShards(shardsArg: string): string[] {
	const stat = statSync(shardsArg)

	if (stat.isDirectory()) {
		return readdirSync(shardsArg)
			.filter((f) => f.endsWith(".parquet"))
			.map((f) => join(shardsArg, f))
	}

	if (stat.isFile() && shardsArg.endsWith(".parquet")) return [shardsArg]

	// Otherwise treat as a literal path list (one per line if it's stdin-friendly).
	return [shardsArg]
}

/**
 * Stream a shard's `tokens`/`labels` columns.
 *
 * `limit` stops the iteration rather than filtering afterwards, so a capped run reads only the row groups it needs.
 */
async function* streamShardRows(
	shardPath: string,
	limit?: number
): AsyncIterable<{ tokens: string[]; labels: string[] }> {
	await using reader = await ParquetReader.openFile<{ tokens: string[]; labels: string[] }>(shardPath)

	let emitted = 0

	for await (const row of reader.project("tokens", "labels")) {
		if (limit !== undefined && emitted >= limit) break

		yield row

		emitted++
	}
}

export async function buildCorpusStats(args: CorpusStatsOptions): Promise<void> {
	const shardPaths = discoverShards(args.shardsArg)

	console.error(`Discovered ${shardPaths.length} parquet shard(s)`)

	const tokenStats = new Map<string, Map<string, number>>()
	const bigramStats = new Map<string, Map<string, number>>()
	let totalRows = 0

	for (const path of shardPaths) {
		console.error(`Reading ${path}...`)

		const before = totalRows

		for await (const { tokens, labels } of streamShardRows(path, args.limitPerShard)) {
			totalRows++

			if (tokens.length !== labels.length) continue

			// skip malformed
			for (let i = 0; i < tokens.length; i++) {
				const tk = tokens[i]!
				const lb = labels[i]!
				let labelMap = tokenStats.get(tk)

				if (!labelMap) {
					labelMap = new Map()
					tokenStats.set(tk, labelMap)
				}

				labelMap.set(lb, (labelMap.get(lb) ?? 0) + 1)

				if (i + 1 < tokens.length) {
					const bigramKey = tk + SEP + tokens[i + 1]!
					const bigramLabel = lb + SEP + labels[i + 1]!
					let bMap = bigramStats.get(bigramKey)

					if (!bMap) {
						bMap = new Map()
						bigramStats.set(bigramKey, bMap)
					}

					bMap.set(bigramLabel, (bMap.get(bigramLabel) ?? 0) + 1)
				}
			}
		}

		console.error(
			`  ${totalRows - before} rows; running totals: ${tokenStats.size} unique tokens, ${bigramStats.size} unique bigrams`
		)
	}

	// Prune bigrams below MIN_BIGRAM_COUNT to keep the output file size sane. Token stats
	// stay complete — they're cheap and we need accuracy at the long tail for label-vacuum
	// detection.
	let prunedBigrams = 0

	for (const [k, labelMap] of bigramStats) {
		let total = 0

		for (const v of labelMap.values()) {
			total += v
		}

		if (total < MIN_BIGRAM_COUNT) {
			bigramStats.delete(k)

			prunedBigrams++
		}
	}

	console.error(`Pruned ${prunedBigrams} singleton bigrams; ${bigramStats.size} remain`)

	const out = {
		row_count: totalRows,
		shard_paths: shardPaths,
		tokens: {} as Record<string, Record<string, number>>,
		bigrams: {} as Record<string, Record<string, number>>,
	}

	for (const [tk, labelMap] of tokenStats) {
		out.tokens[tk] = Object.fromEntries(labelMap)
	}

	for (const [k, labelMap] of bigramStats) {
		out.bigrams[k] = Object.fromEntries(labelMap)
	}

	writeFileSync(args.outputPath, JSON.stringify(out))
	const sizeMB = (Buffer.byteLength(JSON.stringify(out)) / 1024 / 1024).toFixed(1)

	console.error(
		`Wrote ${args.outputPath} (${sizeMB} MB) — ${totalRows} rows, ${tokenStats.size} tokens, ${bigramStats.size} bigrams`
	)
}