package agent

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
	"github.com/tidbcloud/AgentTeam/desktop/internal/consts"
)

const piStartupSteerIdleAbsentGrace = 250 * time.Millisecond

// ReconcileRuntimeInput answers whether a correlated Pi input crossed the
// durable session boundary. Pi does not expose a read-side RPC for arbitrary
// prompt ids, but its JSONL session stores the exact AgentTeam marker. A missing
// marker is never negative evidence in an appendable stream: only a marker or
// an input-specific rejection fence may settle ownership; Process bounds the
// remaining unknown state and fails closed.
func (d *PiRPCDriver) ReconcileRuntimeInput(
	ctx context.Context,
	inputID string,
	sessionID string,
) (RuntimeInputReconciliation, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if err := ctx.Err(); err != nil {
		return RuntimeInputReconciliationUnknown, err
	}
	inputID = strings.TrimSpace(inputID)
	sessionID = strings.TrimSpace(sessionID)
	if inputID == "" || sessionID == "" {
		return RuntimeInputReconciliationUnknown, fmt.Errorf("pi: input and session ids are required for reconciliation")
	}

	d.mu.Lock()
	cfg := clonePiLaunchConfig(d.launchCfg)
	transportClosed := d.transportClosed
	liveTransport := d.stdin != nil && !transportClosed
	active := d.rootRunActive || d.promptAccepted != nil || d.promptInputID == inputID
	correlated := d.steerInFlight[inputID] || d.promptInputID == inputID
	locallyAcked := d.reconciledAckedInputIDs[inputID]
	definitivelyRejected := d.definitiveInputRejections[inputID] == strings.TrimSpace(sessionID)
	d.mu.Unlock()

	// A matching user marker is positive durable evidence.  A Pi session is an
	// appendable stream: a readable EOF, a header, old history, or an assistant
	// entry without this input's marker says nothing about whether an accepted
	// prompt will materialize later.  In particular, a replacement driver has
	// lost the old prompt correlation and must never turn those snapshots into
	// an exact replay.
	path, pathErr := piSessionFilePath(cfg, sessionID)
	if pathErr == nil {
		observed, readErr := piSessionContainsRuntimeInput(path, inputID)
		if readErr == nil && observed {
			return RuntimeInputReconciliationAcked, nil
		}
		if readErr != nil && !os.IsNotExist(readErr) {
			return RuntimeInputReconciliationUnknown, fmt.Errorf("pi: read session for input reconciliation: %w", readErr)
		}
	} else if !errors.Is(pathErr, errPiSessionNotFound) {
		return RuntimeInputReconciliationUnknown, fmt.Errorf("pi: locate session for input reconciliation: %w", pathErr)
	}
	if locallyAcked {
		return RuntimeInputReconciliationAcked, nil
	}

	if definitivelyRejected {
		// This is the narrow exception to Pi's fail-closed default: this same
		// driver observed an input-specific pre-mutation rejection (or a
		// preflight failure before bytes were written).  A replacement cannot
		// inherit this fence because it cannot distinguish that ID from an
		// accepted request whose transport outcome was lost.
		return RuntimeInputReconciliationAbsent, nil
	}
	if transportClosed {
		// A closed stdio transport does not prove that a prompt was rejected:
		// Pi may have consumed the frame before the pipe disappeared. Keep the
		// owner fail-closed until a later authoritative read finds the marker.
		return RuntimeInputReconciliationUnknown, errPiRPCTransportClosed
	}
	if !correlated {
		// This live Pi process did not accept the input. Its in-memory snapshot
		// cannot manufacture negative evidence for a retired process, and any
		// positive marker already has the same authoritative JSONL representation
		// checked above. Avoid issuing control RPCs to an unrelated replacement.
		if active {
			return RuntimeInputReconciliationPending, nil
		}
		return RuntimeInputReconciliationUnknown, nil
	}
	if liveTransport {
		return d.reconcileRuntimeInputLive(ctx, inputID, sessionID)
	}
	if correlated || active {
		return RuntimeInputReconciliationPending, nil
	}
	// Quiescent without a marker is still ambiguous.  Process owns the bounded
	// reconciliation ceiling and will persist failed/wake-failed rather than
	// replaying the original body.
	return RuntimeInputReconciliationUnknown, nil
}

