// SPDX-License-Identifier: MPL-2.0
// (c) Hare authors <https://harelang.org>

use types;

export type tokenizer = struct {
	in: []u8,	// string being tokenized
	delim: []u8,	// delimiter
	p: i64,		// p < 0 for reverse tokenizers, 0 <= p for forward ones.
};

// Tokenizes a byte slice, returning an iterator that yields tokens from the
// slice delimited by one of any number of delimiter bytes. If the input slice
// begins with or ends with a delimiter, an empty slice is returned respectively
// as the first and last call to [[next_token]].
//
// The variadic argument slice is borrowed from the caller, who should take care
// to ensure that it is valid for the lifetime of the tokenizer.
//
// The caller must ensure that at least one delimiter is provided and that the
// length of the slice is less than [[types::I64_MAX]].
export fn tokenize(in: []u8, delim: u8...) tokenizer = {
	assert(len(delim) > 0, "bytes::tokenize called with empty slice");
	assert(len(in) < types::I64_MAX: size,
		"bytes::tokenize: input length exceeds I64_MAX");
	if (len(in) == 0) {
		delim = [];
	};

	return tokenizer {
		in = in,
		delim = delim,
		p = types::I64_MAX, // I64_MAX means we haven't peeked the next token yet.
	};
};

// Like [[tokenize]], but tokenizes the slice in reverse, such that the first
// call to [[next_token]] returns the last token and the last call returns the
// first token.
export fn rtokenize(in: []u8, delim: u8...) tokenizer = {
	assert(len(delim) > 0, "bytes::rtokenize called with empty slice");
	assert(len(in) < types::I64_MAX: size,
		"bytes::rtokenize: input length exceeds I64_MAX");
	if (len(in) == 0) {
		delim = [];
	};

	return tokenizer {
		in = in,
		delim = delim,
		// I64_MIN means we haven't peeked the next token yet. Note that
		// p == -1 corresponds to an index of len(s), and
		// p == -(1 - len(s)) corresponds to an index of 0.
		p = types::I64_MIN,
	};
};

// Returns the next token from a [[tokenizer]] and advances the cursor.
export fn next_token(s: *tokenizer) ([]u8 | done) = {
	const b = match (peek_token(s)) {
	case let b: []u8 =>
		yield b;
	case done => return done;
	};

	const slen = len(s.in): i64;
	const reverse = s.p < 0;
	if (reverse) {
		if (slen + s.p + 1 == 0) {
			s.delim = s.delim[..0];
			s.in = s.in[..0];
		} else {
			const end = (slen + s.p + 1): size - 1;
			s.in = s.in[..end];
		};
		s.p = types::I64_MIN;
	} else {
		if (s.p == slen) {
			s.delim = s.delim[..0];
			s.in = s.in[..0];
		} else {
			s.in = s.in[s.p: size + 1..];
		};
		s.p = types::I64_MAX;
	};

	return b;
};

// Returns the next token from a [[tokenizer]] without advancing the cursor.
export fn peek_token(s: *tokenizer) ([]u8 | done) = {
	if (len(s.delim) == 0) {
		return done;
	};

	const reverse = s.p < 0;
	const ifunc = if (reverse) &rindex else &index;

	const known = ((reverse && s.p != types::I64_MIN) ||
		(!reverse && s.p != types::I64_MAX));
	if (!known) {
		let i = if (reverse) types::I64_MIN else types::I64_MAX;
		let dlen = 0i64;
		const slen = len(s.in): i64;

		for (let d .. s.delim) {
			match (ifunc(s.in, d)) {
			case let ix: size =>
				if (!reverse && ix: i64 < i) {
					i = ix: i64;
					dlen = 1;
				} else if (reverse && ix: i64 > i) {
					i = ix: i64;
					dlen = 1;
				};
			case void =>
				if (!reverse && slen < i: i64) {
					i = slen;
				} else if (reverse && 0 > i: i64) {
					i = 0;
				};
			};
		};

		if (reverse) {
			if (i == slen) {
				s.p = -(slen + 1);
			} else {
				s.p = i + dlen - slen - 1;
			};
		} else {
			s.p = i;
		};
	};

	if (reverse) {
		return s.in[len(s.in) + s.p: size + 1..];
	} else {
		return s.in[..s.p: size];
	};
};

// Returns the remainder of the input slice from a [[tokenizer]] ahead of the
// token cursor.
export fn remaining_tokens(s: *tokenizer) []u8 = {
	return s.in;
};

