package main

import (
	"cmp"
	"context"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"math/bits"
	"os"
	"path/filepath"
	"slices"
	"time"
	"unicode/utf8"
)

// nodeBytes is the size limit of a node text in UTF-8 bytes. A model can go a
// few bytes over it, because the view measures the real sizes.
const nodeBytes = 512

// treeRecordBytes is the size of a record of index/tree.idx: the position of
// the node line, the node index, and the node level.
const treeRecordBytes = 32

// openTree opens the tree of the chat in root and indexes its new lines. A chat
// of total messages cannot have a node that covers a later message.
func openTree(ctx context.Context, root string, total int, clock func() time.Time) (*Tree, error) {
	for _, dir := range []string{"tree", "index"} {
		err := os.MkdirAll(filepath.Join(root, dir), 0o700)
		if err != nil {
			return nil, fmt.Errorf("create %s directory: %w", dir, err)
		}
	}

	t := &Tree{
		clock: clock,
		journal: Journal{
			dir: filepath.Join(root, "tree"),
		},
		dir:    filepath.Join(root, "index"),
		levels: make(map[int]*Records),
	}

	err := t.open(ctx, total)
	if err != nil {
		return nil, errors.Join(fmt.Errorf("index tree: %w", err), t.Close())
	}

	return t, nil
}

// leafSource returns the text that the leaf of message holds or summarizes.
func leafSource(message Message) string {
	return message.Kind + ": " + message.Text
}

// freeLeaf returns the leaf of message when its source fits in a node. Such a
// leaf holds the source word for word and needs no model call.
func freeLeaf(message Message) (Node, bool) {
	text := leafSource(message)
	if len(text) > nodeBytes {
		return Node{}, false
	}

	return Node{
		Level: 0,
		Index: message.ID,
		Text:  text,
		Size:  len(text),
	}, true
}

// freeMerge returns the parent of two sibling nodes when both texts, joined by
// a newline, fit in a node. Such a parent needs no model call.
func freeMerge(left, right Node) (Node, bool) {
	if left.Size+1+right.Size > nodeBytes {
		return Node{}, false
	}

	text := left.Text + "\n" + right.Text

	return Node{
		Level: left.Level + 1,
		Index: left.Index / 2,
		Text:  text,
		Size:  len(text),
	}, true
}

// nodeEnd returns the ID after the last message of node (l, i).
func nodeEnd(level, index int) int {
	return (index + 1) << level
}

// Tree holds the built nodes of one chat. It writes one JSON line for each node
// to tree/YYYY-MM-DD.jsonl, named for the local day of the build. Each node is
// built once and never again. A node of a higher level is built only after both
// of its children. After a failed write, the tree refuses all writes until it
// is opened again. The caller holds the lock of the chat and serializes all
// calls.
//
// The tree keeps no node in memory. The index index/tree.idx lists the node
// lines in file order, and index/tree-L.idx holds the position of node (L, i)
// in record i, with a zero record for a node that is not built. At open the
// tree indexes the lines after the last indexed one, so a missing index is
// built again from the day files.
type Tree struct {
	clock   func() time.Time
	journal Journal
	dir     string
	order   *Records
	levels  map[int]*Records
	failure error
}

// Close closes the files of the tree.
func (t *Tree) Close() error {
	err := t.journal.Close()
	if t.order != nil {
		err = errors.Join(err, t.order.Close())
		t.order = nil
	}

	for level, index := range t.levels {
		err = errors.Join(err, index.Close())
		delete(t.levels, level)
	}

	return err
}

// Add writes a new node to the tree.
func (t *Tree) Add(node Node) error {
	if t.failure != nil {
		return fmt.Errorf("tree needs a reopen after a failed write: %w", t.failure)
	}

	err := t.check(node)
	if err != nil {
		return err
	}

	line, err := json.Marshal(node)
	if err != nil {
		return fmt.Errorf("encode node: %w", err)
	}

	place, err := t.journal.Write(t.clock().Format(time.DateOnly), append(line, '\n'))
	if err == nil {
		err = t.addIndex(NodeKey{node.Level, node.Index}, place)
	}

	if err != nil {
		t.failure = err
		return fmt.Errorf("write node (%d, %d): %w", node.Level, node.Index, err)
	}

	return nil
}

// Built tells if the node (l, i) is built.
func (t *Tree) Built(level, index int) (bool, error) {
	_, built, err := t.place(level, index)
	return built, err
}

// Node reads the node (l, i) if it is built.
func (t *Tree) Node(level, index int) (Node, bool, error) {
	place, built, err := t.place(level, index)
	if err != nil || !built {
		return Node{}, false, err
	}

	line, err := t.journal.Read(place)
	if err != nil {
		return Node{}, false, fmt.Errorf("read node (%d, %d): %w", level, index, err)
	}

	var node Node
	err = json.Unmarshal(line, &node)
	if err != nil {
		return Node{}, false, fmt.Errorf("decode node (%d, %d): %w", level, index, err)
	}

	if node.Level != level || node.Index != index {
		return Node{}, false, fmt.Errorf("index of node (%d, %d) points to node (%d, %d)", level, index, node.Level, node.Index)
	}

	return node, true, nil
}

// Keys returns the keys of the built nodes, ordered by level and then by index.
func (t *Tree) Keys() ([]NodeKey, error) {
	keys := make([]NodeKey, 0, t.order.count)
	for n := range t.order.count {
		record, _, err := t.order.Read(n)
		if err != nil {
			return nil, err
		}

		key, _ := treeDecodeRecord(record)
		keys = append(keys, key)
	}

	slices.SortFunc(keys, func(a, b NodeKey) int {
		return cmp.Or(cmp.Compare(a.level, b.level), cmp.Compare(a.index, b.index))
	})

	return keys, nil
}

