package morph

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"os"
	"sync"

	"golang.org/x/sys/unix"
)

// Magic is the 8-byte magic header of the morph binary format.
const Magic = "JPMORPH\x00"

// CurrentVersion is the binary format version.
const CurrentVersion uint32 = 1

// HeaderSize is the fixed size of the binary file header.
const HeaderSize = 128

// Header holds the section offsets and counts within a morph.bin file.
type Header struct {
	DaOffset, DaCount                 uint32
	DupOffset, DupCount               uint32
	MorphOffset, MorphCount           uint32
	PosOffset, PosCount               uint32
	InflTypeOffset, InflTypeCount     uint32
	InflFormOffset, InflFormCount     uint32
	FeatureOffset, FeatureCount       uint32
	StringOffset, StringSize          uint32
	ConnOffset, ConnRows, ConnCols    uint32
	CharCatOffset, CharCatCount       uint32
	CharClassOffset, CharClassCount   uint32
	UnkMorphOffset, UnkMorphCount     uint32
	UnkFeatureOffset, UnkFeatureCount uint32
}

// CharClassInfo holds character category settings for unknown-word processing.
type CharClassInfo struct {
	Invoke   bool
	Group    bool
	UnkIndex int32
	UnkDup   int32
}

// Dict is an mmap-backed IPADIC dictionary for morphological analysis.
type Dict struct {
	data        []byte
	h           Header
	daBytes     []byte
	dupBytes    []byte
	morphBytes  []byte
	featBytes   []byte
	stringBytes []byte
	connBytes   []byte
	charCat     []byte
	charClasses [11]CharClassInfo
	unkMorphs   []byte

	// Cached string tables populated on Open for zero-alloc tokenization.
	posTuples   [][4]string
	inflTypes   []string
	inflForms   []string
	unkFeatures [][4]string

	pool sync.Pool
}

