package ingest

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

	"git.smesh.lol/iskradb/lattice"
)

var le = binary.LittleEndian

// SaveFlat writes the DB to a flat binary file.
// Format (C=8):
//   [4B]     nodeCount
//   [4B]     recCount
//   [C×4B]   roots (one per axis branch)
//   [4B]     poolLen
//   [nodeCount × nodeSize]  nodes
//   [recCount × 48B]        records
//   [poolLen B]             string pool
//   [recCount × 20B]        RecKey entries: recIdx(4) + key(16)
func SaveFlat(db *DB, path string) error {
	tmp := path | ".tmp"
	f, err := os.Create(tmp)
	if err != nil {
		return err
	}

	nCount := uint32(db.Tree.NodeCount())
	rCount := uint32(db.Tree.RecordCount())
	poolLen := uint32(len(db.StringPool))

	// Header: nCount(4) + rCount(4) + C roots(C*4) + poolLen(4)
	hdrSize := 4 + 4 + lattice.C*4 + 4
	hdr := []byte{:hdrSize:hdrSize}
	le.PutUint32(hdr[0:], nCount)
	le.PutUint32(hdr[4:], rCount)
	for i := 0; i < lattice.C; i++ {
		le.PutUint32(hdr[8+i*4:], db.Tree.RootIdx(lattice.Branch(i)))
	}
	le.PutUint32(hdr[8+lattice.C*4:], poolLen)
	if _, err := f.Write(hdr); err != nil {
		f.Close()
		os.Remove(tmp)
		return err
	}

	// Nodes: 64 bytes each.
	for i := uint32(0); i < nCount; i++ {
		if err := writeNode(f, db.Tree, i); err != nil {
			f.Close()
			os.Remove(tmp)
			return fmt.Errorf("node %d: %w", i, err)
		}
	}

	// Records: 48 bytes each.
	for i := uint32(0); i < rCount; i++ {
		if err := writeRecord(f, db.Tree, i); err != nil {
			f.Close()
			os.Remove(tmp)
			return fmt.Errorf("rec %d: %w", i, err)
		}
	}

	// String pool.
	if _, err := f.Write(db.StringPool); err != nil {
		f.Close()
		os.Remove(tmp)
		return err
	}

	// RecKey map: recIdx(4) + key(16) per entry (128-bit SipHash key).
	var rkBuf [20]byte
	for recIdx, key := range db.Tree.RecKey {
		le.PutUint32(rkBuf[0:], recIdx)
		le.PutUint64(rkBuf[4:], key[0])
		le.PutUint64(rkBuf[12:], key[1])
		if _, err := f.Write(rkBuf[:]); err != nil {
			f.Close()
			os.Remove(tmp)
			return err
		}
	}

	if err := f.Close(); err != nil {
		os.Remove(tmp)
		return err
	}
	return os.Rename(tmp, path)
}

// LoadFlatFromReader reads a DB from an io.Reader containing data written by SaveFlat.
func LoadFlatFromReader(r io.Reader) (*DB, error) {
	hdrSize := 4 + 4 + lattice.C*4 + 4
	hdr := []byte{:hdrSize:hdrSize}
	if _, err := io.ReadFull(r, hdr); err != nil {
		return nil, fmt.Errorf("header: %w", err)
	}
	nCount := le.Uint32(hdr[0:])
	rCount := le.Uint32(hdr[4:])
	var roots [lattice.C]uint32
	for i := 0; i < lattice.C; i++ {
		roots[i] = le.Uint32(hdr[8+i*4:])
	}
	poolLen := le.Uint32(hdr[8+lattice.C*4:])

	tree := lattice.AllocTree(nCount, rCount, roots)

	for i := uint32(0); i < nCount; i++ {
		if err := readNode(r, tree, i); err != nil {
			return nil, fmt.Errorf("node %d: %w", i, err)
		}
	}

	for i := uint32(0); i < rCount; i++ {
		if err := readRecord(r, tree, i); err != nil {
			return nil, fmt.Errorf("rec %d: %w", i, err)
		}
	}

	pool := []byte{:int(poolLen)}
	if poolLen > 0 {
		if _, err := io.ReadFull(r, pool); err != nil {
			return nil, fmt.Errorf("pool: %w", err)
		}
	}

	var rkBuf [20]byte
	for {
		n, err := io.ReadFull(r, rkBuf[:])
		if n == 0 || err == io.EOF || err == io.ErrUnexpectedEOF {
			break
		}
		if err != nil {
			return nil, fmt.Errorf("reckey: %w", err)
		}
		recIdx := le.Uint32(rkBuf[0:])
		key := lattice.Key{le.Uint64(rkBuf[4:]), le.Uint64(rkBuf[12:])}
		tree.RecKey[recIdx] = key
	}

	return &DB{Tree: tree, StringPool: pool}, nil
}

