package lattice

import (
	"encoding/binary"
	"fmt"
	"io"
	"os"
)

var le = binary.LittleEndian

// nodeSize is the on-disk size of a C=8 Node.
// 8×16B keys + 8×4B RecPtrs + 9×4B Children + 4B (Mult, Flags, pad) = 200 bytes.
const nodeSize = C*16 + C*4 + (C+1)*4 + 4

// headerSize for version 4: magic(4) + version(2) + nCount(4) + rCount(4) +
// pendCount(4) + 8×roots(32) + freeHead(4) + freeCount(4) + nodeCap(4) + recCap(4) = 66 → pad to 72.
const headerSize = 72

type iskrHeader struct {
	nCount, rCount, pendCount uint32
	roots                     [C]uint32
	freeHead                  uint32
	freeCount                 uint32
	nodeCap                   uint32
	recCap                    uint32
}

func writeHeaderBuf(t *Tree, nodeCap, recCap uint32) [headerSize]byte {
	var buf [headerSize]byte
	copy(buf[0:4], "ISKR")
	le.PutUint16(buf[4:6], 4) // version 4 = C=8
	le.PutUint32(buf[6:10], t.nCount)
	le.PutUint32(buf[10:14], t.rCount)
	le.PutUint32(buf[14:18], uint32(len(t.Pending)))
	for i := 0; i < C; i++ {
		le.PutUint32(buf[18+i*4:], t.Roots[i])
	}
	// offset 18 + 8*4 = 50
	le.PutUint32(buf[50:54], t.FreeHead)
	le.PutUint32(buf[54:58], t.FreeCount)
	le.PutUint32(buf[58:62], nodeCap)
	le.PutUint32(buf[62:66], recCap)
	// bytes 66-71 = padding (zero)
	return buf
}

func readHeaderBuf(buf [headerSize]byte) (iskrHeader, error) {
	if string(buf[0:4]) != "ISKR" {
		return iskrHeader{}, fmt.Errorf("bad magic: %q", buf[0:4])
	}
	v := le.Uint16(buf[4:6])
	if v != 4 {
		return iskrHeader{}, fmt.Errorf("unsupported version %d (want 4)", v)
	}
	hdr := iskrHeader{
		nCount:    le.Uint32(buf[6:10]),
		rCount:    le.Uint32(buf[10:14]),
		pendCount: le.Uint32(buf[14:18]),
		freeHead:  le.Uint32(buf[50:54]),
		freeCount: le.Uint32(buf[54:58]),
		nodeCap:   le.Uint32(buf[58:62]),
		recCap:    le.Uint32(buf[62:66]),
	}
	for i := 0; i < C; i++ {
		hdr.roots[i] = le.Uint32(buf[18+i*4:])
	}
	return hdr, nil
}

func writeHeader(w io.Writer, t *Tree) error {
	buf := writeHeaderBuf(t, t.nCount, t.rCount)
	_, err := w.Write(buf[:])
	return err
}

func readHeader(r io.Reader) (iskrHeader, error) {
	var buf [headerSize]byte
	if _, err := io.ReadFull(r, buf[:]); err != nil {
		return iskrHeader{}, fmt.Errorf("read header: %w", err)
	}
	return readHeaderBuf(buf)
}

// Node is now nodeSize bytes: 8×16B keys + 8×4B RecPtrs + 9×4B Children + 4B flags/pad.
func writeNodeBuf(n *Node) [nodeSize]byte {
	var buf [nodeSize]byte
	for i := 0; i < C; i++ {
		le.PutUint64(buf[i*16:], n.Keys[i][0])
		le.PutUint64(buf[i*16+8:], n.Keys[i][1])
	}
	// RecPtrs at offset C*16 = 128
	for i := 0; i < C; i++ {
		le.PutUint32(buf[128+i*4:], n.RecPtrs[i])
	}
	// Children at offset 128 + C*4 = 160
	for i := 0; i < C+1; i++ {
		le.PutUint32(buf[160+i*4:], n.Children[i])
	}
	// Mult and Flags at offset 160 + (C+1)*4 = 196
	buf[196] = n.Mult
	buf[197] = n.Flags
	return buf
}

func readNodeBuf(buf [nodeSize]byte) Node {
	var n Node
	for i := 0; i < C; i++ {
		n.Keys[i][0] = le.Uint64(buf[i*16:])
		n.Keys[i][1] = le.Uint64(buf[i*16+8:])
	}
	for i := 0; i < C; i++ {
		n.RecPtrs[i] = le.Uint32(buf[128+i*4:])
	}
	for i := 0; i < C+1; i++ {
		n.Children[i] = le.Uint32(buf[160+i*4:])
	}
	n.Mult = buf[196]
	n.Flags = buf[197]
	return n
}

