package mkmorph

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

	"chirami/internal/morph"
	"github.com/ikawaha/kagome-dict/ipa"
)

// Build extracts ipa.Dict() and writes a serialized morph.bin file to out.
func Build(out string) error {
	d := ipa.Dict()
	if d == nil {
		return fmt.Errorf("mkmorph: nil ipa dict")
	}

	strBuf := &bytes.Buffer{}
	strBuf.WriteByte(0) // offset 0 is ""
	strMap := make(map[string]uint32)
	strMap[""] = 0
	strMap["*"] = 0

	addStr := func(s string) uint32 {
		if s == "" || s == "*" {
			return 0
		}
		if off, ok := strMap[s]; ok {
			return off
		}
		off := uint32(strBuf.Len())
		strBuf.WriteString(s)
		strBuf.WriteByte(0)
		strMap[s] = off
		return off
	}

	for _, s := range d.POSTable.NameList {
		addStr(s)
	}
	for _, c := range d.Contents {
		for _, s := range c {
			addStr(s)
		}
	}
	for _, c := range d.UnkDict.Contents {
		for _, s := range c {
			addStr(s)
		}
	}
	for _, s := range d.CharClass {
		addStr(s)
	}

	// POS table (unique [4]string combinations)
	type posTuple [4]string
	posMap := make(map[posTuple]uint8)
	var posTuples []posTuple
	for _, p := range d.POSTable.POSs {
		var t posTuple
		for i := 0; i < 4; i++ {
			t[i] = d.POSTable.NameList[p[i]]
		}
		if _, ok := posMap[t]; !ok {
			if len(posTuples) >= 256 {
				return fmt.Errorf("mkmorph: POS table overflow (%d > 255)", len(posTuples))
			}
			posMap[t] = uint8(len(posTuples))
			posTuples = append(posTuples, t)
		}
	}

	// InflTypes table
	inflTypeMap := make(map[string]uint8)
	var inflTypes []string
	for _, c := range d.Contents {
		s := c[0]
		if _, ok := inflTypeMap[s]; !ok {
			if len(inflTypes) >= 256 {
				return fmt.Errorf("mkmorph: InflType table overflow (%d > 255)", len(inflTypes))
			}
			inflTypeMap[s] = uint8(len(inflTypes))
			inflTypes = append(inflTypes, s)
		}
	}

	// InflForms table
	inflFormMap := make(map[string]uint8)
	var inflForms []string
	for _, c := range d.Contents {
		s := c[1]
		if _, ok := inflFormMap[s]; !ok {
			if len(inflForms) >= 256 {
				return fmt.Errorf("mkmorph: InflForm table overflow (%d > 255)", len(inflForms))
			}
			inflFormMap[s] = uint8(len(inflForms))
			inflForms = append(inflForms, s)
		}
	}

	tmpOut := out + ".tmp"
	f, err := os.Create(tmpOut)
	if err != nil {
		return fmt.Errorf("mkmorph: create %s: %w", tmpOut, err)
	}
	defer func() {
		f.Close()
		_ = os.Remove(tmpOut)
	}()

	header := make([]byte, morph.HeaderSize)
	copy(header[0:8], []byte(morph.Magic))
	binary.LittleEndian.PutUint32(header[8:12], morph.CurrentVersion)

	offset := uint32(morph.HeaderSize)

	daOffset := offset
	daCount := uint32(len(d.Index.Da))
	offset += daCount * 8

	dupOffset := offset
	dupCount := uint32(len(d.Morphs))
	offset += dupCount * 1

	morphOffset := offset
	morphCount := uint32(len(d.Morphs))
	offset += morphCount * 6

	posTableOffset := offset
	posTableCount := uint32(len(posTuples))
	offset += posTableCount * 16

	inflTypeOffset := offset
	inflTypeCount := uint32(len(inflTypes))
	offset += inflTypeCount * 4

	inflFormOffset := offset
	inflFormCount := uint32(len(inflForms))
	offset += inflFormCount * 4

	featureOffset := offset
	featureCount := uint32(len(d.Contents))
	offset += featureCount * 16

	stringOffset := offset
	stringSize := uint32(strBuf.Len())
	offset += stringSize

	connOffset := offset
	connRows := uint32(d.Connection.Row)
	connCols := uint32(d.Connection.Col)
	offset += connRows * connCols * 2

	charCatOffset := offset
	charCatCount := uint32(len(d.CharCategory))
	offset += charCatCount

	charClassOffset := offset
	charClassCount := uint32(len(d.CharClass))
	offset += charClassCount * 16

	unkMorphOffset := offset
	unkMorphCount := uint32(len(d.UnkDict.Morphs))
	offset += unkMorphCount * 6

	unkFeatureOffset := offset
	unkFeatureCount := uint32(len(d.UnkDict.Contents))
	offset += unkFeatureCount * 16

	// Fill header table of contents
	binary.LittleEndian.PutUint32(header[16:20], daOffset)
	binary.LittleEndian.PutUint32(header[20:24], daCount)
	binary.LittleEndian.PutUint32(header[24:28], dupOffset)
	binary.LittleEndian.PutUint32(header[28:32], dupCount)
	binary.LittleEndian.PutUint32(header[32:36], morphOffset)
	binary.LittleEndian.PutUint32(header[36:40], morphCount)
	binary.LittleEndian.PutUint32(header[40:44], posTableOffset)
	binary.LittleEndian.PutUint32(header[44:48], posTableCount)
	binary.LittleEndian.PutUint32(header[48:52], inflTypeOffset)
	binary.LittleEndian.PutUint32(header[52:56], inflTypeCount)
	binary.LittleEndian.PutUint32(header[56:60], inflFormOffset)
	binary.LittleEndian.PutUint32(header[60:64], inflFormCount)
	binary.LittleEndian.PutUint32(header[64:68], featureOffset)
	binary.LittleEndian.PutUint32(header[68:72], featureCount)
	binary.LittleEndian.PutUint32(header[72:76], stringOffset)
	binary.LittleEndian.PutUint32(header[76:80], stringSize)
	binary.LittleEndian.PutUint32(header[80:84], connOffset)
	binary.LittleEndian.PutUint32(header[84:88], connRows)
	binary.LittleEndian.PutUint32(header[88:92], connCols)
	binary.LittleEndian.PutUint32(header[92:96], charCatOffset)
	binary.LittleEndian.PutUint32(header[96:100], charCatCount)
	binary.LittleEndian.PutUint32(header[100:104], charClassOffset)
	binary.LittleEndian.PutUint32(header[104:108], charClassCount)
	binary.LittleEndian.PutUint32(header[108:112], unkMorphOffset)
	binary.LittleEndian.PutUint32(header[112:116], unkMorphCount)
	binary.LittleEndian.PutUint32(header[116:120], unkFeatureOffset)
	binary.LittleEndian.PutUint32(header[120:124], unkFeatureCount)

	if _, err := f.Write(header); err != nil {
		return err
	}

	// 1. DA
	daBuf := make([]byte, daCount*8)
	for i := range d.Index.Da {
		binary.LittleEndian.PutUint32(daBuf[i*8:i*8+4], uint32(d.Index.Da[i].Base))
		binary.LittleEndian.PutUint32(daBuf[i*8+4:i*8+8], uint32(d.Index.Da[i].Check))
	}
	if _, err := f.Write(daBuf); err != nil {
		return err
	}

	// 2. Dup
	dupBuf := make([]byte, dupCount)
	for i := range dupCount {
		dupBuf[i] = uint8(d.Index.Dup[int32(i)])
	}
	if _, err := f.Write(dupBuf); err != nil {
		return err
	}

	// 3. Morphs
	morphBuf := make([]byte, morphCount*6)
	for i, m := range d.Morphs {
		binary.LittleEndian.PutUint16(morphBuf[i*6:i*6+2], uint16(m.LeftID))
		binary.LittleEndian.PutUint16(morphBuf[i*6+2:i*6+4], uint16(m.RightID))
		binary.LittleEndian.PutUint16(morphBuf[i*6+4:i*6+6], uint16(m.Weight))
	}
	if _, err := f.Write(morphBuf); err != nil {
		return err
	}

	// 4. POS Tuples
	posBuf := make([]byte, posTableCount*16)
	for i, pt := range posTuples {
		for j := 0; j < 4; j++ {
			binary.LittleEndian.PutUint32(posBuf[i*16+j*4:i*16+j*4+4], addStr(pt[j]))
		}
	}
	if _, err := f.Write(posBuf); err != nil {
		return err
	}

	// 5. InflTypes
	itBuf := make([]byte, inflTypeCount*4)
	for i, it := range inflTypes {
		binary.LittleEndian.PutUint32(itBuf[i*4:i*4+4], addStr(it))
	}
	if _, err := f.Write(itBuf); err != nil {
		return err
	}

	// 6. InflForms
	ifBuf := make([]byte, inflFormCount*4)
	for i, ifm := range inflForms {
		binary.LittleEndian.PutUint32(ifBuf[i*4:i*4+4], addStr(ifm))
	}
	if _, err := f.Write(ifBuf); err != nil {
		return err
	}

	// 7. Features
	featBuf := make([]byte, featureCount*16)
	for i, c := range d.Contents {
		p := d.POSTable.POSs[i]
		var t posTuple
		for j := 0; j < 4; j++ {
			t[j] = d.POSTable.NameList[p[j]]
		}
		featBuf[i*16] = posMap[t]
		featBuf[i*16+1] = inflTypeMap[c[0]]
		featBuf[i*16+2] = inflFormMap[c[1]]
		featBuf[i*16+3] = 0 // reserved
		binary.LittleEndian.PutUint32(featBuf[i*16+4:i*16+8], addStr(c[2]))
		binary.LittleEndian.PutUint32(featBuf[i*16+8:i*16+12], addStr(c[3]))
		binary.LittleEndian.PutUint32(featBuf[i*16+12:i*16+16], addStr(c[4]))
	}
	if _, err := f.Write(featBuf); err != nil {
		return err
	}

	// 8. Strings
	if _, err := f.Write(strBuf.Bytes()); err != nil {
		return err
	}

	// 9. Connection
	connBuf := make([]byte, len(d.Connection.Vec)*2)
	for i, v := range d.Connection.Vec {
		binary.LittleEndian.PutUint16(connBuf[i*2:i*2+2], uint16(v))
	}
	if _, err := f.Write(connBuf); err != nil {
		return err
	}

	// 10. CharCat
	if _, err := f.Write(d.CharCategory); err != nil {
		return err
	}

	// 11. CharClass
	ccBuf := make([]byte, charClassCount*16)
	for i := range charClassCount {
		if d.InvokeList[i] {
			ccBuf[i*16] = 1
		}
		if d.GroupList[i] {
			ccBuf[i*16+1] = 1
		}
		binary.LittleEndian.PutUint32(ccBuf[i*16+4:i*16+8], uint32(d.UnkDict.Index[int32(i)]))
		binary.LittleEndian.PutUint32(ccBuf[i*16+8:i*16+12], uint32(d.UnkDict.IndexDup[int32(i)]))
		binary.LittleEndian.PutUint32(ccBuf[i*16+12:i*16+16], addStr(d.CharClass[i]))
	}
	if _, err := f.Write(ccBuf); err != nil {
		return err
	}

	// 12. UnkMorph
	umBuf := make([]byte, unkMorphCount*6)
	for i, m := range d.UnkDict.Morphs {
		binary.LittleEndian.PutUint16(umBuf[i*6:i*6+2], uint16(m.LeftID))
		binary.LittleEndian.PutUint16(umBuf[i*6+2:i*6+4], uint16(m.RightID))
		binary.LittleEndian.PutUint16(umBuf[i*6+4:i*6+6], uint16(m.Weight))
	}
	if _, err := f.Write(umBuf); err != nil {
		return err
	}

	// 13. UnkFeature
	ufBuf := make([]byte, unkFeatureCount*16)
	for i, c := range d.UnkDict.Contents {
		for j := 0; j < 4; j++ {
			binary.LittleEndian.PutUint32(ufBuf[i*16+j*4:i*16+j*4+4], addStr(c[j]))
		}
	}
	if _, err := f.Write(ufBuf); err != nil {
		return err
	}

	if err := f.Sync(); err != nil {
		return err
	}
	if err := f.Close(); err != nil {
		return err
	}

	if err := os.Rename(tmpOut, out); err != nil {
		return fmt.Errorf("mkmorph: rename %s -> %s: %w", tmpOut, out, err)
	}

	return nil
}