// LoadFlat reads a DB from a flat binary file written by SaveFlat.
func LoadFlat(path string) (*DB, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return LoadFlatFromReader(f)
}

// Node: Keys[C](C*16B) RecPtrs[C](C*4B) Children[C+1]((C+1)*4B) Mult(1) Flags(1) pad(2)
const flatNodeSize = lattice.C*16 + lattice.C*4 + (lattice.C+1)*4 + 4

func writeNode(w io.Writer, tree *lattice.Tree, idx uint32) error {
	node := tree.GetNode(idx)
	buf := []byte{:flatNodeSize:flatNodeSize}
	for i := 0; i < lattice.C; i++ {
		le.PutUint64(buf[i*16:], node.Keys[i][0])
		le.PutUint64(buf[i*16+8:], node.Keys[i][1])
	}
	base := lattice.C * 16
	for i := 0; i < lattice.C; i++ {
		le.PutUint32(buf[base+i*4:], node.RecPtrs[i])
	}
	base += lattice.C * 4
	for i := 0; i < lattice.C+1; i++ {
		le.PutUint32(buf[base+i*4:], node.Children[i])
	}
	base += (lattice.C + 1) * 4
	buf[base] = node.Mult
	buf[base+1] = node.Flags
	_, err := w.Write(buf)
	return err
}

func readNode(r io.Reader, tree *lattice.Tree, idx uint32) error {
	buf := []byte{:flatNodeSize:flatNodeSize}
	if _, err := io.ReadFull(r, buf); err != nil {
		return err
	}
	node := tree.GetNode(idx)
	for i := 0; i < lattice.C; i++ {
		node.Keys[i][0] = le.Uint64(buf[i*16:])
		node.Keys[i][1] = le.Uint64(buf[i*16+8:])
	}
	base := lattice.C * 16
	for i := 0; i < lattice.C; i++ {
		node.RecPtrs[i] = le.Uint32(buf[base+i*4:])
	}
	base += lattice.C * 4
	for i := 0; i < lattice.C+1; i++ {
		node.Children[i] = le.Uint32(buf[base+i*4:])
	}
	base += (lattice.C + 1) * 4
	node.Mult = buf[base]
	node.Flags = buf[base+1]
	return nil
}

// Record: DataFile(4) DataOff(4) DataLen(4) Link[2](8) Branch(1) pad(3) Inline[24] = 48B
func writeRecord(w io.Writer, tree *lattice.Tree, idx uint32) error {
	rec := tree.GetRecord(idx)
	var buf [48]byte
	le.PutUint32(buf[0:], rec.DataFile)
	le.PutUint32(buf[4:], rec.DataOff)
	le.PutUint32(buf[8:], rec.DataLen)
	le.PutUint32(buf[12:], rec.Link[0])
	le.PutUint32(buf[16:], rec.Link[1])
	buf[20] = rec.Branch
	copy(buf[24:], rec.Inline[:])
	_, err := w.Write(buf[:])
	return err
}

func readRecord(r io.Reader, tree *lattice.Tree, idx uint32) error {
	var buf [48]byte
	if _, err := io.ReadFull(r, buf[:]); err != nil {
		return err
	}
	rec := tree.GetRecord(idx)
	rec.DataFile = le.Uint32(buf[0:])
	rec.DataOff = le.Uint32(buf[4:])
	rec.DataLen = le.Uint32(buf[8:])
	rec.Link[0] = le.Uint32(buf[12:])
	rec.Link[1] = le.Uint32(buf[16:])
	rec.Branch = buf[20]
	copy(rec.Inline[:], buf[24:])
	return nil
}