package lattice

type ExportFilter func(nodeIdx uint32, n *Node, recIdx uint32, r *Record) bool

type SubLattice struct {
	Tree      *Tree
	OriginKey Key
	SrcBranch Branch
}

type keyRec struct {
	key    Key
	recIdx uint32
	branch Branch
}

func (t *Tree) collectBranch(rootIdx uint32, branch Branch, filter ExportFilter) []keyRec {
	var out []keyRec
	queue := []uint32{rootIdx}
	for len(queue) > 0 {
		idx := queue[0]
		queue = queue[1:]
		if idx == NullIndex {
			continue
		}
		node := t.getNode(idx)
		if node.IsDeleted() {
			continue
		}
		k := node.KeyCount()
		for i := 0; i < k; i++ {
			ri := node.RecPtrs[i]
			if ri == NullRec {
				continue
			}
			if filter != nil {
				rec := t.getRecord(ri)
				if !filter(idx, node, ri, rec) {
					continue
				}
			}
			out = append(out, keyRec{key: node.Keys[i], recIdx: ri, branch: branch})
		}
		for i := 0; i <= k; i++ {
			if node.Children[i] != NullIndex {
				queue = append(queue, node.Children[i])
			}
		}
	}
	return out
}

func (t *Tree) extractPairs(pairs []keyRec) *SubLattice {
	dst := NewTree(len(pairs) + 16)
	remap := map[uint32]uint32{}
	for _, p := range pairs {
		rec := *t.getRecord(p.recIdx)
		rec.Link[0] = NullRec
		rec.Link[1] = NullRec
		dst.Insert(p.branch, p.key, rec)
		newRI := dst.LookupRecIdx(p.branch, p.key)
		remap[p.recIdx] = newRI
	}
	for _, p := range pairs {
		srcRec := t.getRecord(p.recIdx)
		newRI, ok := remap[p.recIdx]
		if !ok {
			continue
		}
		dstRec := dst.getRecord(newRI)
		for li := 0; li < 2; li++ {
			if srcRec.Link[li] == NullRec {
				continue
			}
			if mapped, ok := remap[srcRec.Link[li]]; ok {
				dstRec.Link[li] = mapped
			}
		}
		dstRec.Branch = srcRec.Branch
	}
	return &SubLattice{Tree: dst}
}

func (t *Tree) ExtractBranch(branch Branch) *SubLattice {
	pairs := t.collectBranch(t.Roots[branch], branch, nil)
	sl := t.extractPairs(pairs)
	sl.SrcBranch = branch
	return sl
}

func (t *Tree) ExtractSubtree(rootNodeIdx uint32, branch Branch, filter ExportFilter) *SubLattice {
	pairs := t.collectBranch(rootNodeIdx, branch, filter)
	sl := t.extractPairs(pairs)
	sl.SrcBranch = branch
	return sl
}

func (t *Tree) ExtractLinked(startRecIdx uint32) *SubLattice {
	visited := map[uint32]bool{}
	queue := []uint32{startRecIdx}
	visited[startRecIdx] = true
	for len(queue) > 0 {
		ri := queue[0]
		queue = queue[1:]
		rec := t.getRecord(ri)
		for li := 0; li < 2; li++ {
			partner := rec.Link[li]
			if partner == NullRec || visited[partner] {
				continue
			}
			visited[partner] = true
			queue = append(queue, partner)
		}
	}
	var pairs []keyRec
	for ri := range visited {
		key, ok := t.RecKey[ri]
		if !ok {
			continue
		}
		rec := t.getRecord(ri)
		pairs = append(pairs, keyRec{key: key, recIdx: ri, branch: Branch(rec.Branch)})
	}
	sl := t.extractPairs(pairs)
	if startRecIdx != NullRec {
		sl.OriginKey = t.RecKey[startRecIdx]
		sl.SrcBranch = Branch(t.getRecord(startRecIdx).Branch)
	}
	return sl
}

func (t *Tree) Graft(sl *SubLattice) int {
	var allPairs []keyRec
	for b := 0; b < 3; b++ {
		br := Branch(b)
		pairs := sl.Tree.collectBranch(sl.Tree.Roots[br], br, nil)
		allPairs = append(allPairs, pairs...)
	}
	remap := map[uint32]uint32{}
	inserted := 0
	for _, p := range allPairs {
		existing := t.LookupRecIdx(p.branch, p.key)
		if existing != NullRec {
			remap[p.recIdx] = existing
			continue
		}
		rec := *sl.Tree.getRecord(p.recIdx)
		rec.Link[0] = NullRec
		rec.Link[1] = NullRec
		t.Insert(p.branch, p.key, rec)
		newRI := t.LookupRecIdx(p.branch, p.key)
		remap[p.recIdx] = newRI
		inserted++
	}
	for _, p := range allPairs {
		srcRec := sl.Tree.getRecord(p.recIdx)
		newRI, ok := remap[p.recIdx]
		if !ok {
			continue
		}
		dstRec := t.getRecord(newRI)
		for li := 0; li < 2; li++ {
			if srcRec.Link[li] == NullRec {
				continue
			}
			if mapped, ok := remap[srcRec.Link[li]]; ok {
				dstRec.Link[li] = mapped
				t.markRecDirty(newRI)
			}
		}
		if srcRec.Branch != 0 || dstRec.Branch == 0 {
			dstRec.Branch = srcRec.Branch
		}
	}
	return inserted
}

func (sl *SubLattice) SaveFile(path string) error {
	return SaveFile(path, sl.Tree)
}

func LoadSubLattice(path string) (*SubLattice, error) {
	t, err := LoadFile(path)
	if err != nil {
		return nil, err
	}
	return &SubLattice{Tree: t}, nil
}