package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

// errNoView tells that a chat has messages but no saved view.
var errNoView = errors.New("the chat has messages but no saved view")

// Summarizer builds the text of one node with a model.
type Summarizer interface {
	Summarize(ctx context.Context, task Task) (string, error)
}

// Task is the build of the node (Level, Index). Context is the compaction view
// for the build. Input holds the source "kind: text" of a leaf, or the two
// child lines of a merge with their newlines as spaces.
type Task struct {
	Level   int
	Index   int
	Context string
	Input   []string
}

// Memory is the memory of one chat: the message log, the summary tree, the chat
// view, the compaction view, and the queue of the builds. The compaction view
// is the chat view merged further, at the limits of 16 to 32 KB. It takes the
// same new lines, and it merges again in one batch when it is too large or when
// the chat view merges. So the compactions share its prefix in the prompt
// cache. It saves the views to view.json after each change and loads them at
// start. It never builds the views again from the log, because a new view
// differs from the live one and so loses the prompt cache.
//
// Each write goes to the operating system at once, and a syncer makes the
// writes durable by group commit in the background. Flush waits until the
// writes so far are durable.
//
// The summarizer builds the nodes in the queue in the background. The start
// reads the indexes once to fill the queue again. After that, only a new
// message or a new node adds work. Memory keeps the nodes of the views and the
// tasks of the running builds. It reads all other messages and nodes from their
// day files through the indexes.
type Memory struct {
	mutex      sync.Mutex
	root       string
	log        *Log
	syncer     *syncer
	tree       *Tree
	chat       *View
	cview      *View
	queue      *queue
	summarizer Summarizer
	ctx        context.Context
	cancel     context.CancelFunc
	workers    sync.WaitGroup
	changed    chan struct{}
}

// viewFile is the content of view.json.
type viewFile struct {
	Chat       viewState `json:"chat"`
	Compaction viewState `json:"compaction"`
}

// openMemory opens the chat in root and starts the builds that are ready. The
// clock gives the dates of new messages and the days of new nodes. The builds
// stop when ctx ends or at Close.
func openMemory(ctx context.Context, root string, clock func() time.Time, summarizer Summarizer) (*Memory, error) {
	log, err := openLog(ctx, root, clock)
	if err != nil {
		return nil, err
	}

	m := &Memory{
		root:       root,
		log:        log,
		queue:      newQueue(),
		summarizer: summarizer,
		changed:    make(chan struct{}),
	}

	m.ctx, m.cancel = context.WithCancel(ctx)

	m.syncer, err = openSyncer(root)
	if err != nil {
		return nil, errors.Join(err, m.Close())
	}

	m.tree, err = openTree(ctx, root, log.Len(), clock)
	if err != nil {
		return nil, errors.Join(err, m.Close())
	}

	err = m.load()
	if err != nil {
		return nil, errors.Join(fmt.Errorf("load view: %w", err), m.Close())
	}

	err = m.restoreQueue()
	if err != nil {
		return nil, errors.Join(fmt.Errorf("restore build queue: %w", err), m.Close())
	}

	err = m.catchUp()
	if err != nil {
		return nil, errors.Join(err, m.Close())
	}

	m.mutex.Lock()
	m.syncer.Written()
	m.dispatch()
	m.mutex.Unlock()

	return m, nil
}

// load reads the saved views.
func (m *Memory) load() error {
	data, err := os.ReadFile(m.viewPath())
	if errors.Is(err, os.ErrNotExist) && m.log.Len() == 0 {
		m.chat = newView(chatHigh, chatLow)
		m.cview = newView(compactionHigh, compactionLow)

		return nil
	}

	if errors.Is(err, os.ErrNotExist) {
		return errNoView
	}

	if err != nil {
		return fmt.Errorf("read view: %w", err)
	}

	var file viewFile
	err = json.Unmarshal(data, &file)
	if err != nil {
		return fmt.Errorf("decode view: %w", err)
	}

	m.chat, err = restoreView(file.Chat, chatHigh, chatLow, m.tree)
	if err != nil {
		return fmt.Errorf("chat view: %w", err)
	}

	m.cview, err = restoreView(file.Compaction, compactionHigh, compactionLow, m.tree)
	if err != nil {
		return fmt.Errorf("compaction view: %w", err)
	}

	if m.cview.end != m.chat.end {
		return fmt.Errorf("compaction view covers %d messages, but the chat view %d", m.cview.end, m.chat.end)
	}

	if m.chat.end > m.log.Len() {
		return fmt.Errorf("view covers %d messages, but the log has %d", m.chat.end, m.log.Len())
	}

	return nil
}

