package lattice

const TritEnd = 0

// TritPath packs up to 16 trits into a uint32 (2 bits each).
// Trit 0 is the branch selector (1=noun, 2=verb, 3=modifier).
// Trits 1..15 encode the path within the branch.
type TritPath uint32

func (p TritPath) Trit(i int) uint8 {
	return uint8((p >> (uint(i) * 2)) & 0x3)
}

func (p TritPath) Depth() int {
	for i := 0; i < 16; i++ {
		if p.Trit(i) == TritEnd {
			return i
		}
	}
	return 16
}

func (p TritPath) Branch() uint8 {
	return p.Trit(0)
}

func MakeTritPath(trits ...uint8) TritPath {
	var p TritPath
	for i, t := range trits {
		if i >= 16 {
			break
		}
		p |= TritPath(t&0x3) << (uint(i) * 2)
	}
	return p
}

// BranchTritPath creates a TritPath with just the branch selector set.
func BranchTritPath(b Branch) TritPath {
	return TritPath(b + 1)
}

func Parent(p TritPath) TritPath {
	d := p.Depth()
	if d <= 1 {
		return 0
	}
	// zero out the last trit
	shift := uint(d-1) * 2
	mask := TritPath(0x3) << shift
	return p &^ mask
}

func Child(p TritPath, dir uint8) TritPath {
	d := p.Depth()
	if d >= 16 {
		return p
	}
	return p | TritPath(dir&0x3)<<(uint(d)*2)
}

// LCA returns the longest common prefix of two trit paths.
func LCA(a, b TritPath) TritPath {
	da := a.Depth()
	db := b.Depth()
	limit := da
	if db < limit {
		limit = db
	}
	shared := 0
	for i := 0; i < limit; i++ {
		if a.Trit(i) != b.Trit(i) {
			break
		}
		shared++
	}
	// mask out everything beyond shared prefix
	if shared == 0 {
		return 0
	}
	mask := TritPath((1 << (uint(shared) * 2)) - 1)
	return a & mask
}

// Distance returns the hop count between two trit paths via their LCA.
// Returns -1 if either path is unassigned.
func Distance(a, b TritPath) int {
	da := a.Depth()
	db := b.Depth()
	if da == 0 || db == 0 {
		return -1
	}
	shared := 0
	limit := da
	if db < limit {
		limit = db
	}
	for i := 0; i < limit; i++ {
		if a.Trit(i) != b.Trit(i) {
			break
		}
		shared++
	}
	return (da - shared) + (db - shared)
}

// Resolve applies an expansion multiplier to a TritPath's depth component.
// Each trit beyond the branch selector gets replicated `mult` times.
func Resolve(p TritPath, mult uint8) TritPath {
	if mult <= 1 {
		return p
	}
	d := p.Depth()
	if d <= 1 {
		return p
	}
	var out TritPath
	out = TritPath(p.Trit(0))
	pos := 1
	for i := 1; i < d; i++ {
		t := p.Trit(i)
		for j := uint8(0); j < mult && pos < 16; j++ {
			out |= TritPath(t&0x3) << (uint(pos) * 2)
			pos++
		}
	}
	return out
}