package dict

import (
	"encoding/binary"
	"errors"
	"fmt"
)

const (
	MagicString                      = "JPDICT\x00\x00"
	CurrentVersion            uint32 = 2
	HeaderSize                       = 232
	KeyIndexRecordSize               = 12
	KanjiIndexRecordSize             = 8
	FuriganaIndexRecordSize          = 12
	KradIndexRecordSize              = 12
	RadkIndexRecordSize              = 12
	KanjiWordsIndexRecordSize        = 12
)

var (
	ErrBufferTooSmall = errors.New("buffer too small")
	ErrInvalidVarint  = errors.New("invalid varint")
	ErrInvalidMagic   = errors.New("invalid dictionary magic")
	ErrUnsupportedVer = errors.New("unsupported dictionary version")
)

type SectionHeader struct {
	Offset uint32
	Count  uint32
	Size   uint32
}

func (s *SectionHeader) Write(b []byte) {
	binary.LittleEndian.PutUint32(b[0:4], s.Offset)
	binary.LittleEndian.PutUint32(b[4:8], s.Count)
	binary.LittleEndian.PutUint32(b[8:12], s.Size)
}

func (s *SectionHeader) Read(b []byte) {
	s.Offset = binary.LittleEndian.Uint32(b[0:4])
	s.Count = binary.LittleEndian.Uint32(b[4:8])
	s.Size = binary.LittleEndian.Uint32(b[8:12])
}

type Header struct {
	Magic           [8]byte
	Version         uint32
	Reserved        uint32
	InternedStrings SectionHeader
	TagMap          SectionHeader
	KeyIndex        SectionHeader
	KeyData         SectionHeader
	PostingData     SectionHeader
	EntryIndex      SectionHeader
	EntryData       SectionHeader
	KanjiIndex      SectionHeader
	KanjiData       SectionHeader
	FuriganaIndex   SectionHeader
	FuriganaKeyData SectionHeader
	FuriganaData    SectionHeader
	KradIndex       SectionHeader
	KradData        SectionHeader
	RadkIndex       SectionHeader
	RadkData        SectionHeader
	KanjiWordsIndex SectionHeader
	KanjiWordsData  SectionHeader
}

func (h *Header) Encode() []byte {
	buf := make([]byte, HeaderSize)
	copy(buf[0:8], h.Magic[:])
	binary.LittleEndian.PutUint32(buf[8:12], h.Version)
	binary.LittleEndian.PutUint32(buf[12:16], h.Reserved)
	h.InternedStrings.Write(buf[16:28])
	h.TagMap.Write(buf[28:40])
	h.KeyIndex.Write(buf[40:52])
	h.KeyData.Write(buf[52:64])
	h.PostingData.Write(buf[64:76])
	h.EntryIndex.Write(buf[76:88])
	h.EntryData.Write(buf[88:100])
	h.KanjiIndex.Write(buf[100:112])
	h.KanjiData.Write(buf[112:124])
	h.FuriganaIndex.Write(buf[124:136])
	h.FuriganaKeyData.Write(buf[136:148])
	h.FuriganaData.Write(buf[148:160])
	h.KradIndex.Write(buf[160:172])
	h.KradData.Write(buf[172:184])
	h.RadkIndex.Write(buf[184:196])
	h.RadkData.Write(buf[196:208])
	h.KanjiWordsIndex.Write(buf[208:220])
	h.KanjiWordsData.Write(buf[220:232])
	return buf
}

