// Portions of this file are derived from github.com/ikawaha/kagome
// Copyright (c) 2014-2023 Ikawaha
// Licensed under the MIT License.

package morph

import (
	"encoding/binary"
	"unicode/utf8"
)

type nodeClass uint8

const (
	classDummy nodeClass = iota
	classKnown
	classUnknown
)

type latticeNode struct {
	id       int32
	position int // byte offset in input
	start    int // rune offset in input
	class    nodeClass
	cost     int32
	left     int16
	right    int16
	weight   int16
	surface  string
	prev     int32 // index in nodes slice, -1 for none
}

type lattice struct {
	nodes []latticeNode
	list  [][]int32
	path  []int32
}

func newLattice() *lattice {
	return &lattice{
		nodes: make([]latticeNode, 0, 256),
		list:  make([][]int32, 0, 64),
		path:  make([]int32, 0, 32),
	}
}

// Tokenize segments input string into IPADIC morphemes matching kagome Normal mode.
func (d *Dict) Tokenize(inp string) []Token {
	if len(inp) == 0 {
		return nil
	}
	rc := utf8.RuneCountInString(inp)

	la := d.pool.Get().(*lattice)
	la.nodes = la.nodes[:0]
	if cap(la.list) < rc+2 {
		la.list = make([][]int32, rc+2)
	} else {
		la.list = la.list[:rc+2]
		for i := range la.list {
			la.list[i] = la.list[i][:0]
		}
	}
	la.path = la.path[:0]

	addNode := func(runePos, id, position, start int, class nodeClass, surface string) {
		var left, right, weight int16
		switch class {
		case classKnown:
			off := id * 6
			left = int16(binary.LittleEndian.Uint16(d.morphBytes[off : off+2]))
			right = int16(binary.LittleEndian.Uint16(d.morphBytes[off+2 : off+4]))
			weight = int16(binary.LittleEndian.Uint16(d.morphBytes[off+4 : off+6]))
		case classUnknown:
			off := id * 6
			left = int16(binary.LittleEndian.Uint16(d.unkMorphs[off : off+2]))
			right = int16(binary.LittleEndian.Uint16(d.unkMorphs[off+2 : off+4]))
			weight = int16(binary.LittleEndian.Uint16(d.unkMorphs[off+4 : off+6]))
		}
		nodeIdx := int32(len(la.nodes))
		la.nodes = append(la.nodes, latticeNode{
			id:       int32(id),
			position: position,
			start:    start,
			class:    class,
			left:     left,
			right:    right,
			weight:   weight,
			surface:  surface,
			prev:     -1,
		})
		p := runePos + utf8.RuneCountInString(surface)
		la.list[p] = append(la.list[p], nodeIdx)
	}

	// BOS (rune 0)
	addNode(0, -1, 0, 0, classDummy, inp[0:0])
	// EOS (rune rc+1)
	addNode(rc+1, -1, len(inp), rc, classDummy, inp[len(inp):len(inp)])

	runePos := -1
	for pos, ch := range inp {
		runePos++
		anyMatches := false

		// Known words
		d.commonPrefixSearchCallback(inp[pos:], func(id, l int) {
			addNode(runePos, id, pos, runePos, classKnown, inp[pos:pos+l])
			anyMatches = true
		})

		// Unknown words
		class := d.charCategory(ch)
		ccInfo := &d.charClasses[class]
		if !anyMatches || ccInfo.Invoke {
			var endPos int
			if ch != utf8.RuneError {
				endPos = pos + utf8.RuneLen(ch)
			} else {
				endPos = pos + 1
			}
			unkWordLen := 1
			if ccInfo.Group {
				for i, w, size := endPos, 0, len(inp); i < size; i += w {
					var c rune
					c, w = utf8.DecodeRuneInString(inp[i:])
					if d.charCategory(c) != class {
						break
					}
					endPos += w
					unkWordLen++
					if unkWordLen >= 1024 {
						break
					}
				}
			}

			prev := pos
			if c, size := utf8.DecodeLastRuneInString(inp[pos:endPos]); c != utf8.RuneError {
				prev = endPos - size
			}
			id := ccInfo.UnkIndex
			dup := ccInfo.UnkDup
			for x := range int(dup) + 1 {
				if pos < prev {
					addNode(runePos, int(id)+x, pos, runePos, classUnknown, inp[pos:prev])
				}
				addNode(runePos, int(id)+x, pos, runePos, classUnknown, inp[pos:endPos])
			}
		}
	}

	// Forward pass (Viterbi)
	const maxCost = 1<<31 - 1
	for i := 1; i < len(la.list); i++ {
		currentList := la.list[i]
		for _, targetIdx := range currentList {
			target := &la.nodes[targetIdx]
			prevList := la.list[target.start]
			if len(prevList) == 0 {
				target.cost = maxCost
				continue
			}
			for j, prevIdx := range prevList {
				n := &la.nodes[prevIdx]
				c := d.connCost(n.right, target.left)
				totalCost := int64(c) + int64(target.weight) + int64(n.cost)
				if totalCost > maxCost {
					totalCost = maxCost
				}
				if j == 0 || int32(totalCost) < target.cost {
					target.cost = int32(totalCost)
					target.prev = prevIdx
				}
			}
		}
	}

	// Backward pass
	if len(la.list[rc+1]) == 0 {
		if cap(la.nodes) <= 16384 {
			d.pool.Put(la)
		}
		return nil
	}

	eosIdx := la.list[rc+1][0]
	p := la.nodes[eosIdx].prev
	for p != -1 {
		node := &la.nodes[p]
		if node.id == -1 { // BOS reached
			break
		}
		la.path = append(la.path, p)
		p = node.prev
	}

	tokens := make([]Token, len(la.path))
	for i := range la.path {
		node := &la.nodes[la.path[len(la.path)-1-i]]
		tok := Token{
			Surface: node.surface,
			Start:   node.position,
			End:     node.position + len(node.surface),
		}
		if node.class == classKnown {
			tok.Known = true
			featOff := uint32(node.id) * 16
			posID := d.featBytes[featOff]
			itID := d.featBytes[featOff+1]
			ifID := d.featBytes[featOff+2]
			baseOff := binary.LittleEndian.Uint32(d.featBytes[featOff+4 : featOff+8])
			readOff := binary.LittleEndian.Uint32(d.featBytes[featOff+8 : featOff+12])
			pronOff := binary.LittleEndian.Uint32(d.featBytes[featOff+12 : featOff+16])

			tok.POS = d.posTuples[posID]
			tok.InflType = d.inflTypes[itID]
			tok.InflForm = d.inflForms[ifID]
			tok.Base = d.stringAtOrSurface(baseOff, tok.Surface)
			tok.Reading = d.stringAt(readOff)
			if pronOff == readOff {
				tok.Pron = tok.Reading
			} else {
				tok.Pron = d.stringAt(pronOff)
			}
		} else {
			tok.Known = false
			tok.POS = d.unkFeatures[node.id]
			tok.Base = tok.Surface
		}
		tokens[i] = tok
	}

	if cap(la.nodes) <= 16384 {
		d.pool.Put(la)
	}

	return tokens
}