// Open memory-maps a morph.bin dictionary file and validates its magic,
// version, and section bounds.
func Open(path string) (*Dict, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	fi, err := f.Stat()
	if err != nil {
		return nil, err
	}
	sz := int(fi.Size())
	if sz < HeaderSize {
		return nil, fmt.Errorf("morph: file too small (%d bytes)", sz)
	}

	data, err := unix.Mmap(int(f.Fd()), 0, sz, unix.PROT_READ, unix.MAP_SHARED)
	if err != nil {
		return nil, fmt.Errorf("morph: mmap: %w", err)
	}

	magic := string(data[:8])
	if magic != Magic {
		_ = unix.Munmap(data)
		return nil, fmt.Errorf("morph: invalid magic %q", magic)
	}
	version := binary.LittleEndian.Uint32(data[8:12])
	if version != CurrentVersion {
		_ = unix.Munmap(data)
		return nil, fmt.Errorf("morph: unsupported version %d", version)
	}

	d := &Dict{data: data}
	h := &d.h
	h.DaOffset = binary.LittleEndian.Uint32(data[16:20])
	h.DaCount = binary.LittleEndian.Uint32(data[20:24])
	h.DupOffset = binary.LittleEndian.Uint32(data[24:28])
	h.DupCount = binary.LittleEndian.Uint32(data[28:32])
	h.MorphOffset = binary.LittleEndian.Uint32(data[32:36])
	h.MorphCount = binary.LittleEndian.Uint32(data[36:40])
	h.PosOffset = binary.LittleEndian.Uint32(data[40:44])
	h.PosCount = binary.LittleEndian.Uint32(data[44:48])
	h.InflTypeOffset = binary.LittleEndian.Uint32(data[48:52])
	h.InflTypeCount = binary.LittleEndian.Uint32(data[52:56])
	h.InflFormOffset = binary.LittleEndian.Uint32(data[56:60])
	h.InflFormCount = binary.LittleEndian.Uint32(data[60:64])
	h.FeatureOffset = binary.LittleEndian.Uint32(data[64:68])
	h.FeatureCount = binary.LittleEndian.Uint32(data[68:72])
	h.StringOffset = binary.LittleEndian.Uint32(data[72:76])
	h.StringSize = binary.LittleEndian.Uint32(data[76:80])
	h.ConnOffset = binary.LittleEndian.Uint32(data[80:84])
	h.ConnRows = binary.LittleEndian.Uint32(data[84:88])
	h.ConnCols = binary.LittleEndian.Uint32(data[88:92])
	h.CharCatOffset = binary.LittleEndian.Uint32(data[92:96])
	h.CharCatCount = binary.LittleEndian.Uint32(data[96:100])
	h.CharClassOffset = binary.LittleEndian.Uint32(data[100:104])
	h.CharClassCount = binary.LittleEndian.Uint32(data[104:108])
	h.UnkMorphOffset = binary.LittleEndian.Uint32(data[108:112])
	h.UnkMorphCount = binary.LittleEndian.Uint32(data[112:116])
	h.UnkFeatureOffset = binary.LittleEndian.Uint32(data[116:120])
	h.UnkFeatureCount = binary.LittleEndian.Uint32(data[120:124])

	checkBounds := func(offset, count, elemSize uint32) error {
		total := uint64(offset) + uint64(count)*uint64(elemSize)
		if total > uint64(sz) {
			return fmt.Errorf("morph: section exceeds file bounds")
		}
		return nil
	}

	if err := checkBounds(h.DaOffset, h.DaCount, 8); err != nil {
		_ = unix.Munmap(data)
		return nil, err
	}
	if err := checkBounds(h.DupOffset, h.DupCount, 1); err != nil {
		_ = unix.Munmap(data)
		return nil, err
	}
	if err := checkBounds(h.MorphOffset, h.MorphCount, 6); err != nil {
		_ = unix.Munmap(data)
		return nil, err
	}
	if err := checkBounds(h.PosOffset, h.PosCount, 16); err != nil {
		_ = unix.Munmap(data)
		return nil, err
	}
	if err := checkBounds(h.InflTypeOffset, h.InflTypeCount, 4); err != nil {
		_ = unix.Munmap(data)
		return nil, err
	}
	if err := checkBounds(h.InflFormOffset, h.InflFormCount, 4); err != nil {
		_ = unix.Munmap(data)
		return nil, err
	}
	if err := checkBounds(h.FeatureOffset, h.FeatureCount, 16); err != nil {
		_ = unix.Munmap(data)
		return nil, err
	}
	if err := checkBounds(h.StringOffset, h.StringSize, 1); err != nil {
		_ = unix.Munmap(data)
		return nil, err
	}
	if err := checkBounds(h.ConnOffset, h.ConnRows*h.ConnCols, 2); err != nil {
		_ = unix.Munmap(data)
		return nil, err
	}
	if err := checkBounds(h.CharCatOffset, h.CharCatCount, 1); err != nil {
		_ = unix.Munmap(data)
		return nil, err
	}
	if err := checkBounds(h.CharClassOffset, h.CharClassCount, 16); err != nil {
		_ = unix.Munmap(data)
		return nil, err
	}
	if err := checkBounds(h.UnkMorphOffset, h.UnkMorphCount, 6); err != nil {
		_ = unix.Munmap(data)
		return nil, err
	}
	if err := checkBounds(h.UnkFeatureOffset, h.UnkFeatureCount, 16); err != nil {
		_ = unix.Munmap(data)
		return nil, err
	}

	d.daBytes = data[h.DaOffset : h.DaOffset+h.DaCount*8]
	d.dupBytes = data[h.DupOffset : h.DupOffset+h.DupCount]
	d.morphBytes = data[h.MorphOffset : h.MorphOffset+h.MorphCount*6]
	d.featBytes = data[h.FeatureOffset : h.FeatureOffset+h.FeatureCount*16]
	d.stringBytes = data[h.StringOffset : h.StringOffset+h.StringSize]
	d.connBytes = data[h.ConnOffset : h.ConnOffset+h.ConnRows*h.ConnCols*2]
	d.charCat = data[h.CharCatOffset : h.CharCatOffset+h.CharCatCount]
	d.unkMorphs = data[h.UnkMorphOffset : h.UnkMorphOffset+h.UnkMorphCount*6]

	// Cache POS tuples
	d.posTuples = make([][4]string, h.PosCount)
	posData := data[h.PosOffset : h.PosOffset+h.PosCount*16]
	for i := range h.PosCount {
		for j := 0; j < 4; j++ {
			off := binary.LittleEndian.Uint32(posData[i*16+uint32(j)*4 : i*16+uint32(j)*4+4])
			d.posTuples[i][j] = d.stringAt(off)
		}
	}

	// Cache InflTypes
	d.inflTypes = make([]string, h.InflTypeCount)
	itData := data[h.InflTypeOffset : h.InflTypeOffset+h.InflTypeCount*4]
	for i := range h.InflTypeCount {
		off := binary.LittleEndian.Uint32(itData[i*4 : i*4+4])
		d.inflTypes[i] = d.stringAt(off)
	}

	// Cache InflForms
	d.inflForms = make([]string, h.InflFormCount)
	ifData := data[h.InflFormOffset : h.InflFormOffset+h.InflFormCount*4]
	for i := range h.InflFormCount {
		off := binary.LittleEndian.Uint32(ifData[i*4 : i*4+4])
		d.inflForms[i] = d.stringAt(off)
	}

	// Cache Unknown POS features
	d.unkFeatures = make([][4]string, h.UnkFeatureCount)
	ufData := data[h.UnkFeatureOffset : h.UnkFeatureOffset+h.UnkFeatureCount*16]
	for i := range h.UnkFeatureCount {
		for j := 0; j < 4; j++ {
			off := binary.LittleEndian.Uint32(ufData[i*16+uint32(j)*4 : i*16+uint32(j)*4+4])
			d.unkFeatures[i][j] = d.stringAt(off)
		}
	}

	// Unknown character classes
	ccData := data[h.CharClassOffset : h.CharClassOffset+h.CharClassCount*16]
	for i := range int(h.CharClassCount) {
		if i >= len(d.charClasses) {
			break
		}
		d.charClasses[i].Invoke = ccData[i*16] != 0
		d.charClasses[i].Group = ccData[i*16+1] != 0
		d.charClasses[i].UnkIndex = int32(binary.LittleEndian.Uint32(ccData[i*16+4 : i*16+8]))
		d.charClasses[i].UnkDup = int32(binary.LittleEndian.Uint32(ccData[i*16+8 : i*16+12]))
	}

	d.pool.New = func() any {
		return newLattice()
	}

	return d, nil
}