func DecodeHeader(b []byte) (*Header, error) {
	if len(b) < 16 {
		return nil, ErrBufferTooSmall
	}
	var h Header
	copy(h.Magic[:], b[0:8])
	if string(h.Magic[:]) != MagicString {
		return nil, ErrInvalidMagic
	}
	h.Version = binary.LittleEndian.Uint32(b[8:12])
	if h.Version != CurrentVersion {
		return nil, fmt.Errorf("%w: version %d (want %d); run chirami setup", ErrUnsupportedVer, h.Version, CurrentVersion)
	}
	if len(b) < HeaderSize {
		return nil, ErrBufferTooSmall
	}
	h.Reserved = binary.LittleEndian.Uint32(b[12:16])
	h.InternedStrings.Read(b[16:28])
	h.TagMap.Read(b[28:40])
	h.KeyIndex.Read(b[40:52])
	h.KeyData.Read(b[52:64])
	h.PostingData.Read(b[64:76])
	h.EntryIndex.Read(b[76:88])
	h.EntryData.Read(b[88:100])
	h.KanjiIndex.Read(b[100:112])
	h.KanjiData.Read(b[112:124])
	h.FuriganaIndex.Read(b[124:136])
	h.FuriganaKeyData.Read(b[136:148])
	h.FuriganaData.Read(b[148:160])
	h.KradIndex.Read(b[160:172])
	h.KradData.Read(b[172:184])
	h.RadkIndex.Read(b[184:196])
	h.RadkData.Read(b[196:208])
	h.KanjiWordsIndex.Read(b[208:220])
	h.KanjiWordsData.Read(b[220:232])
	return &h, nil
}

type ByteWriter struct {
	Buf     []byte
	scratch [10]byte
}

func NewByteWriter(capacity int) *ByteWriter {
	return &ByteWriter{Buf: make([]byte, 0, capacity)}
}

func (w *ByteWriter) Reset() {
	w.Buf = w.Buf[:0]
}

func (w *ByteWriter) Uvarint(v uint64) {
	n := binary.PutUvarint(w.scratch[:], v)
	w.Buf = append(w.Buf, w.scratch[:n]...)
}

func (w *ByteWriter) U32(v uint32) {
	w.Uvarint(uint64(v))
}

func (w *ByteWriter) U8(v uint8) {
	w.Buf = append(w.Buf, v)
}

func (w *ByteWriter) Str(s string) {
	w.Uvarint(uint64(len(s)))
	w.Buf = append(w.Buf, s...)
}

func (w *ByteWriter) Interned(s string, internMap map[string]uint32) {
	id, ok := internMap[s]
	if !ok {
		panic(fmt.Sprintf("intern string not found: %q", s))
	}
	w.U32(id)
}

type ByteReader struct {
	B   []byte
	Off int
}

func (r *ByteReader) Uvarint() (uint64, error) {
	if r.Off >= len(r.B) {
		return 0, ErrBufferTooSmall
	}
	v, n := binary.Uvarint(r.B[r.Off:])
	if n <= 0 {
		return 0, ErrInvalidVarint
	}
	r.Off += n
	return v, nil
}

func (r *ByteReader) U32() (uint32, error) {
	v, err := r.Uvarint()
	return uint32(v), err
}

func (r *ByteReader) U8() (uint8, error) {
	if r.Off >= len(r.B) {
		return 0, ErrBufferTooSmall
	}
	b := r.B[r.Off]
	r.Off++
	return b, nil
}

func (r *ByteReader) Str() (string, error) {
	length, err := r.Uvarint()
	if err != nil {
		return "", err
	}
	l := int(length)
	if r.Off+l > len(r.B) {
		return "", ErrBufferTooSmall
	}
	s := string(r.B[r.Off : r.Off+l])
	r.Off += l
	return s, nil
}

func (r *ByteReader) Interned(interned []string) (string, error) {
	id, err := r.U32()
	if err != nil {
		return "", err
	}
	if int(id) < len(interned) {
		return interned[id], nil
	}
	return "", fmt.Errorf("intern id %d out of range %d", id, len(interned))
}