// reconcileRuntimeInputLive reads lifecycle and entries from the same Pi
// process that accepted the input. Unlike the appendable JSONL stream, this is
// authoritative for both a present marker and a stable idle/absent snapshot.
func (d *PiRPCDriver) reconcileRuntimeInputLive(
	ctx context.Context,
	inputID string,
	sessionID string,
) (RuntimeInputReconciliation, error) {
	state, err := d.readRuntimeInputState(ctx)
	if err != nil {
		return RuntimeInputReconciliationUnknown, err
	}
	entries, err := d.readRuntimeInputEntries(ctx, inputID)
	if err != nil {
		return RuntimeInputReconciliationUnknown, err
	}
	if entries {
		d.mu.Lock()
		if d.reconciledAckedInputIDs == nil {
			d.reconciledAckedInputIDs = make(map[string]bool)
		}
		d.reconciledAckedInputIDs[inputID] = true
		delete(d.startupSteerIdleObservedAt, inputID)
		d.mu.Unlock()
		if !state.IsStreaming && !state.IsCompacting &&
			d.settleAcceptedInputFromAuthoritativeIdle(inputID, sessionID, true) {
			// Let Process consume RuntimeIdle before surfacing ownership; otherwise
			// a queued successor can race through its stale busy/steer decision.
			return RuntimeInputReconciliationPending, nil
		}
		return RuntimeInputReconciliationAcked, nil
	}
	if state.IsStreaming || state.IsCompacting {
		d.mu.Lock()
		delete(d.startupSteerIdleObservedAt, inputID)
		d.mu.Unlock()
		return RuntimeInputReconciliationPending, nil
	}

	// Pi can accept a startup steer before committing its queued user entry.
	// Require idle+absent to remain stable beyond a short grace window so a late
	// marker found by the next snapshot wins.
	d.mu.Lock()
	startupSteer := !d.rootRunActive && d.steerInFlight[inputID]
	firstIdle := d.startupSteerIdleObservedAt[inputID]
	if startupSteer && firstIdle.IsZero() {
		if d.startupSteerIdleObservedAt == nil {
			d.startupSteerIdleObservedAt = make(map[string]time.Time)
		}
		d.startupSteerIdleObservedAt[inputID] = time.Now()
	}
	d.mu.Unlock()
	if startupSteer && (firstIdle.IsZero() || time.Since(firstIdle) < piStartupSteerIdleAbsentGrace) {
		return RuntimeInputReconciliationPending, nil
	}

	settled := d.settleAcceptedInputFromAuthoritativeIdle(inputID, sessionID, false)
	d.rememberDefinitiveInputRejection(inputID, sessionID)
	if settled {
		// Sequence RuntimeIdle before the next reconciliation consumes the
		// negative fence and returns the durable body to queued.
		return RuntimeInputReconciliationPending, nil
	}
	return RuntimeInputReconciliationAbsent, nil
}

type piRuntimeInputState struct {
	IsStreaming  bool `json:"isStreaming"`
	IsCompacting bool `json:"isCompacting"`
}

func (d *PiRPCDriver) readRuntimeInputState(ctx context.Context) (piRuntimeInputState, error) {
	response, err := d.requestCommand(ctx, map[string]any{"type": "get_state"})
	if err != nil {
		return piRuntimeInputState{}, fmt.Errorf("pi: get_state for input reconciliation: %w", err)
	}
	var state piRuntimeInputState
	if len(response.Data) == 0 || json.Unmarshal(response.Data, &state) != nil {
		return piRuntimeInputState{}, fmt.Errorf("pi: get_state for input reconciliation returned invalid data")
	}
	return state, nil
}

func (d *PiRPCDriver) readRuntimeInputEntries(ctx context.Context, inputID string) (bool, error) {
	response, err := d.requestCommand(ctx, map[string]any{"type": "get_entries"})
	if err != nil {
		return false, fmt.Errorf("pi: get_entries for input reconciliation: %w", err)
	}
	var snapshot struct {
		Entries []struct {
			Type    string          `json:"type"`
			Message json.RawMessage `json:"message"`
		} `json:"entries"`
	}
	if len(response.Data) == 0 || json.Unmarshal(response.Data, &snapshot) != nil {
		return false, fmt.Errorf("pi: get_entries for input reconciliation returned invalid data")
	}
	for _, entry := range snapshot.Entries {
		if strings.TrimSpace(entry.Type) != "message" || len(entry.Message) == 0 {
			continue
		}
		var message piAgentMessage
		if json.Unmarshal(entry.Message, &message) != nil {
			return false, fmt.Errorf("pi: get_entries contains malformed message data")
		}
		if strings.TrimSpace(message.Role) == "user" && piRuntimeInputIDFromContent(message.Content) == inputID {
			return true, nil
		}
	}
	return false, nil
}

