package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"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")

// Memory is the memory of one chat: the message log, the summary tree, the chat
// view, and the compaction view. 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.
type Memory struct {
	mutex sync.Mutex
	root  string
	log   *Log
	tree  *Tree
	chat  *View
	cview *View
}

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

// openMemory opens the chat in root. The clock gives the dates of new messages
// and the days of new nodes.
func openMemory(ctx context.Context, root string, clock func() time.Time) (*Memory, error) {
	log, err := openLog(ctx, root, clock)
	if err != nil {
		return nil, err
	}

	m := &Memory{
		root: root,
		log:  log,
	}

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

	return m, nil
}

// load reads the saved view. A crash can stop an append after the log write and
// before the view save, so load continues the view with the messages that it
// does not cover yet.
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())
	}

	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 closes the files of the chat and releases its lock.
func (m *Memory) Close() error {
	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 and adds it to the tree and the view.
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
	}

	return message, nil
}

// addMessage builds the leaf of message when it needs no model call and appends
// it to the view.
func (m *Memory) addMessage(message Message) error {
	_, built := m.tree.Node(0, message.ID)
	leaf, free := freeLeaf(message)
	if free && !built {
		err := m.build(leaf)
		if err != nil {
			return err
		}
	}

	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 then each parent above it that needs no
// model call.
func (m *Memory) build(node Node) error {
	for {
		err := m.tree.Add(node)
		if err != nil {
			return err
		}

		sibling, built := m.tree.Node(node.Level, node.Index^1)
		if !built {
			return nil
		}

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

		parent, free := freeMerge(left, right)
		if !free {
			return nil
		}

		node = parent
	}
}

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