func EncodeEntry(w *ByteWriter, e *Entry, internMap map[string]uint32) {
	var idBytes [4]byte
	binary.LittleEndian.PutUint32(idBytes[:], e.ID)
	w.Buf = append(w.Buf, idBytes[:]...)
	w.Str(e.JLPT)

	// Kanji
	w.U32(uint32(len(e.Kanji)))
	for _, f := range e.Kanji {
		var flags uint8
		if f.Common {
			flags |= 1
		}
		w.U8(flags)
		w.Str(f.Text)
		w.U32(uint32(len(f.Tags)))
		for _, t := range f.Tags {
			w.Interned(t, internMap)
		}
	}

	// Kana
	w.U32(uint32(len(e.Kana)))
	for _, f := range e.Kana {
		var flags uint8
		if f.Common {
			flags |= 1
		}
		w.U8(flags)
		w.Str(f.Text)
		w.U32(uint32(len(f.Tags)))
		for _, t := range f.Tags {
			w.Interned(t, internMap)
		}
	}

	// Senses
	w.U32(uint32(len(e.Senses)))
	for _, s := range e.Senses {
		w.U32(uint32(len(s.POS)))
		for _, v := range s.POS {
			w.Interned(v, internMap)
		}
		w.U32(uint32(len(s.Misc)))
		for _, v := range s.Misc {
			w.Interned(v, internMap)
		}
		w.U32(uint32(len(s.Field)))
		for _, v := range s.Field {
			w.Interned(v, internMap)
		}
		w.U32(uint32(len(s.Dialect)))
		for _, v := range s.Dialect {
			w.Interned(v, internMap)
		}
		w.U32(uint32(len(s.Info)))
		for _, v := range s.Info {
			w.Str(v)
		}
		w.U32(uint32(len(s.Gloss)))
		for _, v := range s.Gloss {
			w.Str(v)
		}
		w.U32(uint32(len(s.AppliesToKanji)))
		for _, v := range s.AppliesToKanji {
			w.Interned(v, internMap)
		}
		w.U32(uint32(len(s.AppliesToKana)))
		for _, v := range s.AppliesToKana {
			w.Interned(v, internMap)
		}
	}

	// Pitch
	w.U32(uint32(len(e.Pitch)))
	for _, p := range e.Pitch {
		w.Str(p.Reading)
		w.U32(uint32(len(p.Downsteps)))
		for _, step := range p.Downsteps {
			w.Uvarint(uint64(step))
		}
	}
}

