ollie-9p/session/session.go

package session

import (
	"context"
	"fmt"
	"hash/crc32"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"ollie/pkg/agent"
	"ollie/pkg/backend"
	"ollie/pkg/config"
	olog "ollie/pkg/log"
	"ollie/pkg/paths"
	"ollie/pkg/tools"
	"ollie/pkg/tools/execute"
	"olliesrv/fs"
)

// Session holds all state for one agent session.
type Session struct {
	mu         sync.RWMutex
	id         string
	uname      string // immutable user principal (numeric UID), set at creation
	Core       agent.Core
	Ctx        context.Context
	cancel     context.CancelFunc
	log        []byte
	logVers    uint32
	ChatOffset int
	plan       []byte
	prevPrompt []byte // last submitted prompt; overwritten on each new submission
	peers      map[string]bool // peer session IDs (bidirectional links)

	// Cached model list (expensive API call; refreshed every 24h).
	modelsMu      sync.Mutex
	modelsCache   string
	modelsCacheAt time.Time
}

func NewSession(id string, core agent.Core, ctx context.Context, cancel context.CancelFunc) *Session {
	sess := &Session{id: id, Core: core, Ctx: ctx, cancel: cancel}
	sess.startEventLog()
	return sess
}

const modelsCacheTTL = 24 * time.Hour

// CachedListModels returns the model list, using a 24h cache to avoid
// repeated expensive API calls.
// NOTE: Must NOT be called while holding sess.mu (caller content() holds RLock).
// Uses its own modelsMu to avoid deadlock.
func (sess *Session) CachedListModels() string {
	sess.modelsMu.Lock()
	if sess.modelsCache != "" && time.Since(sess.modelsCacheAt) < modelsCacheTTL {
		result := sess.modelsCache
		sess.modelsMu.Unlock()
		return result
	}
	sess.modelsMu.Unlock()

	// Cache miss — fetch and store.
	result := sess.Core.ListModels()
	sess.modelsMu.Lock()
	sess.modelsCache = result
	sess.modelsCacheAt = time.Now()
	sess.modelsMu.Unlock()
	return result
}

// InvalidateModelsCache clears the cached model list.
func (sess *Session) InvalidateModelsCache() {
	sess.modelsMu.Lock()
	sess.modelsCache = ""
	sess.modelsCacheAt = time.Time{}
	sess.modelsMu.Unlock()
}

// HasPeer reports whether peerID is a registered peer of this session.
func (sess *Session) HasPeer(peerID string) bool {
	sess.mu.RLock()
	defer sess.mu.RUnlock()
	return sess.peers[peerID]
}

func (sess *Session) RunnableID() string { return sess.id }
func (sess *Session) Uname() string     { return sess.uname }

func (sess *Session) Cancel() {
	sess.cancel()
}

func (sess *Session) Interrupt() {
	sess.Core.Interrupt(agent.ErrInterrupted)
}

// AppendLog appends data to the session's log and bumps the version.
func (sess *Session) AppendLog(data []byte) {
	if len(data) == 0 {
		return
	}
	sess.mu.Lock()
	sess.log = append(sess.log, data...)
	sess.logVers++
	sess.mu.Unlock()
}

// EnsureTrailingNewline appends a newline if the log doesn't already end with one.
func (sess *Session) EnsureTrailingNewline() {
	sess.mu.Lock()
	if len(sess.log) > 0 && sess.log[len(sess.log)-1] != '\n' {
		sess.log = append(sess.log, '\n')
	}
	sess.logVers++
	sess.mu.Unlock()
}

// startEventLog subscribes to the agent's "event" bus topic and writes
// all events to the session chat log.
func (sess *Session) startEventLog() {
	streamingRole := "" // tracks current streaming role ("assistant" or "reasoning")
	sess.Core.Bus().Subscribe("event", func(ev agent.Event) {
		switch ev.Role {
		case "assistant", "reasoning":
			if streamingRole != ev.Role {
				if streamingRole != "" {
					sess.AppendLog([]byte("\n"))
				}
				sess.AppendLog([]byte("[" + ev.Role + "]\n"))
				if ev.Role == "assistant" {
					sess.mu.Lock()
					sess.ChatOffset = len(sess.log)
					sess.mu.Unlock()
				}
				streamingRole = ev.Role
			}
			sess.AppendLog(FormatEvent(ev))
		default:
			if streamingRole != "" {
				sess.AppendLog([]byte("\n"))
				streamingRole = ""
			}
			sess.AppendLog(FormatEvent(ev))
		}
	})
}

// LogInfo returns the current log length and version atomically.
func (sess *Session) LogInfo() (length int, vers uint32) {
	sess.mu.RLock()
	defer sess.mu.RUnlock()
	return len(sess.log), sess.logVers
}

// Mu returns the session's mutex for external synchronization (e.g. D-Bus adapter).
func (sess *Session) Mu() *sync.RWMutex { return &sess.mu }

// Log returns the raw chat log bytes. Caller must hold Mu().RLock().
func (sess *Session) Log() []byte { return sess.log }

// Plan returns the plan bytes. Caller must hold Mu().RLock().
func (sess *Session) Plan() []byte { return sess.plan }

// SetPlan sets the plan. Caller must hold Mu().Lock().
func (sess *Session) SetPlan(p []byte) { sess.plan = p }

