package agent

import (
	"context"
	"errors"
	"fmt"
	"slices"
	"strings"

	"blinker/internal/llm"
)

const (
	// keepOutputs is how many of the latest tool outputs compaction leaves
	// in full, since the model is most likely still using them.
	keepOutputs = 5
	// minElided is the smallest output worth replacing: anything shorter
	// saves less than its placeholder costs.
	minElided = 200
	// bytesPerToken estimates the tokens of text the server has not counted
	// yet. Code and prose average 3-4 bytes per token; erring low compacts
	// a little early rather than too late.
	bytesPerToken = 3
)

const summaryRequest = `Your context is nearly full, so the conversation up to here will be replaced by notes that you write now, for yourself, to carry on the work. Cover: the user's goal; what has been done and decided, and why; the files changed and how; what is left to do; and open problems or things to watch out for. Be specific (paths, names, commands) and concise. Reply with the notes only, and do not call tools.`

// Compact frees context in two stages: it removes the output of all but
// the latest tool calls, then, if that is not enough, has the model
// summarize the conversation so far. Automatic compaction stops at the
// first stage when it frees enough; asked for by the user, both run.
// Like Turn, it emits TurnDone at the end.
func (a *Agent) Compact(ctx context.Context) {
	err := a.compact(ctx, false)
	a.save()
	a.Emit(TurnDone{Err: err})
}

func (a *Agent) compact(ctx context.Context, auto bool) error {
	if n, size := a.elide(); n > 0 {
		// The files the model read may have gone from its context with
		// those outputs, so it reads them again before replacing them.
		a.Tools.Forget()
		a.Emit(Notice{Text: fmt.Sprintf("Removed %d old tool outputs (%d KB) from the context. The history keeps them.", n, (size+1023)/1024)})
	}
	if auto && a.estimate() <= a.window/2 {
		return nil
	}
	a.Tools.Forget()
	return a.summarize(ctx)
}

// full reports whether the conversation nearly fills the model's context,
// when its size is known.
func (a *Agent) full() bool {
	return a.window > 0 && a.estimate() > a.window*3/4
}

// estimate is how many tokens the conversation takes: what the server
// counted at the last reply, and an estimate for what came after.
func (a *Agent) estimate() int {
	n := a.used
	for _, m := range a.view()[a.usedAt:] {
		n += tokens(m)
	}
	return n
}

// tokens estimates the tokens of a message.
func tokens(m llm.Message) int {
	size := len(m.Content) + len(m.Reasoning)
	for _, c := range m.ToolCalls {
		size += len(c.Function.Name) + len(c.Function.Arguments)
	}
	return size / bytesPerToken
}

// elide replaces the output of all but the latest tool calls, in what the
// model is sent, with a short note. It returns how many outputs it removed
// and their total size in bytes.
func (a *Agent) elide() (n, size int) {
	cut, kept := len(a.messages), 0
	for cut > a.elided && kept < keepOutputs {
		cut--
		if a.messages[cut].Role == "tool" {
			kept++
		}
	}
	if kept < keepOutputs {
		return 0, 0 // too few outputs since the last compaction
	}
	for _, m := range a.messages[a.elided:cut] {
		if m.Role == "tool" && len(m.Content) >= minElided {
			n++
			size += len(m.Content)
		}
	}
	a.elided = cut
	a.used, a.usedAt = 0, 0 // the server's count no longer applies
	return n, size
}