func DecodeEntry(b []byte, interned []string) (Entry, int, error) {
	if len(b) < 4 {
		return Entry{}, 0, ErrBufferTooSmall
	}
	r := &ByteReader{B: b}
	var e Entry
	e.ID = binary.LittleEndian.Uint32(b[0:4])
	r.Off = 4

	var err error
	if e.JLPT, err = r.Str(); err != nil {
		return Entry{}, 0, err
	}

	numKanji, err := r.U32()
	if err != nil {
		return Entry{}, 0, err
	}
	if numKanji > 0 {
		e.Kanji = make([]Form, numKanji)
		for i := range e.Kanji {
			flags, err := r.U8()
			if err != nil {
				return Entry{}, 0, err
			}
			e.Kanji[i].Common = (flags & 1) != 0
			if e.Kanji[i].Text, err = r.Str(); err != nil {
				return Entry{}, 0, err
			}
			numTags, err := r.U32()
			if err != nil {
				return Entry{}, 0, err
			}
			if numTags > 0 {
				e.Kanji[i].Tags = make([]string, numTags)
				for t := range e.Kanji[i].Tags {
					if e.Kanji[i].Tags[t], err = r.Interned(interned); err != nil {
						return Entry{}, 0, err
					}
				}
			}
		}
	}

	numKana, err := r.U32()
	if err != nil {
		return Entry{}, 0, err
	}
	if numKana > 0 {
		e.Kana = make([]Form, numKana)
		for i := range e.Kana {
			flags, err := r.U8()
			if err != nil {
				return Entry{}, 0, err
			}
			e.Kana[i].Common = (flags & 1) != 0
			if e.Kana[i].Text, err = r.Str(); err != nil {
				return Entry{}, 0, err
			}
			numTags, err := r.U32()
			if err != nil {
				return Entry{}, 0, err
			}
			if numTags > 0 {
				e.Kana[i].Tags = make([]string, numTags)
				for t := range e.Kana[i].Tags {
					if e.Kana[i].Tags[t], err = r.Interned(interned); err != nil {
						return Entry{}, 0, err
					}
				}
			}
		}
	}

	numSenses, err := r.U32()
	if err != nil {
		return Entry{}, 0, err
	}
	if numSenses > 0 {
		e.Senses = make([]Sense, numSenses)
		for s := range e.Senses {
			nPOS, err := r.U32()
			if err != nil {
				return Entry{}, 0, err
			}
			if nPOS > 0 {
				e.Senses[s].POS = make([]string, nPOS)
				for i := range e.Senses[s].POS {
					if e.Senses[s].POS[i], err = r.Interned(interned); err != nil {
						return Entry{}, 0, err
					}
				}
			}

			nMisc, err := r.U32()
			if err != nil {
				return Entry{}, 0, err
			}
			if nMisc > 0 {
				e.Senses[s].Misc = make([]string, nMisc)
				for i := range e.Senses[s].Misc {
					if e.Senses[s].Misc[i], err = r.Interned(interned); err != nil {
						return Entry{}, 0, err
					}
				}
			}

			nField, err := r.U32()
			if err != nil {
				return Entry{}, 0, err
			}
			if nField > 0 {
				e.Senses[s].Field = make([]string, nField)
				for i := range e.Senses[s].Field {
					if e.Senses[s].Field[i], err = r.Interned(interned); err != nil {
						return Entry{}, 0, err
					}
				}
			}

			nDialect, err := r.U32()
			if err != nil {
				return Entry{}, 0, err
			}
			if nDialect > 0 {
				e.Senses[s].Dialect = make([]string, nDialect)
				for i := range e.Senses[s].Dialect {
					if e.Senses[s].Dialect[i], err = r.Interned(interned); err != nil {
						return Entry{}, 0, err
					}
				}
			}

			nInfo, err := r.U32()
			if err != nil {
				return Entry{}, 0, err
			}
			if nInfo > 0 {
				e.Senses[s].Info = make([]string, nInfo)
				for i := range e.Senses[s].Info {
					if e.Senses[s].Info[i], err = r.Str(); err != nil {
						return Entry{}, 0, err
					}
				}
			}

			nGloss, err := r.U32()
			if err != nil {
				return Entry{}, 0, err
			}
			if nGloss > 0 {
				e.Senses[s].Gloss = make([]string, nGloss)
				for i := range e.Senses[s].Gloss {
					if e.Senses[s].Gloss[i], err = r.Str(); err != nil {
						return Entry{}, 0, err
					}
				}
			}

			nAKanji, err := r.U32()
			if err != nil {
				return Entry{}, 0, err
			}
			if nAKanji > 0 {
				e.Senses[s].AppliesToKanji = make([]string, nAKanji)
				for i := range e.Senses[s].AppliesToKanji {
					if e.Senses[s].AppliesToKanji[i], err = r.Interned(interned); err != nil {
						return Entry{}, 0, err
					}
				}
			}

			nAKana, err := r.U32()
			if err != nil {
				return Entry{}, 0, err
			}
			if nAKana > 0 {
				e.Senses[s].AppliesToKana = make([]string, nAKana)
				for i := range e.Senses[s].AppliesToKana {
					if e.Senses[s].AppliesToKana[i], err = r.Interned(interned); err != nil {
						return Entry{}, 0, err
					}
				}
			}
		}
	}

	numPitch, err := r.U32()
	if err != nil {
		return Entry{}, 0, err
	}
	if numPitch > 0 {
		e.Pitch = make([]Pitch, numPitch)
		for p := range e.Pitch {
			if e.Pitch[p].Reading, err = r.Str(); err != nil {
				return Entry{}, 0, err
			}
			nSteps, err := r.U32()
			if err != nil {
				return Entry{}, 0, err
			}
			if nSteps > 0 {
				e.Pitch[p].Downsteps = make([]int, nSteps)
				for i := range e.Pitch[p].Downsteps {
					step, err := r.Uvarint()
					if err != nil {
						return Entry{}, 0, err
					}
					e.Pitch[p].Downsteps[i] = int(step)
				}
			}
		}
	}

	return e, r.Off, nil
}