// PrevPrompt returns the last submitted prompt. Caller must hold Mu().RLock().
func (sess *Session) PrevPrompt() []byte { return sess.prevPrompt }



var sessionStoreOrder = []string{"new", "idx", "ls", "kill", "sh", "b", "bfg", "bbg", "cleanup"}

// FileMode returns the mode for a fixed session file,
// or 0 and false if the name is not a fixed file.
func FileMode(name string) (os.FileMode, bool) {
	m, ok := fs.Perms[fs.PathSessions].Files[name]
	return m, ok
}

// ManagerConfig holds the dependencies for a Manager.
type ManagerConfig struct {
	AgentsDir   string
	SessionsDir string
	Log         *olog.Logger
	Sink        *olog.Sink
	ReadFile    func(string) ([]byte, error)
	MkdirAll    func(string, os.FileMode) error
	// NewCore, if non-nil, replaces the default backend.New + agent.NewAgentCore
	// path. It receives the session ID, agent name, and cwd, and returns a Core.
	NewCore func(sessionID, agentName, cwd string) (agent.Core, error)
	// Strict rejects inline code steps; only tool steps are allowed.
	Strict bool
	// Yolo skips the landrun sandbox.
	Yolo bool
	// NoMount disables the per-session FUSE mount (e.g. when listening on TCP).
	NoMount bool
	// InvalidateModels clears the model cache, forcing a refresh.
	InvalidateModels func()
	// OnSessionCreated is called after a new session is added to the manager.
	// Receives the session ID and the Session pointer.
	OnSessionCreated func(id string, sess *Session)
	// OnSessionKilled is called after a session is removed from the manager.
	OnSessionKilled func(id string)
	// OnSessionRenamed is called after a session is renamed.
	OnSessionRenamed func(oldID, newID string)
}

// Manager manages session lifecycle and exposes sessions as a Tree.
type Manager struct {
	tree     *fs.Tree
	cfg       ManagerConfig
	mu        sync.RWMutex
	sessions  map[string]*Session
	nextUID   atomic.Uint32 // incrementing principal counter
}

// Tree returns the Tree view of the session namespace.
func (s *Manager) Tree() *fs.Tree { return s.tree }
// nextUname generates the next uname atomically.
func (s *Manager) nextUname() string {
	return fmt.Sprintf("%d", s.nextUID.Add(1))
}

func NewManager(cfg ManagerConfig) *Manager {
	if cfg.ReadFile == nil {
		cfg.ReadFile = os.ReadFile
	}
	if cfg.MkdirAll == nil {
		cfg.MkdirAll = os.MkdirAll
	}
	ss := &Manager{
		cfg:      cfg,
		sessions: make(map[string]*Session),
	}
	ss.nextUID.Store(9999)
	ss.restoreAllSessions()
	ss.tree = fs.NewTree(nil, 0,
		fs.WithStat(func(_ []string, name string) (os.FileInfo, error) { return ss.stat(name) }),
		fs.WithOpener(func(_ []string, name string) (fs.File, error) { return ss.openEntry(name) }),
		fs.WithLister(func(_ []string) ([]os.DirEntry, error) { return ss.list() }),
		fs.WithReaddir(func(_ []string, name string) ([]os.DirEntry, error) { return ss.Readdir(name) }),
		fs.WithDeleter(func(_ []string, name string) error { return ss.del(name) }),
		fs.WithCreator(func(_ []string, name string, _ os.FileMode) error { return ss.create(name) }),
		fs.WithRenamer(func(_ []string, old, new string) error { return ss.renameSession(old, new) }),
	)
	return ss
}

// AddSession inserts a pre-built session into the fs.
func (s *Manager) AddSession(sess *Session) {
	s.mu.Lock()
	s.sessions[sess.RunnableID()] = sess
	s.mu.Unlock()
}

// List returns all root-level entries (fixed files + session directories).
func (s *Manager) List() ([]os.DirEntry, error) {
	return s.list()
}

func (s *Manager) list() ([]os.DirEntry, error) {
	entries := make([]os.DirEntry, 0, len(sessionStoreOrder))
	for _, name := range sessionStoreOrder {
		entries = append(entries, fs.FileEntry(name, fs.Perms[fs.PathSessions].Files[name]))
	}
	s.mu.RLock()
	for id := range s.sessions {
		entries = append(entries, fs.DirEntry(id, fs.Perms[fs.PathSessionDir].DirMode))
	}
	s.mu.RUnlock()
	return entries, nil
}

// Readdir lists entries in a subdirectory (e.g. "{id}", "{id}/peer").
func (s *Manager) Readdir(name string) ([]os.DirEntry, error) {
	parts := strings.SplitN(name, "/", 3)
	sessID := parts[0]
	sess := s.Session(sessID)
	if sess == nil {
		return nil, fmt.Errorf("session not found: %s", sessID)
	}
	// {id} — list session files + peer/
	if len(parts) == 1 {
		sfs, err := s.openStore(sess)
		if err != nil {
			return nil, err
		}
		entries, err := sfs.List()
		if err != nil {
			return nil, err
		}
		// Append peer/ and proc/ directories
		entries = append(entries, fs.DirEntry("peer", 0755))
		entries = append(entries, fs.DirEntry("proc", 0755))
		return entries, nil
	}
	// {id}/peer — list peer session IDs
	if len(parts) == 2 && parts[1] == "peer" {
		peers, _ := s.PeerList(sessID)
		entries := make([]os.DirEntry, len(peers))
		for i, p := range peers {
			entries[i] = fs.FileEntry(p, 0666)
		}
		return entries, nil
	}
	// {id}/proc — list detached process PIDs
	if len(parts) == 2 && parts[1] == "proc" {
		procs := sess.Core.ListDetached()
		entries := make([]os.DirEntry, len(procs))
		for i, p := range procs {
			entries[i] = fs.FileEntry(fmt.Sprintf("%d", p.PID), 0666)
		}
		return entries, nil
	}
	return nil, fmt.Errorf("%s: not a directory", name)
}

