package winpwn

import "bytes"

const (
	cyclicAlphabet = "abcdefghijklmnopqrstuvwxyz"
	cyclicDefaultN = 4
)

// Cyclic returns the first length bytes of a de Bruijn sequence over the
// lowercase alphabet with subsequence length 4 (pwntools' default), the Go
// analogue of pwnlib.util.cyclic.cyclic -- a buffer where every 4-byte
// window is unique, so a crash address/register read back out of it tells
// you exactly how far into the buffer it pointed.
func Cyclic(length int) []byte {
	return CyclicN(length, cyclicDefaultN)
}

// CyclicN is Cyclic with an explicit subsequence length n instead of the
// default 4 (e.g. 8 to locate an offset into a 64-bit pointer-sized
// overwrite).
func CyclicN(length, n int) []byte {
	if length <= 0 {
		return []byte{}
	}
	out := make([]byte, 0, length)
	deBruijnEach(cyclicAlphabet, n, func(b byte) bool {
		out = append(out, b)
		return len(out) < length
	})
	if len(out) >= length || len(out) == 0 {
		return out[:length]
	}
	// The sequence's period (len(alphabet)^n) was shorter than the
	// requested length -- only possible for a tiny alphabet/n combination,
	// never in practice with the default 26-letter alphabet. Wrap.
	full := make([]byte, length)
	for i := range full {
		full[i] = out[i%len(out)]
	}
	return full
}

// CyclicFind returns the offset of subseq within the default (alphabet,
// n=4) de Bruijn sequence, or -1 if it can't appear in it -- the analogue
// of pwnlib.util.cyclic.cyclic_find. subseq is typically the 4 bytes read
// back from a crashed return address/register.
func CyclicFind(subseq []byte) int {
	return CyclicFindN(subseq, cyclicDefaultN)
}

// CyclicFindN is CyclicFind with an explicit subsequence length n, matching
// whatever n the buffer was generated with via CyclicN. Stops generating as
// soon as a match is found, so this stays fast even for n=8 (where the
// theoretical period, 26^8, is far too large to ever materialize) as long
// as subseq actually came from a real CyclicN(_, n) buffer, which is the
// only case this is ever used for.
func CyclicFindN(subseq []byte, n int) int {
	m := len(subseq)
	if m == 0 {
		return -1
	}

	window := make([]byte, 0, m)
	total := 0
	found := -1

	deBruijnEach(cyclicAlphabet, n, func(b byte) bool {
		if len(window) < m {
			window = append(window, b)
		} else {
			copy(window, window[1:])
			window[m-1] = b
		}
		total++
		if len(window) == m && bytes.Equal(window, subseq) {
			found = total - m
			return false
		}
		return true
	})

	return found
}

// deBruijnEach generates the de Bruijn sequence B(k, n) over alphabet (k =
// len(alphabet)) one byte at a time via yield, stopping as soon as yield
// returns false -- the Go equivalent of pwntools' de_bruijn() being a
// Python generator. This early-exit property is the entire point: B(k, n)'s
// period is k^n, which is astronomical for n=8 even with a small alphabet
// (26^8 ~ 2*10^11) and must never be materialized in full. Callers only
// ever need a bounded prefix (CyclicN) or an early match (CyclicFindN), and
// this lets both stop without walking the rest of the sequence. Uses the
// same Fredricksen-Kessler-Maiorana algorithm pwntools' cyclic() does.
func deBruijnEach(alphabet string, n int, yield func(byte) bool) {
	k := len(alphabet)
	a := make([]int, k*n)
	stop := false

	var db func(t, p int)
	db = func(t, p int) {
		if stop {
			return
		}
		if t > n {
			if n%p == 0 {
				for _, idx := range a[1 : p+1] {
					if !yield(alphabet[idx]) {
						stop = true
						return
					}
				}
			}
			return
		}
		a[t] = a[t-p]
		db(t+1, p)
		for j := a[t-p] + 1; j < k && !stop; j++ {
			a[t] = j
			db(t+1, t)
		}
	}
	db(1, 1)
}