func writeRecBuf(r *Record) [48]byte {
	var buf [48]byte
	le.PutUint32(buf[0:], r.DataFile)
	le.PutUint32(buf[4:], r.DataOff)
	le.PutUint32(buf[8:], r.DataLen)
	le.PutUint32(buf[12:], r.Link[0])
	le.PutUint32(buf[16:], r.Link[1])
	buf[20] = uint8(r.Branch)
	copy(buf[24:], r.Inline[:])
	return buf
}

func readRecBuf(buf [48]byte) Record {
	var r Record
	r.DataFile = le.Uint32(buf[0:])
	r.DataOff = le.Uint32(buf[4:])
	r.DataLen = le.Uint32(buf[8:])
	r.Link[0] = le.Uint32(buf[12:])
	r.Link[1] = le.Uint32(buf[16:])
	r.Branch = buf[20]
	copy(r.Inline[:], buf[24:])
	return r
}

func writeNode(w io.Writer, n *Node) error {
	buf := writeNodeBuf(n)
	_, err := w.Write(buf[:])
	return err
}

func readNode(r io.Reader) (Node, error) {
	var buf [nodeSize]byte
	if _, err := io.ReadFull(r, buf[:]); err != nil {
		return Node{}, err
	}
	return readNodeBuf(buf), nil
}

func writeRecord(w io.Writer, r *Record) error {
	buf := writeRecBuf(r)
	_, err := w.Write(buf[:])
	return err
}

func readRecord(r io.Reader) (Record, error) {
	var buf [48]byte
	if _, err := io.ReadFull(r, buf[:]); err != nil {
		return Record{}, err
	}
	return readRecBuf(buf), nil
}

func Save(w io.Writer, t *Tree) error {
	if err := writeHeader(w, t); err != nil {
		return fmt.Errorf("header: %w", err)
	}
	for i := uint32(0); i < t.nCount; i++ {
		if err := writeNode(w, &t.Nodes[i]); err != nil {
			return fmt.Errorf("node %d: %w", i, err)
		}
	}
	for i := uint32(0); i < t.rCount; i++ {
		if err := writeRecord(w, &t.Records[i]); err != nil {
			return fmt.Errorf("record %d: %w", i, err)
		}
	}
	for i := range t.Pending {
		if err := writePending(w, &t.Pending[i]); err != nil {
			return fmt.Errorf("pending %d: %w", i, err)
		}
	}
	return nil
}

func Load(r io.Reader) (*Tree, error) {
	hdr, err := readHeader(r)
	if err != nil {
		return nil, err
	}
	t := &Tree{
		Nodes:     []Node{:0:int(hdr.nCount)},
		Records:   []Record{:0:int(hdr.rCount)},
		Roots:     hdr.roots,
		FreeHead:  hdr.freeHead,
		FreeCount: hdr.freeCount,
		nCount:    hdr.nCount,
		rCount:    hdr.rCount,
		RecKey:    map[uint32]Key{},
	}
	for i := uint32(0); i < hdr.nCount; i++ {
		n, err := readNode(r)
		if err != nil {
			return nil, fmt.Errorf("node %d: %w", i, err)
		}
		t.Nodes = append(t.Nodes, n)
	}
	for i := uint32(0); i < hdr.rCount; i++ {
		rec, err := readRecord(r)
		if err != nil {
			return nil, fmt.Errorf("record %d: %w", i, err)
		}
		t.Records = append(t.Records, rec)
	}
	if hdr.pendCount > 0 {
		t.Pending = []PendingExp{:0:int(hdr.pendCount)}
		for i := uint32(0); i < hdr.pendCount; i++ {
			p, err := readPending(r)
			if err != nil {
				return nil, fmt.Errorf("pending %d: %w", i, err)
			}
			t.Pending = append(t.Pending, p)
		}
	}
	t.rebuildRecKey()
	return t, nil
}

func SaveFile(path string, t *Tree) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer f.Close()
	return Save(f, t)
}

func LoadFile(path string) (*Tree, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return Load(f)
}

func writePending(w io.Writer, p *PendingExp) error {
	var buf [4]byte
	buf[0] = uint8(p.Branch)
	le.PutUint16(buf[1:3], p.Depth)
	buf[3] = p.Factor
	_, err := w.Write(buf[:])
	return err
}

func readPending(r io.Reader) (PendingExp, error) {
	var buf [4]byte
	var p PendingExp
	if _, err := io.ReadFull(r, buf[:]); err != nil {
		return p, err
	}
	p.Branch = Branch(buf[0])
	p.Depth = le.Uint16(buf[1:3])
	p.Factor = buf[3]
	return p, nil
}