// --- Peer Management ---

// PeerAdd creates a bidirectional peer link between two sessions.
func (s *Manager) PeerAdd(sessID, peerID string) error {
	s.mu.RLock()
	sess, ok := s.sessions[sessID]
	peer, peerOk := s.sessions[peerID]
	s.mu.RUnlock()
	if !ok {
		return fmt.Errorf("session not found: %s", sessID)
	}
	if !peerOk {
		return fmt.Errorf("peer session not found: %s", peerID)
	}
	if sessID == peerID {
		return fmt.Errorf("cannot peer a session with itself")
	}

	sess.mu.Lock()
	if sess.peers == nil {
		sess.peers = make(map[string]bool)
	}
	sess.peers[peerID] = true
	sess.mu.Unlock()

	peer.mu.Lock()
	if peer.peers == nil {
		peer.peers = make(map[string]bool)
	}
	peer.peers[sessID] = true
	peer.mu.Unlock()
	return nil
}

// PeerRemove removes a bidirectional peer link between two sessions.
func (s *Manager) PeerRemove(sessID, peerID string) error {
	s.mu.RLock()
	sess, ok := s.sessions[sessID]
	peer, peerOk := s.sessions[peerID]
	s.mu.RUnlock()
	if !ok {
		return fmt.Errorf("session not found: %s", sessID)
	}

	sess.mu.Lock()
	delete(sess.peers, peerID)
	sess.mu.Unlock()

	if peerOk {
		peer.mu.Lock()
		delete(peer.peers, sessID)
		peer.mu.Unlock()
	}
	return nil
}

// PeerList returns the peer session IDs for a session.
func (s *Manager) PeerList(sessID string) ([]string, error) {
	s.mu.RLock()
	sess, ok := s.sessions[sessID]
	s.mu.RUnlock()
	if !ok {
		return nil, fmt.Errorf("session not found: %s", sessID)
	}
	sess.mu.RLock()
	result := make([]string, 0, len(sess.peers))
	for id := range sess.peers {
		result = append(result, id)
	}
	sess.mu.RUnlock()
	return result, nil
}

// PeerSubmit sends a prompt to a peer session.
func (s *Manager) PeerSubmit(sessID, peerID, prompt string) error {
	s.mu.RLock()
	sess, ok := s.sessions[sessID]
	peer, peerOk := s.sessions[peerID]
	s.mu.RUnlock()
	if !ok {
		return fmt.Errorf("session not found: %s", sessID)
	}
	if !peerOk {
		return fmt.Errorf("peer session not found: %s", peerID)
	}
	sess.mu.RLock()
	isPeer := sess.peers[peerID]
	sess.mu.RUnlock()
	if !isPeer {
		return fmt.Errorf("%s is not a peer of %s", peerID, sessID)
	}
	go peer.Core.Submit(peer.Ctx, prompt)
	return nil
}

func (s *Manager) stat(name string) (os.FileInfo, error) {
	// Top-level fixed files (new, idx, sh, etc.)
	if mode, ok := fs.Perms[fs.PathSessions].Files[name]; ok {
		return &fs.SyntheticFileInfo{Name_: name, Mode_: mode}, nil
	}
	parts := strings.SplitN(name, "/", 3)
	sessID := parts[0]
	s.mu.RLock()
	sess, ok := s.sessions[sessID]
	s.mu.RUnlock()
	if !ok {
		return nil, fmt.Errorf("%s: not found", name)
	}
	// Session directory: {id}
	if len(parts) == 1 {
		return &fs.SyntheticFileInfo{Name_: sessID, Mode_: fs.Perms[fs.PathSessionDir].DirMode, IsDir_: true}, nil
	}
	// Peer directory: {id}/peer
	if parts[1] == "peer" {
		if len(parts) == 2 {
			return &fs.SyntheticFileInfo{Name_: "peer", Mode_: 0755, IsDir_: true}, nil
		}
		// Peer file: {id}/peer/{peer-id}
		peerID := parts[2]
		sess.mu.RLock()
		isPeer := sess.peers[peerID]
		sess.mu.RUnlock()
		if !isPeer {
			return nil, fmt.Errorf("%s: not found", name)
		}
		return &fs.SyntheticFileInfo{Name_: peerID, Mode_: 0666}, nil
	}
	// Proc directory: {id}/proc
	if parts[1] == "proc" {
		if len(parts) == 2 {
			return &fs.SyntheticFileInfo{Name_: "proc", Mode_: 0755, IsDir_: true}, nil
		}
		// Proc file: {id}/proc/{pid} — read-only output
		return &fs.SyntheticFileInfo{Name_: parts[2], Mode_: 0444}, nil
	}
	// Session file: {id}/{file}
	sfs, err := s.openStore(sess)
	if err != nil {
		return nil, err
	}
	return sfs.Stat(parts[1])
}