// restoreQueue adds the work of the messages in the view to the queue. It
// builds each node that needs no model call and that a crash left out.
func (m *Memory) restoreQueue() error {
	for id := range m.chat.end {
		built, err := m.tree.Built(0, id)
		if err != nil {
			return err
		}

		if built {
			continue
		}

		message, err := m.log.Message(id)
		if err != nil {
			return err
		}

		leaf, free := freeLeaf(message)
		if !free {
			m.queue.AddLeaf(id)
			continue
		}

		err = m.build(leaf)
		if err != nil {
			return err
		}
	}

	keys, err := m.tree.Keys()
	if err != nil {
		return err
	}

	for _, key := range keys {
		if key.index%2 != 0 {
			continue
		}

		err = m.settle(key)
		if err != nil {
			return err
		}
	}

	return nil
}

// catchUp adds the messages that the view does not cover yet. A crash can stop
// an append after the log write and before the view save.
func (m *Memory) catchUp() error {
	if m.chat.end == m.log.Len() {
		return nil
	}

	for m.chat.end < m.log.Len() {
		message, err := m.log.Message(m.chat.end)
		if err != nil {
			return err
		}

		err = m.addMessage(message)
		if err != nil {
			return err
		}
	}

	return m.save()
}

// Close stops the builds, waits for them, makes the writes durable, closes the
// files of the chat, and releases its lock.
func (m *Memory) Close() error {
	m.cancel()
	m.workers.Wait()

	m.mutex.Lock()
	defer m.mutex.Unlock()

	var err error
	if m.syncer != nil {
		err = m.syncer.Close()
	}

	if m.tree != nil {
		err = errors.Join(err, m.tree.Close())
	}

	return errors.Join(err, m.log.Close())
}

// Append logs a new message, adds it to the tree and the views, and starts the
// builds that are ready, also the failed ones.
func (m *Memory) Append(kind, text string) (Message, error) {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	message, err := m.log.Append(kind, text)
	if err != nil {
		return Message{}, err
	}

	defer m.syncer.Written()

	err = m.addMessage(message)
	if err != nil {
		return Message{}, err
	}

	err = m.save()
	if err != nil {
		return Message{}, err
	}

	m.queue.Retry()
	m.dispatch()

	return message, nil
}

// addMessage builds the leaf of message when it needs no model call, or else
// adds the leaf to the queue. Then it appends the leaf to the views. A crash can
// stop an append after the leaf build and before the view save, so a leaf that
// is built already stays as it is, and restoreQueue settles its parent.
func (m *Memory) addMessage(message Message) error {
	built, err := m.tree.Built(0, message.ID)
	if err != nil {
		return err
	}

	leaf, free := freeLeaf(message)
	switch {
	case built:
	case free:
		err = m.build(leaf)
		if err != nil {
			return err
		}
	default:
		m.queue.AddLeaf(message.ID)
	}

	merged, err := m.chat.Append(m.tree)
	if err != nil {
		return fmt.Errorf("append to chat view: %w", err)
	}

	_, err = m.cview.Append(m.tree)
	if err == nil && merged {
		err = m.cview.Coarsen(m.chat, m.tree)
	}

	if err != nil {
		return fmt.Errorf("append to compaction view: %w", err)
	}

	return nil
}

// build adds node to the tree, keeps it in the views when it is a pending line,
// and settles its parent.
func (m *Memory) build(node Node) error {
	err := m.tree.Add(node)
	if err != nil {
		return err
	}

	m.chat.Refresh(node)
	m.cview.Refresh(node)

	return m.settle(nodeKey{node.Level, node.Index})
}

// settle looks at the parent of the built node key when both children are
// built. It builds each parent that needs no model call, up the tree, and adds
// the first other parent to the queue.
func (m *Memory) settle(key nodeKey) error {
	for {
		parent := nodeKey{key.level + 1, key.index / 2}
		ready, err := m.ready(parent)
		if err != nil || !ready {
			return err
		}

		left, right, err := m.children(parent)
		if err != nil {
			return err
		}

		node, free := freeMerge(left, right)
		if !free {
			m.queue.AddMerge(parent)
			return nil
		}

		err = m.tree.Add(node)
		if err != nil {
			return err
		}

		key = parent
	}
}

// ready tells if both children of the node key are built and the node is not.
func (m *Memory) ready(key nodeKey) (bool, error) {
	for _, index := range []int{2 * key.index, 2*key.index + 1} {
		built, err := m.tree.Built(key.level-1, index)
		if err != nil || !built {
			return false, err
		}
	}

	built, err := m.tree.Built(key.level, key.index)

	return !built, err
}