// settleAcceptedInputFromAuthoritativeIdle repairs the local generation that
// handleLine opened at prompt preflight success when Pi later reports that no
// agent run is active. Without this repair Process would requeue the Absent
// owner but continue treating the phantom root as busy forever.
func (d *PiRPCDriver) settleAcceptedInputFromAuthoritativeIdle(inputID, sessionID string, acked bool) bool {
	d.mu.Lock()
	correlated := d.promptInputID == inputID || d.steerInFlight[inputID]
	if !correlated {
		d.mu.Unlock()
		return false
	}
	// Pi accepts `steer` even when no agent run is active. That startup race has
	// no root generation to settle, but the live idle snapshot is still
	// authoritative negative evidence for the accepted steer. Retire its
	// correlation and publish RuntimeIdle so Process stops routing the replay
	// back through the busy-steer path.
	if !d.rootRunActive {
		d.promptInputID = ""
		delete(d.steerInFlight, inputID)
		delete(d.startupSteerIdleObservedAt, inputID)
		d.mu.Unlock()
		d.emit(translatedEvent{
			Kind:      consts.ParsedEventKindRuntimeIdle,
			SessionID: sessionID,
		})
		return true
	}
	gen := d.rootRunGeneration
	turnID := d.rootRunTurnID
	settled := d.rootRunSettled
	rootStarted := d.rootStartEmitted
	d.rootRunActive = false
	d.promptInputID = ""
	delete(d.steerInFlight, inputID)
	d.assistant = nil
	d.mu.Unlock()

	errorMessage := "pi runtime became idle without an agent_settled event"
	if !acked {
		errorMessage = "pi prompt was accepted but did not ack before the runtime became idle"
	}
	if rootStarted {
		// A real streamed root start requires its matching terminal boundary.
		d.emitTerminal(gen, turnID, sessionID, "error", errorMessage, "error", false)
	} else {
		// Preflight opened only a local generation; Pi never emitted agent_start.
		// Emitting a synthetic TurnEnd would trigger between-turn compaction and
		// delay the durable retry for a turn that never existed. An Error event
		// would also terminalize the wake that is about to return to queued, so
		// publish only the authoritative idle boundary; the durable ownership
		// transition records the absent/retry outcome.
		d.emit(translatedEvent{
			Kind:      consts.ParsedEventKindRuntimeIdle,
			SessionID: sessionID,
			Metadata:  map[string]any{runtimeCompletedTurnIDKey: firstNonEmptyString(turnID, fmt.Sprintf("pi-gen-%d", gen))},
		})
	}
	d.settleRootRunLocked(settled)
	return true
}

// rememberDefinitiveInputRejection installs the only negative fence that Pi
// can provide for a runtime input.  It is deliberately scoped to this driver
// instance/session; no replacement may manufacture it from a transcript
// snapshot.
func (d *PiRPCDriver) rememberDefinitiveInputRejection(inputID, sessionID string) {
	inputID = strings.TrimSpace(inputID)
	sessionID = strings.TrimSpace(sessionID)
	if inputID == "" || sessionID == "" {
		return
	}
	d.mu.Lock()
	if d.definitiveInputRejections == nil {
		d.definitiveInputRejections = make(map[string]string)
	}
	d.definitiveInputRejections[inputID] = sessionID
	d.mu.Unlock()
}

func (d *PiRPCDriver) forgetDefinitiveInputRejection(inputID, sessionID string) {
	inputID = strings.TrimSpace(inputID)
	sessionID = strings.TrimSpace(sessionID)
	if inputID == "" || sessionID == "" {
		return
	}
	d.mu.Lock()
	if d.definitiveInputRejections[inputID] == sessionID {
		delete(d.definitiveInputRejections, inputID)
	}
	d.mu.Unlock()
}