func (s *Manager) openEntry(name string) (fs.File, error) {
	notBlocking := func(context.Context, string) ([]byte, string, error) {
		return nil, "", fmt.Errorf("blocking read not supported")
	}

	// Peer file: {id}/peer/{peer-id} — write submits prompt to peer
	parts := strings.SplitN(name, "/", 3)

	// Proc file: {id}/proc/{pid} — read returns ring buffer output
	if len(parts) == 3 && parts[1] == "proc" {
		sessID := parts[0]
		sess := s.Session(sessID)
		if sess == nil {
			return nil, fmt.Errorf("%s: not found", name)
		}
		pidStr := parts[2]
		pid, err := strconv.Atoi(pidStr)
		if err != nil {
			return nil, fmt.Errorf("invalid pid: %s", pidStr)
		}
		return &fs.FileConfig{
			StatFn: func() (os.FileInfo, error) {
				return &fs.SyntheticFileInfo{Name_: pidStr, Mode_: 0444}, nil
			},
			ReadFn: func() ([]byte, error) {
				output, err := sess.Core.GetDetachedOutput(pid)
				if err != nil {
					return nil, err
				}
				return []byte(output), nil
			},
			WriteFn:        func([]byte) error { return fmt.Errorf("read-only") },
			BlockingReadFn: notBlocking,
		}, nil
	}

	if len(parts) == 3 && parts[1] == "peer" {
		sessID := parts[0]
		peerID := parts[2]
		return &fs.FileConfig{
			StatFn: func() (os.FileInfo, error) {
				return &fs.SyntheticFileInfo{Name_: peerID, Mode_: 0666}, nil
			},
			ReadFn: func() ([]byte, error) {
				return nil, nil // empty read
			},
			WriteFn: func(data []byte) error {
				prompt := strings.TrimSpace(string(data))
				if prompt == "" {
					return nil
				}
				return s.PeerSubmit(sessID, peerID, prompt)
			},
			BlockingReadFn: notBlocking,
		}, nil
	}

	// Top-level fixed files.
	switch name {
	case "new":
		return &fs.FileConfig{
			StatFn: func() (os.FileInfo, error) { return &fs.SyntheticFileInfo{Name_: "new", Mode_: fs.Perms[fs.PathSessions].Files["new"]}, nil },
			ReadFn: func() ([]byte, error) {
				return []byte("name=\ncwd=\nbackend=\nmodel=\nagent=\nmaxTokens=\nmaxCompletionTokens=\ntemperature=\ntopP=\ntopK=\nminP=\ntopA=\nfrequencyPenalty=\npresencePenalty=\nrepetitionPenalty=\nreasoning=\nreasoningEffort=\nincludeReasoning=\nresponseFormat=\nstop=\nverbosity=\n"), nil
			},
			WriteFn: func(data []byte) error {
				_, err := s.CreateSession(strings.Fields(strings.TrimSpace(string(data))))
				return err
			},
			BlockingReadFn: notBlocking,
		}, nil
	case "idx":
		return &fs.FileConfig{
			StatFn:         func() (os.FileInfo, error) { return &fs.SyntheticFileInfo{Name_: "idx", Mode_: fs.Perms[fs.PathSessions].Files["idx"]}, nil },
			ReadFn:         func() ([]byte, error) { return s.index(), nil },
			WriteFn:        func([]byte) error { return fmt.Errorf("idx: read-only") },
			BlockingReadFn: notBlocking,
		}, nil
	}
	if _, ok := fs.Perms[fs.PathSessions].Files[name]; ok {
		return &fs.FileConfig{
			StatFn: func() (os.FileInfo, error) { return &fs.SyntheticFileInfo{Name_: name, Mode_: fs.Perms[fs.PathSessions].Files[name]}, nil },
			ReadFn: func() ([]byte, error) {
				return s.cfg.ReadFile(paths.CfgDir() + "/scripts/s/" + name)
			},
			WriteFn:        func([]byte) error { return fmt.Errorf("%s: not writable", name) },
			BlockingReadFn: notBlocking,
		}, nil
	}

	// Hierarchical paths: {id}/{file} or {id}/t/{tool}
	parts = strings.SplitN(name, "/", 3)
	sessID := parts[0]
	sess := s.Session(sessID)
	if sess == nil {
		return nil, fmt.Errorf("%s: not found", name)
	}
	if len(parts) == 1 {
		return nil, fmt.Errorf("%s: is a directory", name)
	}
	// Session file: {id}/{file}
	sfs, err := s.openStore(sess)
	if err != nil {
		return nil, err
	}
	return sfs.Open(parts[1])
}

func (s *Manager) create(name string) error {
	// Peer creation: {id}/peer/{peer-id}
	parts := strings.SplitN(name, "/", 3)
	if len(parts) == 3 && parts[1] == "peer" {
		return s.PeerAdd(parts[0], parts[2])
	}
	return fmt.Errorf("create not supported: %s", name)
}