// Splits a byte slice into tokens delimited by 'delim', starting at the
// beginning of the slice, and returning a slice of up to N tokens. The caller
// must free the returned slice. The slices within the result are borrowed from
// 'in'.
//
// The caller must ensure that 'delim' is not empty. The caller must free the
// result after use.
export fn splitn(in: []u8, delim: []u8, n: size) ([][]u8 | nomem) = {
	assert(len(delim) > 0, "bytes::splitn must not be called with an empty delimiter");
	let toks: [][]u8 = [];
	let ok = false;
	defer if (!ok) free(toks);
	let tok = tokenize(in, delim...);
	for (let i = 0z; i < n - 1z; i += 1) {
		match (next_token(&tok)) {
		case let s: []u8 =>
			append(toks, s)?;
		case done =>
			ok = true;
			return toks;
		};
	};
	match (peek_token(&tok)) {
	case done => void;
	case []u8 =>
		append(toks, remaining_tokens(&tok))?;
	};
	ok = true;
	return toks;
};

// Splits a byte slice into tokens delimited by 'delim', starting at the end of
// the slice, and returning a slice of up to N tokens. The caller must free the
// returned slice. The slices within the result are borrowed from 'in'.
//
// The caller must ensure that 'delim' is not empty. The caller must free the
// result after use.
export fn rsplitn(in: []u8, delim: []u8, n: size) ([][]u8 | nomem) = {
	assert(len(delim) > 0, "bytes::rsplitn called with empty delimiter");
	let toks: [][]u8 = [];
	let tok = rtokenize(in, delim...);
	let ok = false;
	defer if (!ok) free(toks);
	for (let i = 0z; i < n - 1z; i += 1) {
		match (next_token(&tok)) {
		case let s: []u8 =>
			append(toks, s)?;
		case done =>
			ok = true;
			return toks;
		};
	};
	match (peek_token(&tok)) {
	case done => void;
	case []u8 =>
		append(toks, remaining_tokens(&tok))?;
	};

	for (let i = 0z; i < len(toks) / 2; i += 1) {
		const j = len(toks) - i - 1;
		const tmp = toks[i];
		toks[i] = toks[j];
		toks[j] = tmp;
	};

	ok = true;
	return toks;
};

// Splits a byte slice into tokens delimited by byte present in 'delim' (see
// [[tokenize]] for details).
//
// The caller must free the returned slice. The byte slices within the slice are
// borrowed from 'in'.
//
// The caller must ensure that 'delim' is not empty.
export fn split(in: []u8, delim: []u8) ([][]u8 | nomem) =
	splitn(in, delim, types::SIZE_MAX);

fn tokenize_test(
	func: *fn(_: []u8, _: u8...) tokenizer,
	in: []u8,
	delim: []u8,
	tokens: [][]u8,
	iters: size = types::SIZE_MAX,
) tokenizer = {
	const tok = func(in, delim...);
	let n = 0z;
	for (const want .. tokens) {
		if (n >= iters) {
			return tok;
		};
		n += 1;

		const p = peek_token(&tok) as []u8;
		const n = next_token(&tok) as []u8;
		assert(equal(p, n));
		assert(equal(n, want));
	};

	if (n >= iters) {
		return tok;
	};

	assert(peek_token(&tok) is done);
	assert(next_token(&tok) is done);
	return tok;
};

@test fn tokenize() void = {
	// simple case
	tokenize_test(&tokenize, [1, 2, 0, 3, 4], [0], [
		[1, 2],
		[3, 4],
	]);

	// multiple delimiters
	tokenize_test(&tokenize, [1, 2, 0, 3, 4, 42, 5, 6], [0, 42], [
		[1, 2],
		[3, 4],
		[5, 6],
	]);

	// empty tokens
	tokenize_test(&tokenize, [1, 2, 0, 0, 0, 3, 4], [0], [
		[1, 2],
		[],
		[],
		[3, 4],
	]);

	// leading empty tokens
	tokenize_test(&tokenize, [0, 1, 2, 3, 0], [0], [
		[],
		[1, 2, 3],
		[],
	]);

	// remaining_tokens
	const tok = tokenize_test(&tokenize, [1, 2, 0, 3, 4], [0], [
		[1, 2],
	], 1);
	assert(equal(remaining_tokens(&tok), [3, 4]));
};

@test fn rtokenize() void = {
	// simple case
	tokenize_test(&rtokenize, [1, 2, 0, 3, 4], [0], [
		[3, 4],
		[1, 2],
	]);

	// multiple delimiters
	tokenize_test(&rtokenize, [1, 2, 0, 3, 4, 42, 5, 6], [0, 42], [
		[5, 6],
		[3, 4],
		[1, 2],
	]);

	// empty tokens
	tokenize_test(&rtokenize, [1, 2, 0, 0, 0, 3, 4], [0], [
		[3, 4],
		[],
		[],
		[1, 2],
	]);

	// leading empty tokens
	tokenize_test(&rtokenize, [0, 1, 2, 3, 0], [0], [
		[],
		[1, 2, 3],
		[],
	]);

	// remaining_tokens
	const tok = tokenize_test(&rtokenize, [1, 2, 0, 3, 4], [0], [
		[3, 4],
	], 1);
	assert(equal(remaining_tokens(&tok), [1, 2]));
};