func EncodeKanji(w *ByteWriter, k *Kanji) {
	w.Str(k.Literal)
	w.Uvarint(uint64(k.Strokes))
	w.Uvarint(uint64(k.Grade))
	w.Uvarint(uint64(k.JLPT))
	w.Uvarint(uint64(k.Freq))

	w.U32(uint32(len(k.Meanings)))
	for _, m := range k.Meanings {
		w.Str(m)
	}

	w.U32(uint32(len(k.On)))
	for _, o := range k.On {
		w.Str(o)
	}

	w.U32(uint32(len(k.Kun)))
	for _, ku := range k.Kun {
		w.Str(ku)
	}

	w.U32(uint32(len(k.Nanori)))
	for _, na := range k.Nanori {
		w.Str(na)
	}
}

func DecodeKanji(b []byte) (Kanji, int, error) {
	r := &ByteReader{B: b}
	var k Kanji
	var err error
	if k.Literal, err = r.Str(); err != nil {
		return Kanji{}, 0, err
	}
	strokes, err := r.Uvarint()
	if err != nil {
		return Kanji{}, 0, err
	}
	k.Strokes = int(strokes)

	grade, err := r.Uvarint()
	if err != nil {
		return Kanji{}, 0, err
	}
	k.Grade = int(grade)

	jlpt, err := r.Uvarint()
	if err != nil {
		return Kanji{}, 0, err
	}
	k.JLPT = int(jlpt)

	freq, err := r.Uvarint()
	if err != nil {
		return Kanji{}, 0, err
	}
	k.Freq = int(freq)

	nMeanings, err := r.U32()
	if err != nil {
		return Kanji{}, 0, err
	}
	if nMeanings > 0 {
		k.Meanings = make([]string, nMeanings)
		for i := range k.Meanings {
			if k.Meanings[i], err = r.Str(); err != nil {
				return Kanji{}, 0, err
			}
		}
	}

	nOn, err := r.U32()
	if err != nil {
		return Kanji{}, 0, err
	}
	if nOn > 0 {
		k.On = make([]string, nOn)
		for i := range k.On {
			if k.On[i], err = r.Str(); err != nil {
				return Kanji{}, 0, err
			}
		}
	}

	nKun, err := r.U32()
	if err != nil {
		return Kanji{}, 0, err
	}
	if nKun > 0 {
		k.Kun = make([]string, nKun)
		for i := range k.Kun {
			if k.Kun[i], err = r.Str(); err != nil {
				return Kanji{}, 0, err
			}
		}
	}

	nNanori, err := r.U32()
	if err != nil {
		return Kanji{}, 0, err
	}
	if nNanori > 0 {
		k.Nanori = make([]string, nNanori)
		for i := range k.Nanori {
			if k.Nanori[i], err = r.Str(); err != nil {
				return Kanji{}, 0, err
			}
		}
	}

	return k, r.Off, nil
}

func EncodeRubies(w *ByteWriter, rubies []Ruby) {
	w.U32(uint32(len(rubies)))
	for _, r := range rubies {
		w.Str(r.Text)
		w.Str(r.Ruby)
	}
}

func DecodeRubies(b []byte) ([]Ruby, int, error) {
	r := &ByteReader{B: b}
	n, err := r.U32()
	if err != nil {
		return nil, 0, err
	}
	rubies := make([]Ruby, n)
	for i := range rubies {
		if rubies[i].Text, err = r.Str(); err != nil {
			return nil, 0, err
		}
		if rubies[i].Ruby, err = r.Str(); err != nil {
			return nil, 0, err
		}
	}
	return rubies, r.Off, nil
}