func (s *Manager) del(name string) error {
	parts := strings.SplitN(name, "/", 3)
	sessID := parts[0]
	// Delete session directory itself.
	if len(parts) == 1 {
		s.mu.RLock()
		_, ok := s.sessions[sessID]
		s.mu.RUnlock()
		if ok {
			s.KillSession(sessID)
			return nil
		}
		return fmt.Errorf("session not found: %s", sessID)
	}
	// Peer removal: {id}/peer/{peer-id}
	if len(parts) == 3 && parts[1] == "peer" {
		return s.PeerRemove(sessID, parts[2])
	}
	// Proc dismiss: {id}/proc/{pid}
	if len(parts) == 3 && parts[1] == "proc" {
		sess := s.Session(sessID)
		if sess == nil {
			return fmt.Errorf("session not found: %s", sessID)
		}
		pid, err := strconv.Atoi(parts[2])
		if err != nil {
			return fmt.Errorf("invalid pid: %s", parts[2])
		}
		if !sess.Core.DismissDetached(pid) {
			return fmt.Errorf("process %d not found or still running", pid)
		}
		return nil
	}
	// Session files are synthetic; allow rm -r to continue.
	return nil
}

// Session returns the session for the given ID, or nil.
func (s *Manager) Session(id string) *Session {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.sessions[id]
}

// SessionByUname returns the session with the given uname (principal), or nil.
func (s *Manager) SessionByUname(uname string) *Session {
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, sess := range s.sessions {
		if sess.uname == uname {
			return sess
		}
	}
	return nil
}

// OpenStore returns a Tree for the given session ID.
func (s *Manager) OpenStore(id string) (*fs.Tree, error) {
	if sess := s.Session(id); sess != nil {
		return s.openStore(sess)
	}
	return nil, fmt.Errorf("session not found: %s", id)
}

func (s *Manager) openStore(sess *Session) (*fs.Tree, error) {
	return NewSessionTree(
		sess,
		s.cfg.Log,
		func() { s.KillSession(sess.id) },
		func(newID string) error { return s.renameSession(sess.id, newID) },
		nil, // no transcript saving
		s.cfg.InvalidateModels,
	), nil
}

// InterruptAll interrupts every active session.
func (s *Manager) InterruptAll() {
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, sess := range s.sessions {
		sess.Core.Interrupt(agent.ErrInterrupted)
	}
}

// Shutdown kills all active sessions.
// --- Session Persistence ---

func (s *Manager) activeSessionsDir() string {
	return filepath.Join(paths.DataDir(), "sessions", "active-9p")
}

func (s *Manager) persistSession(id string) {
	s.mu.RLock()
	sess, ok := s.sessions[id]
	s.mu.RUnlock()
	if !ok {
		return
	}
	dir := s.activeSessionsDir()
	os.MkdirAll(dir, 0700)
	path := filepath.Join(dir, id+".json")
	if err := sess.Core.SaveSession(path); err != nil {
		s.cfg.Log.Error("persist session %s: %v", id, err)
	}
}

func (s *Manager) removePersistedSession(id string) {
	path := filepath.Join(s.activeSessionsDir(), id+".json")
	os.Remove(path)
}

func (s *Manager) saveAllSessions() {
	s.mu.RLock()
	ids := make([]string, 0, len(s.sessions))
	for id := range s.sessions {
		ids = append(ids, id)
	}
	s.mu.RUnlock()
	for _, id := range ids {
		s.persistSession(id)
	}
}

func (s *Manager) restoreAllSessions() {
	dir := s.activeSessionsDir()
	entries, err := os.ReadDir(dir)
	if err != nil {
		return
	}
	for _, e := range entries {
		if !strings.HasSuffix(e.Name(), ".json") {
			continue
		}
		path := filepath.Join(dir, e.Name())
		ps, err := agent.LoadPersistedSession(path)
		if err != nil {
			s.cfg.Log.Error("restore session %s: %v", e.Name(), err)
			continue
		}
		if err := s.restoreSession(ps); err != nil {
			s.cfg.Log.Error("restore session %s: %v", ps.ID, err)
		}
	}
	// Clean up persisted files after successful restore
	for _, e := range entries {
		os.Remove(filepath.Join(dir, e.Name()))
	}
}