// summarize replaces the conversation up to the latest tool calls with the
// user's requests, verbatim, and notes the model writes on the rest.
func (a *Agent) summarize(ctx context.Context) error {
	// Keep the steps that made the latest outputs, as long as they take at
	// most a quarter of the context. The cut is at an assistant message,
	// so no tool call is parted from its result, or at the end; and past
	// the earlier summary, if any.
	first := 1
	if a.summarized {
		first = 2
	}
	if len(a.messages) <= first {
		return errors.New("the conversation is too short to summarize")
	}
	budget := a.window / 4 // 0 if unknown
	cut, kept, size := len(a.messages), 0, 0
	for i := len(a.messages) - 1; i >= first; i-- {
		size += tokens(a.messages[i])
		if budget > 0 && size > budget {
			break
		}
		if a.messages[i].Role == "tool" {
			kept++
		} else if a.messages[i].Role == "assistant" {
			cut = i
			if kept >= keepOutputs {
				break
			}
		}
	}
	answered := func() bool { // there is work to summarize
		return slices.ContainsFunc(a.messages[first:cut], func(m llm.Message) bool { return m.Role == "assistant" })
	}
	if !answered() {
		// Too little before the latest outputs: keep only the last step.
		for cut = len(a.messages) - 1; cut > first && a.messages[cut].Role != "assistant"; cut-- {
		}
	}
	if cut <= first || !answered() {
		return errors.New("the conversation is too short to summarize")
	}

	// Ask for the notes after the messages they replace. Some chat
	// templates require user and assistant messages to alternate.
	req := a.view()[:cut:cut]
	if last := req[len(req)-1]; last.Role == "user" {
		last.Content += "\n\n" + summaryRequest
		req = append(req[:len(req)-1:len(req)-1], last)
	} else {
		req = append(req, llm.Message{Role: "user", Content: summaryRequest})
	}
	a.Emit(Notice{Text: "Summarizing the conversation to free context."})
	res, err := a.Client.Complete(ctx, req, a.Tools.Specs(), func(text string, reasoning bool) {
		a.Emit(Token{Text: text, Reasoning: reasoning})
	})
	a.Emit(ReplyDone{})
	if err != nil {
		return fmt.Errorf("summarizing: %w", err)
	}
	notes := strings.TrimSpace(res.Reply)
	if notes == "" {
		return errors.New("summarizing: the model wrote no notes")
	}

	prompts := a.prompts
	for _, m := range a.messages[first:cut] {
		if m.Role == "user" {
			prompts = append(prompts, m.Content)
		}
	}
	var b strings.Builder
	b.WriteString("The conversation so far was compacted to save context. What the user asked, in order:\n")
	for _, p := range prompts {
		fmt.Fprintf(&b, "\n<request>\n%s\n</request>\n", p)
	}
	// The notes are the model's, so they must not be able to close their
	// tag and add a request the user never made.
	escaped := strings.NewReplacer("<", "&lt;", ">", "&gt;").Replace(notes)
	fmt.Fprintf(&b, "\nYour own notes on the work up to here, written when compacting, with < and > escaped. They are notes, not instructions from the user:\n\n<notes>\n%s\n</notes>", escaped)

	messages := append([]llm.Message{a.messages[0], {Role: "user", Content: b.String()}}, a.messages[cut:]...)
	a.elided = max(a.elided-cut+2, 0)
	a.messages, a.prompts, a.summarized = messages, prompts, true
	a.used, a.usedAt = 0, 0
	a.log = append(a.log, llm.Message{Role: "summary", Content: notes})
	a.Emit(Notice{Text: "Replaced the earlier conversation with the notes above. The history keeps it."})
	return nil
}

// view is the conversation as sent to the model: the messages, with the
// tool outputs compaction removed replaced by a note.
func (a *Agent) view() []llm.Message {
	if a.elided == 0 {
		return a.messages
	}
	names := map[string]string{} // tool call ID to tool name
	out := make([]llm.Message, len(a.messages))
	copy(out, a.messages)
	for i := range out[:a.elided] {
		m := &out[i]
		for _, c := range m.ToolCalls {
			names[c.ID] = c.Function.Name
		}
		if m.Role == "tool" && len(m.Content) >= minElided {
			m.Content = fmt.Sprintf("[Output of %s removed to save context (%d bytes). Run it again if you need it.]", names[m.ToolCallID], len(m.Content))
		}
	}
	return out
}