@test fn split() void = {
	const input = *(&"Hello, my name is Drew": *[]u8);
	const space: [1]u8 = [' '];

	const expected = [
		*(&"Hello,": *[]u8),
		*(&"my": *[]u8),
		*(&"name": *[]u8),
		*(&"is Drew": *[]u8),
	];
	const actual = splitn(input, space[..], 4z)!;
	assert(len(expected) == len(actual));
	for (let i = 0z; i < len(expected); i += 1) {
		assert(equal(expected[i], actual[i]));
	};
	free(actual);

	const expected2 = [
		*(&"Hello,": *[]u8),
		*(&"my": *[]u8),
		*(&"name": *[]u8),
		*(&"is": *[]u8),
		*(&"Drew": *[]u8),
	];
	const actual2 = split(input, space[..])!;
	assert(len(expected2) == len(actual2));
	for (let i = 0z; i < len(expected2); i += 1) {
		assert(equal(expected2[i], actual2[i]));
	};
	free(actual2);

	const eq: [1]u8 = ['='];
	const single = *(&"one": *[]u8);
	const actual3 = splitn(single, eq[..], 2z)!;
	const expected3 = [
		*(&"one": *[]u8),
	];
	assert(len(expected3) == len(actual3));
	for (let i = 0z; i < len(expected3); i += 1) {
		assert(equal(expected3[i], actual3[i]));
	};
	free(actual3);

	const expected4 = [
		*(&"Hello, my": *[]u8),
		*(&"name": *[]u8),
		*(&"is": *[]u8),
		*(&"Drew": *[]u8),
	];
	const actual4 = rsplitn(input, space[..], 4z)!;
	assert(len(expected4) == len(actual4));
	for (let i = 0z; i < len(expected4); i += 1) {
		assert(equal(expected4[i], actual4[i]));
	};
	free(actual4);
};

// Returns the input slice "cut" along the first instance of a delimiter,
// returning everything up to the delimiter, and everything after the delimiter,
// in a tuple. The contents are borrowed from the input slice.
//
// The caller must ensure that 'delimiter' is not an empty slice.
export fn cut(in: []u8, delim: ([]u8 | u8)) ([]u8, []u8) = {
	let ln = if (delim is u8) {
		yield 1z;
	} else {
		let ln = len(delim: []u8);
		assert(ln > 0, "bytes::cut called with empty delimiter");
		yield ln;
	};
	match (index(in, delim)) {
	case let i: size =>
		return (in[..i], in[i + ln..]);
	case void =>
		return (in, []);
	};
};

// Returns the input slice "cut" along the last instance of a delimiter,
// returning everything up to the delimiter, and everything after the delimiter,
// in a tuple. The contents are borrowed from the input slice.
//
// The caller must ensure that 'delimiter' is not an empty slice.
export fn rcut(in: []u8, delim: ([]u8 | u8)) ([]u8, []u8) = {
	let ln = if (delim is u8) {
		yield 1z;
	} else {
		let ln = len(delim: []u8);
		assert(ln > 0, "bytes::rcut called with empty delimiter");
		yield ln;
	};
	match (rindex(in, delim)) {
	case let i: size =>
		return (in[..i], in[i + ln..]);
	case void =>
		return (in, []);
	};
};

@test fn cut() void = {
	const c = cut(['a', 'b', 'c'], ['b']);
	assert(equal(c.0, ['a']) && equal(c.1, ['c']));
	const c = cut(['a', 'b', 'c'], 'b');
	assert(equal(c.0, ['a']) && equal(c.1, ['c']));
	const c = cut(['a', 'b', 'c', 'b', 'a'], 'b');
	assert(equal(c.0, ['a']) && equal(c.1, ['c', 'b', 'a']));
	const c = cut(['a', 'b', 'c'], 'x');
	assert(equal(c.0, ['a', 'b', 'c']) && equal(c.1, []));
	const c = cut([], 'x');
	assert(equal(c.0, []) && equal(c.1, []));

	const c = rcut(['a', 'b', 'c'], ['b']);
	assert(equal(c.0, ['a']) && equal(c.1, ['c']));
	const c = rcut(['a', 'b', 'c', 'b', 'a'], 'b');
	assert(equal(c.0, ['a', 'b', 'c']) && equal(c.1, ['a']));
};