func (s *Manager) restoreSession(ps *agent.PersistedSession) error {
	cwd := ps.CWD
	if cwd == "" {
		cwd, _ = os.Getwd()
	}
	agentName := ps.Agent
	if agentName == "" {
		agentName = "default"
	}
	sessID := ps.ID

	cfg := LoadAgentConfig(s.cfg.AgentsDir, agentName, nil)

	backendName := ps.Backend
	if backendName == "" && cfg != nil && cfg.Backend != "" {
		backendName = cfg.Backend
	}
	be, err := backend.NewWithName(backendName)
	if err != nil {
		return fmt.Errorf("backend: %w", err)
	}
	modelName := ps.Model
	if modelName == "" && cfg != nil && cfg.Model != "" {
		modelName = cfg.Model
	}
	if modelName == "" {
		modelName = os.Getenv("OLLIE_MODEL")
	}
	if modelName != "" {
		be.SetModel(modelName)
	}

	uname := s.nextUname()
	var execOpts []execute.Option
	if !s.cfg.NoMount {
		execOpts = append(execOpts, WithMount9P())
	}
	if s.cfg.Strict {
		execOpts = append(execOpts, execute.WithStrict())
	}
	if s.cfg.Yolo {
		execOpts = append(execOpts, execute.WithYolo())
	}
	newDisp := tools.NewDispatcherFunc(map[string]func() tools.Server{
		"execute": execute.Decl(cwd, execOpts...),
	})
	rt := agent.BuildRuntime(cfg, newDisp(), cwd, []string{"OLLIE_SESSION_ID=" + sessID, "OLLIE_UNAME=" + uname, "OLLIE=" + MountPathFor(uname)})

	restoredSession := agent.RestoreSession(ps.Messages)
	if ps.TaskState != nil {
		restoredSession.TaskState = ps.TaskState
	}

	var sessPtr *Session
	core := agent.NewAgentCore(agent.AgentCoreConfig{
		Backend:       be,
		AgentName:     agentName,
		AgentsDir:     s.cfg.AgentsDir,
		SessionsDir:   s.cfg.SessionsDir,
		SessionID:     sessID,
		Uname:         uname,
		CWD:           cwd,
		Session:       restoredSession,
		Runtime:       rt,
		NewDispatcher: newDisp,
		Log:           s.cfg.Sink.NewLogger("core"),
		ReadPlanStep: func() string {
			if sessPtr == nil {
				return ""
			}
			sessPtr.mu.RLock()
			data := make([]byte, len(sessPtr.plan))
			copy(data, sessPtr.plan)
			sessPtr.mu.RUnlock()
			return agent.NextUncheckedStep(data)
		},
	})

	ctx, cancel := context.WithCancel(context.Background())
	sess := NewSession(sessID, core, ctx, cancel)
	sessPtr = sess
	sess.uname = uname

	// Replay tail of persisted messages into the chat log so the GUI
	// and `chat` file show recent history on restore.
	replayMessagesToLog(sess, ps.Messages)

	s.mu.Lock()
	s.sessions[sessID] = sess
	s.mu.Unlock()

	s.cfg.Log.Info("restored session %s (backend=%s model=%s agent=%s)", sessID, backendName, modelName, agentName)
	if s.cfg.OnSessionCreated != nil {
		s.cfg.OnSessionCreated(sessID, sess)
	}
	return nil
}

func (s *Manager) Shutdown() {
	s.saveAllSessions()
	s.mu.Lock()
	ids := make([]string, 0, len(s.sessions))
	for id := range s.sessions {
		ids = append(ids, id)
	}
	s.mu.Unlock()
	for _, id := range ids {
		s.mu.Lock()
		sess := s.sessions[id]
		delete(s.sessions, id)
		s.mu.Unlock()
		if sess != nil {
			sess.Cancel()
			sess.Core.Close()
			s.cfg.Log.Info("shutdown session %s", id)
		}
	}
}

func (s *Manager) KillSession(id string) {
	s.mu.Lock()
	sess := s.sessions[id]
	delete(s.sessions, id)
	s.mu.Unlock()
	if sess != nil {
		sess.Cancel()
		sess.Core.Close()
		s.removePersistedSession(id)
		s.cfg.Log.Info("killed session %s", id)
		if s.cfg.OnSessionKilled != nil {
			s.cfg.OnSessionKilled(id)
		}
	}
}

func (s *Manager) index() []byte {
	var sb strings.Builder
	s.mu.RLock()
	ids := make([]string, 0, len(s.sessions))
	for id := range s.sessions {
		ids = append(ids, id)
	}
	sort.Strings(ids)
	for _, id := range ids {
		sess := s.sessions[id]
		sess.mu.RLock()
		state := sess.Core.State()
		cwd := sess.Core.CWD()
		be := sess.Core.BackendName()
		model := sess.Core.ModelName()
		agent := sess.Core.AgentName()
		sess.mu.RUnlock()
		fmt.Fprintf(&sb, "%s\t%s\t%s\t%s\t%s\t%s\n", id, state, cwd, be, model, agent)
	}
	s.mu.RUnlock()
	return []byte(sb.String())
}

