package main

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"strings"
)

// layoutFileName is the file of the view layout in the directory of an agent.
const layoutFileName = "layout.json"

// viewLayout splits the view of a turn request into user messages that only
// grow between batches. The view goes out as the same messages as in the
// request before, while their lines give the same text, and the new lines
// follow as one new message. So the end of the view of each earlier turn stays
// a message end. On the ChatGPT backend, the first request of a turn with this
// layout read the view of the turn before from the prompt cache in each of 10
// turns, and with the view in one message it missed once in 5.
//
// The first message starts with "<chat>\n". The closing tag goes into the
// state message after the view, so that no view message changes when lines
// come. The layout keeps the line count and a hash of the text of each message
// in layout.json, so it outlives a restart.
type viewLayout struct {
	path     string
	Messages []layoutMessage `json:"messages"`
}

// layoutMessage is one view message: its number of lines and the hash of its
// text.
type layoutMessage struct {
	Lines int    `json:"lines"`
	Hash  string `json:"hash"`
}

// loadLayout reads the layout at path. A missing file gives an empty layout.
func loadLayout(path string) (*viewLayout, error) {
	layout := &viewLayout{
		path: path,
	}

	data, err := os.ReadFile(path)
	if errors.Is(err, os.ErrNotExist) {
		return layout, nil
	}

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

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

	return layout, nil
}

// Split returns the view messages for lines and saves the layout. It keeps the
// earlier messages up to the first one whose text differs, and puts the rest
// of the lines into one new message.
func (l *viewLayout) Split(lines []string) ([]string, error) {
	var texts []string
	var kept []layoutMessage
	start := 0
	for _, message := range l.Messages {
		if start+message.Lines > len(lines) {
			break
		}

		text := viewMessage(len(texts) == 0, lines[start:start+message.Lines])
		if textHash(text) != message.Hash {
			break
		}

		texts = append(texts, text)
		kept = append(kept, message)
		start += message.Lines
	}

	if start < len(lines) || len(texts) == 0 {
		text := viewMessage(len(texts) == 0, lines[start:])
		texts = append(texts, text)
		kept = append(kept, layoutMessage{
			Lines: len(lines) - start,
			Hash:  textHash(text),
		})
	}

	l.Messages = kept
	data, err := json.Marshal(l)
	if err != nil {
		return nil, fmt.Errorf("encode view layout: %w", err)
	}

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

	return texts, nil
}

// viewMessage returns the text of a view message with lines. The first one
// opens the chat.
func viewMessage(first bool, lines []string) string {
	text := strings.Join(lines, "")
	if first {
		text = "<chat>\n" + text
	}

	return text
}

// textHash returns a short hash of text.
func textHash(text string) string {
	sum := sha256.Sum256([]byte(text))
	return hex.EncodeToString(sum[:16])
}