// piSessionContainsRuntimeInput reads only user message content from a Pi
// session. A byte-level search is not sufficient: an assistant/provider
// response can quote the marker and must not settle ownership for the user's
// delivery. Malformed or unsafe session files are unknown rather than absent,
// so the Process layer never retries an input after losing authoritative
// evidence.
func piSessionContainsRuntimeInput(path, inputID string) (bool, error) {
	info, err := os.Lstat(path)
	if err != nil {
		return false, err
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
		return false, fmt.Errorf("pi: session path %q is not a safe regular file", path)
	}
	file, err := os.Open(path)
	if err != nil {
		return false, err
	}
	defer file.Close()

	markerID := strings.TrimSpace(inputID)
	found := false
	scanner := bufio.NewScanner(file)
	scanner.Buffer(make([]byte, 64*1024), piRPCMaxLineBytes)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" {
			continue
		}
		var record struct {
			Type    string          `json:"type"`
			Message json.RawMessage `json:"message"`
		}
		if err := json.Unmarshal([]byte(line), &record); err != nil {
			return false, fmt.Errorf("pi: malformed session entry: %w", err)
		}
		if record.Type != "message" {
			continue
		}
		if len(record.Message) == 0 {
			return false, fmt.Errorf("pi: message entry has no message payload")
		}
		var message struct {
			Role    string          `json:"role"`
			Content json.RawMessage `json:"content"`
		}
		if err := json.Unmarshal(record.Message, &message); err != nil {
			return false, fmt.Errorf("pi: malformed session message: %w", err)
		}
		if len(message.Content) == 0 || string(message.Content) == "null" {
			return false, fmt.Errorf("pi: session message has no content")
		}
		if strings.TrimSpace(message.Role) != "user" {
			continue
		}
		var text string
		if json.Unmarshal(message.Content, &text) == nil {
			if piRuntimeInputIDFromText(text) == markerID {
				found = true
			}
			continue
		}
		var blocks []struct {
			Type string `json:"type"`
			Text string `json:"text"`
		}
		if json.Unmarshal(message.Content, &blocks) != nil {
			return false, fmt.Errorf("pi: session message content is not text or content blocks")
		}
		for _, block := range blocks {
			if block.Type == "text" && piRuntimeInputIDFromText(block.Text) == markerID {
				found = true
			}
		}
	}
	if err := scanner.Err(); err != nil {
		return false, err
	}
	// Do not return early on a positive marker. A malformed/truncated suffix
	// means the JSONL snapshot is not authoritative; returning acked from a
	// partial read would let a damaged session advance ownership without a
	// trustworthy transcript boundary.
	return found, nil
}

// captureInputBoundary records the current session leaf entry id before a AgentTeam
// message is delivered, so PrepareRecoveryPointForMessage can restore to the
// exact pre-message boundary. It is best-effort: a transient read failure keeps
// the previous boundary and the caller still proceeds with delivery.
func (d *PiRPCDriver) captureInputBoundary(sessionID string) {
	sessionID = strings.TrimSpace(sessionID)
	if sessionID == "" {
		return
	}
	leafID := ""
	// Prefer the live transport (authoritative current leaf), then fall back to
	// the durable session file. Both are best-effort; a fresh session before
	// the first delivery legitimately has no leaf yet.
	if leaf, err := d.currentLeafEntryID(d.launchCfg, sessionID); err == nil && leaf != "" {
		leafID = leaf
	}
	if leafID == "" {
		if path, err := piSessionFilePath(d.launchCfg, sessionID); err == nil {
			if leaf, leafErr := piSessionLeafID(path); leafErr == nil && leaf != "" {
				leafID = leaf
			}
		}
	}
	if leafID == "" {
		return
	}
	d.mu.Lock()
	d.entryIDAtLastInput = leafID
	d.mu.Unlock()
}

// currentLeafEntryID queries the running pi process for the current session
// leaf entry id via the get_entries RPC. When the transport is unavailable
// (for example during detached recovery preparation), it falls back to reading
// the durable session file.
func (d *PiRPCDriver) currentLeafEntryID(cfg drivercontract.LaunchConfig, sessionID string) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), piRPCRequestTimeout)
	defer cancel()
	// A live transport get_entries round-trip is authoritative when available.
	if leafID, err := d.getEntries(ctx); err == nil && leafID != "" {
		return leafID, nil
	}
	// Fall back to the durable file.
	path, err := piSessionFilePath(cfg, sessionID)
	if err != nil {
		return "", err
	}
	leafID, err := piSessionLeafID(path)
	if err != nil {
		return "", err
	}
	return leafID, nil
}