// CreateSession creates a new agent session from key=value args.
// Returns the session ID on success.
func (s *Manager) CreateSession(args []string) (string, error) {
	name := ""
	backendOverride := ""
	modelOverride := ""
	agentName := "default"
	cwd := ""
	for _, arg := range args {
		k, v, ok := strings.Cut(arg, "=")
		if !ok {
			return "", fmt.Errorf("invalid option %q (expected key=value)", arg)
		}
		if v == "" {
			continue
		}
		switch k {
		case "name":
			name = v
		case "backend":
			backendOverride = v
		case "model":
			modelOverride = v
		case "agent":
			agentName = v
		case "cwd":
			cwd = v
		default:
			return "", fmt.Errorf("unknown option %q (valid: name, backend, model, agent, cwd)", k)
		}
	}
	if agentName == "default" {
		if v := os.Getenv("OLLIE_DEFAULT_AGENT"); v != "" {
			agentName = v
		}
	}

	cwd = paths.ExpandHome(os.ExpandEnv(cwd))
	if cwd == "" {
		return "", fmt.Errorf("cwd is required (e.g. new cwd=/path/to/project)")
	}

	sessID := name
	if sessID == "" {
		sessID = agent.NewSessionID()
	}

	s.mu.RLock()
	_, exists := s.sessions[sessID]
	s.mu.RUnlock()
	if exists {
		return "", fmt.Errorf("session already exists: %s", sessID)
	}

	var core agent.Core
	var sessPtr *Session
	uname := s.nextUname()
	if s.cfg.NewCore != nil {
		var err error
		core, err = s.cfg.NewCore(sessID, agentName, cwd)
		if err != nil {
			return "", err
		}
	} else {
		cfg := LoadAgentConfig(s.cfg.AgentsDir, agentName, nil)
		if cfg != nil {
			if backendOverride == "" && cfg.Backend != "" {
				backendOverride = cfg.Backend
			}
			if modelOverride == "" && cfg.Model != "" {
				modelOverride = cfg.Model
			}
		}

		be, err := backend.NewWithName(backendOverride)
		if err != nil {
			return "", fmt.Errorf("backend: %w", err)
		}

		if modelOverride == "" {
			modelOverride = os.Getenv("OLLIE_MODEL")
		}
		if modelOverride != "" {
			be.SetModel(modelOverride)
		}

		if err := s.cfg.MkdirAll(s.cfg.SessionsDir, fs.PermMkdirPrivate); err != nil {
			return "", fmt.Errorf("sessions dir: %w", err)
		}

		var execOpts []execute.Option
			if !s.cfg.NoMount {
				execOpts = append(execOpts, WithMount9P())
			}
		if s.cfg.Strict {
			execOpts = append(execOpts, execute.WithStrict())
		}
		if s.cfg.Yolo {
			execOpts = append(execOpts, execute.WithYolo())
		}
		if cfg != nil && len(cfg.AllowExecutors) > 0 {
			execOpts = append(execOpts, execute.WithAllowExecutors(cfg.AllowExecutors))
		}
		if cfg != nil && len(cfg.AllowTools) > 0 {
			execOpts = append(execOpts, execute.WithAllowTools(cfg.AllowTools))
		}
		newDisp := tools.NewDispatcherFunc(map[string]func() tools.Server{
			"execute": execute.Decl(cwd, execOpts...),
		})

		rt := agent.BuildRuntime(cfg, newDisp(), cwd, []string{"OLLIE_SESSION_ID=" + sessID, "OLLIE_UNAME=" + uname, "OLLIE=" + MountPathFor(uname)})

		// sessPtr is set after NewSession; the ReadPlanStep closure captures it.
		core = agent.NewAgentCore(agent.AgentCoreConfig{
			Backend:       be,
			AgentName:     agentName,
			AgentsDir:     s.cfg.AgentsDir,
			SessionsDir:   s.cfg.SessionsDir,
			SessionID:     sessID,
			Uname:         uname,
			CWD:           cwd,
			Runtime:       rt,
			NewDispatcher: newDisp,
			Log:           s.cfg.Sink.NewLogger("core"),
			ReadPlanStep: func() string {
				if sessPtr == nil {
					return ""
				}
				sessPtr.mu.RLock()
				data := make([]byte, len(sessPtr.plan))
				copy(data, sessPtr.plan)
				sessPtr.mu.RUnlock()
				return agent.NextUncheckedStep(data)
			},
		})
	}

	ctx, cancel := context.WithCancel(context.Background())
	sess := NewSession(sessID, core, ctx, cancel)
	sessPtr = sess

	s.mu.Lock()
	sess.uname = uname
	s.sessions[sessID] = sess
	s.mu.Unlock()

	// HACK: Load the agent now that the session mount is active.
	core.Submit(ctx, "/agent "+agentName)

	s.cfg.Log.Info("new session %s (backend=%s model=%s agent=%s)",
		sessID, core.BackendName(), core.ModelName(), core.AgentName())
	if s.cfg.OnSessionCreated != nil {
		s.cfg.OnSessionCreated(sessID, sess)
	}
	return sessID, nil
}

// RenameSession renames a session (exported for D-Bus adapter).
func (s *Manager) RenameSession(old, new string) error {
	return s.renameSession(old, new)
}

func (s *Manager) renameSession(old, new string) error {
	// Session rename: {oldID} -> {newID}
	oldID := old
	newID := new
	s.mu.Lock()

	sess, ok := s.sessions[oldID]
	if !ok {
		s.mu.Unlock()
		return fmt.Errorf("session not found: %s", oldID)
	}
	if _, exists := s.sessions[newID]; exists {
		s.mu.Unlock()
		return fmt.Errorf("session already exists: %s", newID)
	}
	if sess.Core.IsRunning() {
		s.mu.Unlock()
		return fmt.Errorf("cannot rename while agent is running")
	}

	if err := sess.Core.SetSessionID(newID); err != nil {
		s.mu.Unlock()
		return err
	}

	sess.id = newID
	s.sessions[newID] = sess
	delete(s.sessions, oldID)
	s.mu.Unlock()

	sess.AppendLog([]byte(fmt.Sprintf(":: session renamed: %s -> %s\n", oldID, newID)))
	s.cfg.Log.Info("renamed session %s -> %s", oldID, newID)
	if s.cfg.OnSessionRenamed != nil {
		s.cfg.OnSessionRenamed(oldID, newID)
	}
	return nil
}

// LoadAgentConfig resolves and loads the config for a named agent.
// Returns nil if the config file does not exist; BuildRuntime handles nil configs.
func LoadAgentConfig(agentsDir, name string, open func(string) (*os.File, error)) *config.Config {
	if open == nil {
		open = os.Open
	}
	path := agent.AgentConfigPath(agentsDir, name)
	f, err := open(path)
	if err != nil {
		return nil
	}
	defer f.Close()
	cfg, _ := config.Load(f)
	return cfg
}