// Close unmaps the dictionary data.
func (d *Dict) Close() error {
	if d.data == nil {
		return nil
	}
	err := unix.Munmap(d.data)
	d.data = nil
	return err
}

func (d *Dict) stringAt(off uint32) string {
	if off == 0 || int(off) >= len(d.stringBytes) {
		return ""
	}
	b := d.stringBytes[off:]
	end := bytes.IndexByte(b, 0)
	if end < 0 {
		return ""
	}
	return string(b[:end])
}

func (d *Dict) stringAtOrSurface(off uint32, surface string) string {
	if off == 0 || int(off) >= len(d.stringBytes) {
		return ""
	}
	b := d.stringBytes[off:]
	slen := len(surface)
	if len(b) > slen && b[slen] == 0 && string(b[:slen]) == surface {
		return surface
	}
	end := bytes.IndexByte(b, 0)
	if end < 0 {
		return ""
	}
	return string(b[:end])
}

func (d *Dict) daBase(i int) int32 {
	return int32(binary.LittleEndian.Uint32(d.daBytes[i*8 : i*8+4]))
}

func (d *Dict) daCheck(i int) int32 {
	return int32(binary.LittleEndian.Uint32(d.daBytes[i*8+4 : i*8+8]))
}

func (d *Dict) charCategory(r rune) byte {
	if int(r) < len(d.charCat) {
		return d.charCat[r]
	}
	return d.charCat[0]
}

func (d *Dict) connCost(right, left int16) int16 {
	idx := int(left)*int(d.h.ConnRows) + int(right)
	return int16(binary.LittleEndian.Uint16(d.connBytes[idx*2 : idx*2+2]))
}

func (d *Dict) commonPrefixSearchCallback(input string, callback func(id, l int)) {
	var p, q int
	bufLen := int(d.h.DaCount)
	for i, size := 0, len(input); i < size; i++ {
		c := input[i]
		if c == 0 {
			return
		}
		p = q
		q = int(d.daBase(p)) + int(c)
		if q < 0 || q >= bufLen || int(d.daCheck(q)) != p {
			break
		}
		ahead := int(d.daBase(q))
		if ahead >= 0 && ahead < bufLen && int(d.daCheck(ahead)) == q {
			baseVal := d.daBase(ahead)
			if baseVal <= 0 {
				id := int(-baseVal)
				dup := int(d.dupBytes[id])
				for k := id; k <= id+dup; k++ {
					callback(k, i+1)
				}
			}
		}
	}
}