package iskra

import "git.smesh.lol/iskradb/lattice"

// WalkKind classifies the intent of a walk step.
type WalkKind uint8

const (
	WalkTranslate WalkKind = 0 // width stays constant
	WalkElaborate WalkKind = 1 // width grows toward MaxWidth
	WalkNarrow    WalkKind = 2 // width shrinks toward MinWidth
)

// WalkState is a single position in an in-progress lattice walk.
type WalkState struct {
	Key    lattice.Key // current lattice position
	Domain uint8       // current domain
	Coord  uint64      // accumulated coord context
	Word   string      // surface form at this position
	Score  uint32      // accumulated confidence (higher = better)
	Depth  uint8       // steps taken
}

// WalkCandidate is a scored next-step from WalkStep.
type WalkCandidate struct {
	State  WalkState
	Weight uint32 // bigram weight that produced this candidate
}

// WalkStep computes the set of candidate next positions from a current state.
// Uses BigramIdx for O(1) lookup of continuations from state.Word, then scores
// each by bigram weight with coord relaxation. Returns candidates ordered by
// descending weight, capped at maxCandidates.
func WalkStep(t *Tree, state WalkState, maxCandidates int32) []WalkCandidate {
	idxEntries := t.BigramIdx[state.Word]
	if len(idxEntries) == 0 {
		return nil
	}

	var candidates []WalkCandidate
	prefix := state.Word | "|"

	for _, ri := range idxEntries {
		if int32(ri) >= len(t.RecMeta) {
			continue
		}
		meta := &t.RecMeta[ri]
		if meta.Count == 0 {
			continue
		}
		rec := t.db.GetRecord(ri)
		if rec == nil {
			continue
		}
		form := FormFromRecord(rec, t.StringPool)
		if len(form) <= len(prefix) {
			continue
		}
		nextWord := form[len(prefix):]

		w, matchedCoord := BigramWeightRelaxed(t, state.Domain, state.Coord, state.Word, nextWord)
		if w == 0 {
			w = meta.Count
			matchedCoord = 0
		}

		nextKey := MakeKey(state.Domain, matchedCoord, nextWord)
		candidates = append(candidates, WalkCandidate{
			State: WalkState{
				Key:    nextKey,
				Domain: state.Domain,
				Coord:  matchedCoord,
				Word:   nextWord,
				Score:  state.Score + w,
				Depth:  state.Depth + 1,
			},
			Weight: w,
		})
	}

	sortCandidates(candidates)
	if maxCandidates > 0 && len(candidates) > maxCandidates {
		candidates = candidates[:maxCandidates]
	}
	return candidates
}

// WalkStepCrossDomain performs a walk step that crosses from srcDomain to dstDomain.
// This is the "translate" step: find the best match for state.Word in dstDomain
// using coord-relaxation and cross-domain links.
func WalkStepCrossDomain(t *Tree, state WalkState, dstDomain uint8) WalkState {
	src := NewSubLattice(t.db, t.StringPool, state.Domain, t.Reg)
	dst := NewSubLattice(t.db, t.StringPool, dstDomain, t.Reg)
	translated := Translate(src, dst, state.Word, state.Coord)
	if translated == "" {
		translated = state.Word
	}
	newKey := MakeKey(dstDomain, state.Coord, translated)
	return WalkState{
		Key:    newKey,
		Domain: dstDomain,
		Coord:  state.Coord,
		Word:   translated,
		Score:  state.Score,
		Depth:  state.Depth + 1,
	}
}

// Beam holds the top-K active walks in progress.
type Beam struct {
	Walks    []WalkState
	Width    int32
	MaxWidth int32
	MinWidth int32
	Kind     WalkKind
}

// NewBeam creates a beam with the initial state.
func NewBeam(initial WalkState, width int32, kind WalkKind) *Beam {
	return &Beam{
		Walks:    []WalkState{initial},
		Width:    width,
		MaxWidth: width * 4,
		MinWidth: 1,
		Kind:     kind,
	}
}

// Step advances all walks in the beam by one position.
// For WalkElaborate: width grows by 1 each step (up to MaxWidth).
// For WalkNarrow: width shrinks by 1 each step (down to MinWidth).
// For WalkTranslate: width stays constant.
func (b *Beam) Step(t *Tree) {
	var allCandidates []WalkCandidate

	for _, walk := range b.Walks {
		candidates := WalkStep(t, walk, b.Width)
		allCandidates = append(allCandidates, candidates...)
	}

	if len(allCandidates) == 0 {
		return
	}

	sortCandidates(allCandidates)

	switch b.Kind {
	case WalkElaborate:
		if b.Width < b.MaxWidth {
			b.Width++
		}
	case WalkNarrow:
		if b.Width > b.MinWidth {
			b.Width--
		}
	}

	limit := b.Width
	if len(allCandidates) < limit {
		limit = len(allCandidates)
	}
	b.Walks = []WalkState{:0:limit}
	for i := 0; i < limit; i++ {
		b.Walks = append(b.Walks, allCandidates[i].State)
	}
}

// Best returns the highest-scoring walk in the beam.
func (b *Beam) Best() WalkState {
	if len(b.Walks) == 0 {
		return WalkState{}
	}
	best := b.Walks[0]
	for i := 1; i < len(b.Walks); i++ {
		if b.Walks[i].Score > best.Score {
			best = b.Walks[i]
		}
	}
	return best
}

// Run executes the beam for maxSteps iterations, returning the best final state.
func (b *Beam) Run(t *Tree, maxSteps int32) WalkState {
	for step := 0; step < maxSteps; step++ {
		b.Step(t)
		if len(b.Walks) == 0 {
			break
		}
	}
	return b.Best()
}

// sortCandidates orders by descending weight (highest first).
func sortCandidates(cs []WalkCandidate) {
	for i := 1; i < len(cs); i++ {
		j := i
		for j > 0 && cs[j].Weight > cs[j-1].Weight {
			cs[j], cs[j-1] = cs[j-1], cs[j]
			j--
		}
	}
}