// FormatEvent converts an agent Event to bytes for appending to a chat log.
// Streaming roles (assistant, reasoning) return only the content chunk;
// the caller (startEventLog) is responsible for writing the [role] header.
func FormatEvent(ev agent.Event) []byte {
	switch ev.Role {
	case "user":
		return []byte("[user]\n" + ev.Content + "\n")
	case "assistant":
		return []byte(ev.Content)
	case "reasoning":
		return []byte(ev.Content)
	case "call":
		args := squashWhitespace(ev.Content)
		return []byte("[call:" + ev.Name + "]\n" + args + "\n")
	case "tool":
		return []byte("[tool:" + ev.Name + "]\n" + strings.TrimRight(ev.Content, "\n") + "\n")
	case "error":
		return []byte("[error]\n" + ev.Content + "\n")
	case "info":
		return []byte("[info]\n" + ev.Content)
	case "retry":
		return []byte("[retry]\n" + ev.Content + "s\n")
	case "stalled":
		return []byte("[stalled]\n")
	case "maxsteps":
		return []byte("[maxsteps]\n" + ev.Content + "\n")
	case "exec":
		return []byte("[exec]\n" + ev.Content + "\n")
	default:
		return nil
	}
}

func squashWhitespace(s string) string {
	return strings.Join(strings.Fields(s), " ")
}

// replayMessagesToLog renders the tail of a persisted message list into the
// session's chat log so that restored sessions show recent history.
func replayMessagesToLog(sess *Session, messages []backend.Message) {
	const maxReplay = 20
	start := len(messages) - maxReplay
	if start < 0 {
		start = 0
	}
	// Skip system messages at the start of the replay window.
	for start < len(messages) && messages[start].Role == "system" {
		start++
	}
	for _, m := range messages[start:] {
		switch m.Role {
		case "system":
			// skip
		case "user":
			sess.AppendLog([]byte("[user]\n" + m.Content + "\n"))
		case "assistant":
			sess.AppendLog([]byte("[assistant]\n"))
			if m.Content != "" {
				sess.AppendLog([]byte(m.Content + "\n"))
			}
			for _, tc := range m.ToolCalls {
				args := squashWhitespace(string(tc.Arguments))
				sess.AppendLog([]byte("[call:" + tc.Name + "]\n" + args + "\n"))
			}
		case "tool":
			sess.AppendLog([]byte("[tool]\n" + strings.TrimRight(m.Content, "\n") + "\n"))
		}
	}
}

// --- Code Completion ---

// Complete performs a stateless one-shot code completion via backend.Generate.
func (s *Manager) Complete(cwd, filePath, prefix, suffix, extraContext string) (string, error) {
	if cwd == "" {
		cwd, _ = os.Getwd()
	}
	cwd = paths.ExpandHome(os.ExpandEnv(cwd))

	model := os.Getenv("OLLIE_COMPLETE_MODEL")
	backendName := os.Getenv("OLLIE_COMPLETE_BACKEND")
	if model == "" || backendName == "" {
		return "", fmt.Errorf("OLLIE_COMPLETE_MODEL and OLLIE_COMPLETE_BACKEND must be set")
	}

	const prefixMax = 12000
	const suffixMax = 1000
	if len(prefix) > prefixMax {
		prefix = prefix[len(prefix)-prefixMax:]
	}
	if len(suffix) > suffixMax {
		suffix = suffix[:suffixMax]
	}

	fileHint := ""
	if filePath != "" {
		fileHint = " in " + filePath
	}
	contextBlock := ""
	if extraContext != "" {
		contextBlock = "\n" + extraContext
	}
	prompt := fmt.Sprintf(`Implement the code at the cursor%s. The prefix ends at the point where new code is needed. Write the implementation — do not echo stubs, TODOs, or placeholder returns from the prefix. Output ONLY raw code. No reasoning, no shell commands, no explanations, no markdown fences, no backticks, no preamble. Your entire response must be valid code that can be inserted directly into the file.
%s
<prefix>
%s
</prefix>
<suffix>
%s
</suffix>`, fileHint, contextBlock, prefix, suffix)

	result, err := backend.Generate(context.Background(), backend.GenerateRequest{
		Prompt:  prompt,
		Backend: backendName,
		Model:   model,
	})
	if err != nil {
		return "", err
	}

	result = stripCompletionNoise(result)
	result = stripPrefixEcho(prefix, result)
	return result, nil
}


func stripCompletionNoise(s string) string {
	var lines []string
	for _, line := range strings.Split(s, "\n") {
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, "```") {
			continue
		}
		if trimmed == "<prefix>" || trimmed == "</prefix>" ||
			trimmed == "<suffix>" || trimmed == "</suffix>" {
			continue
		}
		if strings.HasPrefix(trimmed, ":: ") {
			continue
		}
		lines = append(lines, line)
	}
	return strings.Join(lines, "\n")
}

func stripPrefixEcho(prefix, result string) string {
	tailMax := 200
	if len(prefix) < tailMax {
		tailMax = len(prefix)
	}
	for i := tailMax; i > 0; i-- {
		tail := prefix[len(prefix)-i:]
		if strings.HasPrefix(result, tail) {
			return result[len(tail):]
		}
	}
	return result
}

func crc32Str(s string) uint32 {
	return crc32.ChecksumIEEE([]byte(s))
}