// getEntries issues a get_entries RPC and returns the current leaf id.
func (d *PiRPCDriver) getEntries(ctx context.Context) (string, error) {
	envelope, err := d.syncCommand(ctx, map[string]any{"type": "get_entries"})
	if err != nil {
		return "", err
	}
	var result struct {
		LeafID string `json:"leafId"`
	}
	if len(envelope.Data) > 0 {
		_ = json.Unmarshal(envelope.Data, &result)
	}
	return result.LeafID, nil
}

// syncCommand sends a command and waits for its correlated response.
func (d *PiRPCDriver) syncCommand(ctx context.Context, command map[string]any) (piRPCFrame, error) {
	if err := ctx.Err(); err != nil {
		return piRPCFrame{}, err
	}
	d.mu.Lock()
	if d.stdin == nil || d.transportClosed {
		d.mu.Unlock()
		return piRPCFrame{}, io.ErrClosedPipe
	}
	d.nextID++
	id := d.nextID
	command["id"] = id
	responseCh := make(chan piRPCFrame, 1)
	d.pending[id] = responseCh
	d.mu.Unlock()

	data, err := json.Marshal(command)
	if err != nil {
		d.forgetPending(id)
		return piRPCFrame{}, err
	}
	data = append(data, '\n')
	if err := d.writeFrame(data); err != nil {
		d.forgetPending(id)
		return piRPCFrame{}, errPiRPCTransportClosed
	}
	select {
	case <-ctx.Done():
		d.forgetPending(id)
		return piRPCFrame{}, ctx.Err()
	case response := <-responseCh:
		if response.Success != nil && !*response.Success {
			return response, fmt.Errorf("pi RPC %s: %s", command["type"], firstNonEmptyString(response.Error, "request failed"))
		}
		return response, nil
	}
}

// piSessionLeafID reads the last entry id from a session JSONL file.
func piSessionLeafID(path string) (string, error) {
	f, err := os.Open(path)
	if err != nil {
		return "", err
	}
	defer f.Close()
	var leafID string
	scanner := bufio.NewScanner(f)
	scanner.Buffer(make([]byte, 64*1024), piRPCMaxLineBytes)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" {
			continue
		}
		var entry struct {
			Type string `json:"type"`
			ID   string `json:"id"`
		}
		if json.Unmarshal([]byte(line), &entry) != nil {
			continue
		}
		if entry.Type == "session" {
			continue
		}
		if entry.ID != "" {
			leafID = entry.ID
		}
	}
	if err := scanner.Err(); err != nil {
		return "", err
	}
	return leafID, nil
}

// piTruncateSessionAtBoundary truncates the session JSONL file so it keeps the
// header and every entry through boundaryEntryID (inclusive), dropping all
// later context. This is the Pi rewind restore primitive.
func piTruncateSessionAtBoundary(ctx context.Context, sessionPath string, boundaryEntryID string, boundaryKind string) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	sessionPath = strings.TrimSpace(sessionPath)
	if sessionPath == "" {
		return fmt.Errorf("pi: session path is required")
	}
	info, err := os.Lstat(sessionPath)
	if err != nil {
		return fmt.Errorf("pi: inspect session file: %w", err)
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
		return fmt.Errorf("pi: session path %q is not a safe regular file", sessionPath)
	}
	f, err := os.Open(sessionPath)
	if err != nil {
		return fmt.Errorf("pi: open session file: %w", err)
	}
	defer f.Close()

	var header string
	var keep []string
	var hitBoundary bool
	dropMessages := strings.TrimSpace(boundaryEntryID) == ""
	scanner := bufio.NewScanner(f)
	scanner.Buffer(make([]byte, 64*1024), piRPCMaxLineBytes)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" {
			continue
		}
		var entry struct {
			Type string `json:"type"`
			ID   string `json:"id"`
		}
		if json.Unmarshal([]byte(line), &entry) != nil {
			continue
		}
		if entry.Type == "session" {
			header = line
			continue
		}
		// before_first_input boundaries drop every message (and every entry
		// that depends on one) while keeping harmless metadata.
		if dropMessages {
			if entry.Type == "message" {
				continue
			}
			keep = append(keep, line)
			continue
		}
		keep = append(keep, line)
		if strings.TrimSpace(boundaryEntryID) != "" && entry.ID == boundaryEntryID {
			hitBoundary = true
			break
		}
	}
	if err := scanner.Err(); err != nil {
		return fmt.Errorf("pi: read session file: %w", err)
	}
	if strings.TrimSpace(boundaryEntryID) != "" && !hitBoundary {
		return fmt.Errorf("pi: boundary entry %q not found in session", boundaryEntryID)
	}
	_ = header
	if err := ctx.Err(); err != nil {
		return err
	}
	return piWriteSessionTruncation(ctx, sessionPath, header, keep)
}