func (t *Tree) open(ctx context.Context, total int) error {
	var err error
	t.order, err = openRecords(filepath.Join(t.dir, "tree.idx"), treeRecordBytes)
	if err != nil {
		return err
	}

	start := Position{}
	if t.order.count == 0 {
		err = t.removeLevels()
		if err != nil {
			return err
		}
	} else {
		record, _, err := t.order.Read(t.order.count - 1)
		if err != nil {
			return err
		}

		// A crash can stop an add between the two index writes, so the
		// last node gets its level record again.
		node, last := treeDecodeRecord(record)
		err = t.writeLevel(node, last)
		if err != nil {
			return err
		}

		start = Position{
			day:    last.day,
			offset: last.offset + uint64(last.length),
		}
	}

	return t.journal.Scan(ctx, start, func(line []byte, place Position) error {
		var node Node
		err := json.Unmarshal(line, &node)
		if err != nil {
			return err
		}

		err = t.check(node)
		if err != nil {
			return err
		}

		if nodeEnd(node.Level, node.Index) > total {
			return fmt.Errorf("node (%d, %d) covers messages after the log", node.Level, node.Index)
		}

		return t.addIndex(NodeKey{node.Level, node.Index}, place)
	})
}

// check makes sure that node is valid and can join the tree now.
func (t *Tree) check(node Node) error {
	if node.Level < 0 || node.Level >= bits.UintSize-2 || node.Index < 0 {
		return fmt.Errorf("node (%d, %d) is out of range", node.Level, node.Index)
	}

	if node.Text == "" || !utf8.ValidString(node.Text) || node.Size != len(node.Text) {
		return fmt.Errorf("node (%d, %d) has no valid text of its size", node.Level, node.Index)
	}

	built, err := t.Built(node.Level, node.Index)
	if err != nil {
		return err
	}

	if built {
		return fmt.Errorf("node (%d, %d) is already built", node.Level, node.Index)
	}

	if node.Level == 0 {
		return nil
	}

	for _, index := range []int{2 * node.Index, 2*node.Index + 1} {
		built, err = t.Built(node.Level-1, index)
		if err != nil {
			return err
		}

		if !built {
			return fmt.Errorf("node (%d, %d) has a child that is not built", node.Level, node.Index)
		}
	}

	return nil
}

// addIndex adds the node line at place to both indexes.
func (t *Tree) addIndex(node NodeKey, place Position) error {
	record := make([]byte, treeRecordBytes)
	encodePosition(record, place)
	binary.LittleEndian.PutUint64(record[positionBytes:], uint64(node.index))
	binary.LittleEndian.PutUint32(record[positionBytes+8:], uint32(node.level))

	err := t.order.Append(record)
	if err != nil {
		return err
	}

	return t.writeLevel(node, place)
}

// place returns the position of the node (l, i) if it is built.
func (t *Tree) place(level, index int) (Position, bool, error) {
	if level < 0 || level >= bits.UintSize-2 || index < 0 {
		return Position{}, false, nil
	}

	records, err := t.level(level)
	if err != nil {
		return Position{}, false, err
	}

	record, found, err := records.Read(index)
	if err != nil || !found {
		return Position{}, false, err
	}

	place := decodePosition(record)

	return place, place.length > 0, nil
}

// removeLevels removes the level indexes, which are not valid without the order
// index.
func (t *Tree) removeLevels() error {
	paths, err := filepath.Glob(filepath.Join(t.dir, "tree-*.idx"))
	if err != nil {
		return fmt.Errorf("list level indexes: %w", err)
	}

	for _, path := range paths {
		err = os.Remove(path)
		if err != nil {
			return fmt.Errorf("remove level index: %w", err)
		}
	}

	return nil
}

// writeLevel writes the position of node to its level index.
func (t *Tree) writeLevel(node NodeKey, place Position) error {
	level, err := t.level(node.level)
	if err != nil {
		return err
	}

	record := make([]byte, positionBytes)
	encodePosition(record, place)

	return level.Write(node.index, record)
}

// level returns the index of level, and opens it at first use.
func (t *Tree) level(level int) (*Records, error) {
	index, open := t.levels[level]
	if open {
		return index, nil
	}

	index, err := openRecords(filepath.Join(t.dir, fmt.Sprintf("tree-%d.idx", level)), positionBytes)
	if err != nil {
		return nil, err
	}

	t.levels[level] = index

	return index, nil
}

// Count returns the number of built nodes.
func (t *Tree) Count() int {
	return t.order.count
}

// Node is one line of the summary tree. Node (l, i) covers the 2^l messages
// from i*2^l on. A node of level 0 holds one message. A node of a higher level
// merges its two children, (l-1, 2i) and (l-1, 2i+1). Size is the UTF-8 byte
// length of Text.
type Node struct {
	Level int    `json:"l"`
	Index int    `json:"i"`
	Text  string `json:"text"`
	Size  int    `json:"size"`
}

// NodeKey names the node (l, i).
type NodeKey struct {
	level int
	index int
}

// treeDecodeRecord reads the key and the position of a record of
// index/tree.idx.
func treeDecodeRecord(record []byte) (NodeKey, Position) {
	key := NodeKey{
		level: int(binary.LittleEndian.Uint32(record[positionBytes+8:])),
		index: int(binary.LittleEndian.Uint64(record[positionBytes:])),
	}

	return key, decodePosition(record)
}