// Package agent holds the turn loop and the conversation state it runs on.
//
// It must not import the ui package: the eval runner in Phase 3 drives this
// loop headless, and anything terminal-shaped here would make that impossible.
package agent

import (
	"encoding/json"
	"strings"
	"sync"
	"time"

	"github.com/johannSo/codey/internal/llm"
	"github.com/johannSo/codey/internal/tools"
)

// ReadRecord lives in tools so the edit tool can query the tracker without
// importing this package.
type ReadRecord = tools.ReadRecord

// State is the conversation history plus the per-session file-read tracking.
type State struct {
	mu       sync.Mutex
	messages []llm.Message
	reads    map[string]ReadRecord
	usage    llm.Usage
	turns    int
}

func NewState() *State {
	return &State{reads: map[string]ReadRecord{}}
}

// MarkRead implements tools.ReadTracker.
func (s *State) MarkRead(path string, mod time.Time) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.reads[path] = ReadRecord{At: time.Now(), Mod: mod}
}

// LastRead returns the read record for path, if the file has been read.
func (s *State) LastRead(path string) (ReadRecord, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	r, ok := s.reads[path]
	return r, ok
}

// Append adds a message to the history.
func (s *State) Append(m llm.Message) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.messages = append(s.messages, m)
}

// Replace swaps the whole history, which is how compaction lands.
func (s *State) Replace(msgs []llm.Message) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.messages = append([]llm.Message(nil), msgs...)
}

// SupersedeReads blanks the results of every read_file call for path except the
// most recent one. A file read twenty turns ago and edited since is not merely
// wasted context — it is wrong, and the model will reason from it.
func (s *State) SupersedeReads(path string) int {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Collect the tool_use ids of every read of this path, in order.
	var ids []string
	for _, m := range s.messages {
		if m.Role != llm.RoleAssistant {
			continue
		}
		for _, b := range m.Content {
			if b.Type == llm.BlockToolUse && b.Name == "read_file" && readPath(b.Input) == path {
				ids = append(ids, b.ID)
			}
		}
	}
	if len(ids) <= 1 {
		return 0
	}
	stale := map[string]bool{}
	for _, id := range ids[:len(ids)-1] {
		stale[id] = true
	}

	var n int
	for i, m := range s.messages {
		if m.Role != llm.RoleUser {
			continue
		}
		for j, b := range m.Content {
			if b.Type != llm.BlockToolResult || !stale[b.ToolUseID] {
				continue
			}
			if strings.HasPrefix(b.Content, stalePrefix) {
				continue
			}
			s.messages[i].Content[j].Content = stalePrefix + path +
				" has been modified since this read. The contents shown here are out of date; read it again if you need them.]"
			// An image of the old file is as stale as its text, and far
			// bigger; the notice above is all that goes back.
			s.messages[i].Content[j].Blocks = nil
			n++
		}
	}
	return n
}

const stalePrefix = "[superseded: "

// readPath pulls the path argument out of a read_file call.
func readPath(input json.RawMessage) string {
	var args struct {
		Path string `json:"path"`
	}
	if json.Unmarshal(input, &args) != nil {
		return ""
	}
	return args.Path
}

// History returns a copy of the conversation so far.
func (s *State) History() []llm.Message {
	s.mu.Lock()
	defer s.mu.Unlock()
	return append([]llm.Message(nil), s.messages...)
}

// Len is the number of messages in the history.
func (s *State) Len() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.messages)
}

// AddUsage accumulates the token counts of one response.
func (s *State) AddUsage(u llm.Usage) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.usage.InputTokens += u.InputTokens
	s.usage.OutputTokens += u.OutputTokens
	s.usage.CacheReadInputTokens += u.CacheReadInputTokens
	s.usage.CacheCreationInputTokens += u.CacheCreationInputTokens
	s.turns++
}

// RestoreUsage sets the running totals outright, which is what resuming a
// saved session needs: AddUsage would count the whole restored history as one
// more turn.
func (s *State) RestoreUsage(u llm.Usage, turns int) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.usage, s.turns = u, turns
}

// Usage is the running total for the session, and the number of model turns it
// took to get there.
func (s *State) Usage() (llm.Usage, int) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.usage, s.turns
}