// children reads the two children of the node key.
func (m *Memory) children(key nodeKey) (Node, Node, error) {
	var children [2]Node
	for n := range children {
		child, built, err := m.tree.Node(key.level-1, 2*key.index+n)
		if err != nil {
			return Node{}, Node{}, err
		}

		if !built {
			return Node{}, Node{}, fmt.Errorf("node (%d, %d) is not built", key.level-1, 2*key.index+n)
		}

		children[n] = child
	}

	return children[0], children[1], nil
}

// dispatch starts the builds that are ready, up to the limit. A build whose
// input cannot be read fails and waits for the next message.
func (m *Memory) dispatch() {
	for m.ctx.Err() == nil {
		key, ready := m.queue.Next()
		if !ready {
			return
		}

		task, err := m.task(key)
		if err != nil {
			m.queue.Fail(key)
			continue
		}

		m.workers.Add(1)
		go m.run(key, task)
	}
}

// task returns the build task of the node key, with its input read from the log
// or the tree.
func (m *Memory) task(key nodeKey) (Task, error) {
	task := Task{
		Level:   key.level,
		Index:   key.index,
		Context: m.context(key.level, key.index),
	}

	if key.level == 0 {
		message, err := m.log.Message(key.index)
		if err != nil {
			return Task{}, err
		}

		task.Input = []string{leafSource(message)}

		return task, nil
	}

	left, right, err := m.children(key)
	if err != nil {
		return Task{}, err
	}

	for _, child := range []Node{left, right} {
		task.Input = append(task.Input, strings.ReplaceAll(child.Text, "\n", " "))
	}

	return task, nil
}

// run builds one node with the summarizer. A failed build stays in the queue
// until the next message.
func (m *Memory) run(key nodeKey, task Task) {
	defer m.workers.Done()

	text, err := m.summarizer.Summarize(m.ctx, task)

	m.mutex.Lock()
	defer m.mutex.Unlock()

	if err == nil {
		err = m.build(Node{
			Level: key.level,
			Index: key.index,
			Text:  text,
			Size:  len(text),
		})
		m.syncer.Written()
	}

	if err != nil {
		m.queue.Fail(key)
	} else {
		m.queue.Done(key)
	}

	close(m.changed)
	m.changed = make(chan struct{})
	m.dispatch()
}

// Wait returns when no build runs. The failed builds wait for the next message.
func (m *Memory) Wait(ctx context.Context) error {
	for {
		m.mutex.Lock()
		running := m.queue.Running()
		changed := m.changed
		m.mutex.Unlock()

		if running == 0 {
			return nil
		}

		select {
		case <-changed:
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

// Flush returns when every write before the call is durable.
func (m *Memory) Flush(ctx context.Context) error {
	return m.syncer.Flush(ctx)
}

// save writes the views to view.json.
func (m *Memory) save() error {
	data, err := json.Marshal(viewFile{
		Chat:       m.chat.State(),
		Compaction: m.cview.State(),
	})
	if err != nil {
		return fmt.Errorf("encode view: %w", err)
	}

	err = writeFileAtomic(m.viewPath(), append(data, '\n'))
	if err != nil {
		return fmt.Errorf("save view: %w", err)
	}

	return nil
}

func (m *Memory) viewPath() string {
	return filepath.Join(m.root, "view.json")
}

// View returns the chat view as the model reads it.
func (m *Memory) View() string {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	return m.chat.Render()
}

// ViewLines returns the lines of the chat view, each with its newline, without
// the leaf lines of the messages in skip.
func (m *Memory) ViewLines(skip map[int]bool) []string {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	return m.chat.LinesWithout(skip)
}

// Len returns the number of messages in the log.
func (m *Memory) Len() int {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	return m.log.Len()
}

// Message reads the message with the given ID.
func (m *Memory) Message(id int) (Message, error) {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	return m.log.Message(id)
}

// Node reads the node (l, i) if it is built.
func (m *Memory) Node(level, index int) (Node, bool, error) {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	return m.tree.Node(level, index)
}

// Context returns the compaction view for the build of node (l, i). It ends
// before the message of a leaf, or after the last message of a merge, and it
// stops at the first line that is not built.
func (m *Memory) Context(level, index int) string {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	return m.context(level, index)
}

func (m *Memory) context(level, index int) string {
	end := index
	if level > 0 {
		end = nodeEnd(level, index)
	}

	return m.cview.Context(end)
}