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.
//
// The summarizer builds the nodes in the queue in the background. The start
// reads the log and the tree once to fill the queue again. After that, only a
// new message or a new node adds work.
type Memory struct {
	mutex      sync.Mutex
	root       string
	log        *Log
	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.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.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 {
		if m.tree.Built(0, id) {
			continue
		}

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

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

	for _, node := range m.tree.Nodes() {
		if node.Index%2 != 0 {
			continue
		}

		err := m.settle(node)
		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() {
		err := m.addMessage(m.log.Message(m.chat.end))
		if err != nil {
			return err
		}
	}

	return m.save()
}

// Close stops the builds, waits for them, 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.tree != nil {
		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
	}

	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.
func (m *Memory) addMessage(message Message) error {
	leaf, free := freeLeaf(message)
	if free {
		err := m.build(leaf)
		if err != nil {
			return err
		}
	} else {
		m.queue.AddLeaf(message.ID)
	}

	merged := m.chat.Append(m.tree)
	m.cview.Append(m.tree)
	if merged {
		m.cview.Coarsen(m.chat, m.tree)
	}

	return nil
}

// build adds node to the tree and settles its parent.
func (m *Memory) build(node Node) error {
	err := m.tree.Add(node)
	if err != nil {
		return err
	}

	return m.settle(node)
}

// settle looks at the parent of node 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(node Node) error {
	for {
		sibling, built := m.tree.Node(node.Level, node.Index^1)
		if !built || m.tree.Built(node.Level+1, node.Index/2) {
			return nil
		}

		left, right := node, sibling
		if node.Index%2 != 0 {
			left, right = sibling, node
		}

		parent, free := freeMerge(left, right)
		if !free {
			m.queue.AddMerge(nodeKey{node.Level + 1, node.Index / 2})
			return nil
		}

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

		node = parent
	}
}

// dispatch starts the builds that are ready, up to the limit.
func (m *Memory) dispatch() {
	for m.ctx.Err() == nil {
		key, ready := m.queue.Next()
		if !ready {
			return
		}

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

// task returns the build task of the node key.
func (m *Memory) task(key nodeKey) Task {
	task := Task{
		Level:   key.level,
		Index:   key.index,
		Context: m.context(key.level, key.index),
	}

	if key.level == 0 {
		task.Input = []string{leafSource(m.log.Message(key.index))}
		return task
	}

	for _, index := range []int{2 * key.index, 2*key.index + 1} {
		child, _ := m.tree.Node(key.level-1, index)
		task.Input = append(task.Input, strings.ReplaceAll(child.Text, "\n", " "))
	}

	return task
}

// 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),
		})
	}

	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()
		}
	}
}

// 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(m.tree)
}

// 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(m.tree, end)
}

// writeFileAtomic replaces the file at path with data. A reader or a crash sees
// the old file or the new file, never a part of one.
func writeFileAtomic(path string, data []byte) error {
	dir := filepath.Dir(path)
	file, err := os.CreateTemp(dir, "."+filepath.Base(path)+".*")
	if err != nil {
		return fmt.Errorf("create temporary file: %w", err)
	}

	_, err = file.Write(data)
	if err == nil {
		err = file.Sync()
	}

	err = errors.Join(err, file.Close())
	if err == nil {
		err = os.Rename(file.Name(), path)
	}

	if err != nil {
		return errors.Join(fmt.Errorf("replace file: %w", err), os.Remove(file.Name()))
	}

	return syncDir(dir)
}

// syncDir flushes the entries of dir to the disk.
func syncDir(dir string) (err error) {
	handle, err := os.Open(dir)
	if err != nil {
		return fmt.Errorf("open directory: %w", err)
	}

	defer func() {
		err = errors.Join(err, handle.Close())
	}()

	err = handle.Sync()
	if err != nil {
		return fmt.Errorf("flush directory: %w", err)
	}

	return nil
}