// piWriteSessionTruncation atomically rewrites the session file with only the
// header and kept entries.
func piWriteSessionTruncation(ctx context.Context, sessionPath string, header string, keep []string) error {
	dir := filepath.Dir(sessionPath)
	tmp, err := os.CreateTemp(dir, ".pi-session-trunc-*")
	if err != nil {
		return fmt.Errorf("pi: create truncation temp file: %w", err)
	}
	tmpName := tmp.Name()
	cleanup := func() {
		_ = tmp.Close()
		_ = os.Remove(tmpName)
	}
	writer := bufio.NewWriter(tmp)
	write := func(line string) error {
		if line == "" {
			return nil
		}
		_, err := writer.WriteString(line + "\n")
		return err
	}
	if strings.TrimSpace(header) != "" {
		if err := write(header); err != nil {
			cleanup()
			return fmt.Errorf("pi: write truncated header: %w", err)
		}
	}
	for _, line := range keep {
		if err := write(line); err != nil {
			cleanup()
			return fmt.Errorf("pi: write truncated entry: %w", err)
		}
	}
	if err := writer.Flush(); err != nil {
		cleanup()
		return fmt.Errorf("pi: flush truncated session: %w", err)
	}
	if err := tmp.Sync(); err != nil {
		cleanup()
		return fmt.Errorf("pi: sync truncated session: %w", err)
	}
	if err := ctx.Err(); err != nil {
		cleanup()
		return err
	}
	if err := tmp.Close(); err != nil {
		_ = os.Remove(tmpName)
		return fmt.Errorf("pi: close truncated session: %w", err)
	}
	if err := os.Chmod(tmpName, 0o600); err != nil {
		_ = os.Remove(tmpName)
		return fmt.Errorf("pi: chmod truncated session: %w", err)
	}
	if err := os.Rename(tmpName, sessionPath); err != nil {
		_ = os.Remove(tmpName)
		return fmt.Errorf("pi: replace session file: %w", err)
	}
	// Best-effort directory sync so a crash does not lose the rename.
	if dirHandle, err := os.Open(dir); err == nil {
		_ = dirHandle.Sync()
		_ = dirHandle.Close()
	}
	return nil
}

// piFindSessionFiles returns all session files under the runtime state dir for
// a session id prefix, used by resume discovery.
func piFindSessionFiles(cfg drivercontract.LaunchConfig, sessionID string) ([]string, error) {
	stateDir := strings.TrimSpace(piSessionDir(cfg))
	if stateDir == "" {
		return nil, nil
	}
	entries, err := os.ReadDir(stateDir)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return nil, nil
		}
		return nil, err
	}
	var matches []string
	sessionID = strings.TrimSpace(sessionID)
	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}
		name := entry.Name()
		if !strings.HasSuffix(name, ".jsonl") {
			continue
		}
		if sessionID != "" && !strings.Contains(name, sessionID) {
			continue
		}
		matches = append(matches, filepath.Join(stateDir, name))
	}
	sort.Strings(matches)
	return matches, nil
}

// piResolveLatestSessionFile returns the most recently modified matching
// session file.
func piResolveLatestSessionFile(cfg drivercontract.LaunchConfig, sessionID string) (string, error) {
	matches, err := piFindSessionFiles(cfg, sessionID)
	if err != nil {
		return "", err
	}
	if len(matches) == 0 {
		return "", os.ErrNotExist
	}
	best := matches[0]
	var bestMod time.Time
	for _, match := range matches {
		if info, err := os.Stat(match); err == nil && info.ModTime().After(bestMod) {
			best = match
			bestMod = info.ModTime()
		}
	}
	return best, nil
}
