/internal/agent/pi.go

package agent

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os/exec"
	"strings"
	"sync"

	"github.com/kunchenguid/no-mistakes/internal/shellenv"
)

// piAgent spawns the pi CLI for each invocation. Pi reads its prompt from
// stdin and emits JSONL on stdout when --mode json is set. The lifecycle is
// codex-shaped: one process per Run, no managed server.
type piAgent struct {
	bin       string
	extraArgs []string
}

func (a *piAgent) Name() string { return "pi" }

func (a *piAgent) ReportsAgentAttempts() bool { return true }

func (a *piAgent) Run(ctx context.Context, opts RunOpts) (*Result, error) {
	return runWithRetry(ctx, "pi", opts, claudeMaxRetries, classifyTransient, nil, func() (*Result, error) {
		return a.runOnce(ctx, opts)
	})
}

func (a *piAgent) Close() error { return nil }

func (a *piAgent) runOnce(ctx context.Context, opts RunOpts) (*Result, error) {
	args := a.buildArgs()
	cmd := exec.CommandContext(ctx, a.bin, args...)
	cmd.Dir = opts.CWD
	cmd.Env = gitSafeEnv(opts.CWD)
	shellenv.ConfigureShellCommand(cmd)

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, fmt.Errorf("pi stdin pipe: %w", err)
	}

	started, err := startNativeAgentCommand(cmd)
	if err != nil {
		return nil, fmt.Errorf("pi start: %w", err)
	}
	defer started.closePipes()
	pid := started.pid()
	emitAgentStarted(opts, "pi", pid)

	prompt := buildPiPrompt(opts.Prompt, opts.JSONSchema)
	go func() {
		defer stdin.Close()
		_, _ = io.WriteString(stdin, prompt)
	}()

	var stderrBuf []byte
	var stderrWG sync.WaitGroup
	stderrWG.Add(1)
	go func() {
		defer stderrWG.Done()
		stderrBuf, _ = io.ReadAll(started.stderr)
	}()

	pp := &piParser{onChunk: opts.OnChunk}
	if err := pp.parse(ctx, started.stdout); err != nil {
		err = started.waitAfterParseError(err)
		stderrWG.Wait()
		retErr := fmt.Errorf("pi parse events: %w", err)
		emitAgentExited(opts, "pi", pid, retErr)
		return nil, retErr
	}

	waitErr := started.wait()
	stderrWG.Wait()
	if waitErr != nil {
		stderr := strings.TrimSpace(string(stderrBuf))
		if stderr != "" {
			retErr := fmt.Errorf("pi exited: %w: %s", waitErr, stderr)
			emitAgentExited(opts, "pi", pid, retErr)
			return nil, retErr
		}
		retErr := fmt.Errorf("pi exited: %w", waitErr)
		emitAgentExited(opts, "pi", pid, retErr)
		return nil, retErr
	}

	if pp.assistantError != "" {
		retErr := fmt.Errorf("pi reported error: %s", pp.assistantError)
		emitAgentExited(opts, "pi", pid, retErr)
		return nil, retErr
	}

	text := pp.finalText()
	res, err := finalizeTextResult("pi", text, opts.JSONSchema, pp.usage)
	emitAgentExited(opts, "pi", pid, err)
	return res, err
}

// buildArgs returns the Pi argv for one invocation. User extras come first
// (so user --provider/--model take effect), then the managed flags that
// no-mistakes requires for JSONL parsing.
func (a *piAgent) buildArgs() []string {
	args := make([]string, 0, len(a.extraArgs)+3)
	args = append(args, a.extraArgs...)
	args = append(args, "--mode", "json", "--no-session")
	return args
}

// buildPiPrompt appends a JSON-output contract to the user prompt when a
// schema is provided. Pi has no equivalent of codex's --output-schema flag,
// so we inline the schema in the prompt the same way gnhf does.
func buildPiPrompt(prompt string, schema json.RawMessage) string {
	if len(schema) == 0 {
		return prompt
	}
	pretty, err := json.MarshalIndent(json.RawMessage(schema), "", "  ")
	if err != nil {
		pretty = []byte(schema)
	}
	return prompt + "\n\n## no-mistakes final output contract\n\n" +
		"When the iteration is complete, your final assistant response must be only valid JSON matching this JSON Schema. " +
		"Do not wrap it in Markdown fences. Do not include prose before or after the JSON object.\n\n" +
		string(pretty)
}

// piParser tracks the streaming state of one Pi run. It accumulates text
// deltas, captures the final assistant text and usage, and surfaces any
// reported assistant error.
type piParser struct {
	onChunk func(string)

	streamText     map[int]string
	completeText   map[int]string
	finalAssistant map[string]any
	usage          TokenUsage
	seenUsage      map[string]struct{}
	assistantError string
}

func (p *piParser) parse(ctx context.Context, r io.Reader) error {
	scanner := bufio.NewScanner(r)
	scanner.Buffer(make([]byte, 0, 64*1024), 256*1024*1024)

	p.streamText = make(map[int]string)
	p.completeText = make(map[int]string)
	p.seenUsage = make(map[string]struct{})

	for scanner.Scan() {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}

		line := scanner.Bytes()
		if len(line) == 0 {
			continue
		}

		var event map[string]any
		if err := json.Unmarshal(line, &event); err != nil {
			continue
		}
		p.handleEvent(event)
	}

	return scanner.Err()
}

func (p *piParser) handleEvent(event map[string]any) {
	typ, _ := event["type"].(string)
	switch typ {
	case "message_update":
		p.rememberAssistant(event["message"])
		p.handleAssistantEvent(event["assistantMessageEvent"])
	case "message_end", "turn_end":
		p.rememberAssistant(event["message"])
		p.recordAssistantUsage(event["message"])
	case "agent_end":
		p.rememberAgentEnd(event["messages"])
	}
}

func (p *piParser) rememberAssistant(raw any) {
	msg, ok := raw.(map[string]any)
	if !ok {
		return
	}
	if role, _ := msg["role"].(string); role != "assistant" {
		return
	}
	p.finalAssistant = msg

	if reason, _ := msg["stopReason"].(string); reason == "error" || reason == "aborted" {
		p.assistantError = piFirstString(msg, "errorMessage", "error", "message")
		if p.assistantError == "" {
			p.assistantError = fmt.Sprintf("stopReason=%s", reason)
		}
	} else {
		p.assistantError = ""
	}
}

func (p *piParser) rememberAgentEnd(raw any) {
	messages, ok := raw.([]any)
	if !ok {
		return
	}

	total := TokenUsage{}
	seen := make(map[string]struct{})
	hasUsage := false
	for i, rawMsg := range messages {
		msg, ok := rawMsg.(map[string]any)
		if !ok || msg["role"] != "assistant" {
			continue
		}
		usageMap, ok := msg["usage"].(map[string]any)
		if !ok {
			continue
		}
		usage := piUsageFrom(usageMap)
		if piUsageIsZero(usage) {
			continue
		}
		key := piUsageKey(msg)
		if key == "" {
			key = fmt.Sprintf("agent_end:%d", i)
		}
		if _, ok := seen[key]; ok {
			continue
		}
		seen[key] = struct{}{}
		total = piUsageAdd(total, usage)
		hasUsage = true
	}
	if hasUsage {
		p.usage = total
		p.seenUsage = make(map[string]struct{}, len(seen))
		for key := range seen {
			p.seenUsage[key] = struct{}{}
		}
	}

	for i := len(messages) - 1; i >= 0; i-- {
		if msg, ok := messages[i].(map[string]any); ok && msg["role"] == "assistant" {
			p.rememberAssistant(msg)
			return
		}
	}
}

func (p *piParser) recordAssistantUsage(raw any) {
	msg, ok := raw.(map[string]any)
	if !ok || msg["role"] != "assistant" {
		return
	}
	usageMap, ok := msg["usage"].(map[string]any)
	if !ok {
		return
	}
	usage := piUsageFrom(usageMap)
	if piUsageIsZero(usage) {
		return
	}
	key := piUsageKey(msg)
	if key == "" {
		encoded, err := json.Marshal([]any{msg["role"], msg["stopReason"], msg["content"], msg["usage"]})
		if err != nil {
			return
		}
		key = string(encoded)
	}
	if p.seenUsage == nil {
		p.seenUsage = make(map[string]struct{})
	}
	if _, ok := p.seenUsage[key]; ok {
		return
	}
	p.seenUsage[key] = struct{}{}
	p.usage = piUsageAdd(p.usage, usage)
}

func (p *piParser) handleAssistantEvent(raw any) {
	evt, ok := raw.(map[string]any)
	if !ok {
		return
	}
	idx := piIntField(evt, "contentIndex", "content_index")
	switch evt["type"] {
	case "text_delta":
		// Emit just the incremental delta. no-mistakes' OnChunk consumers
		// (TUI log line buffer, file logger) expect appended text, not
		// cumulative state.
		delta := piFirstString(evt, "delta", "text", "content")
		if delta == "" {
			return
		}
		p.streamText[idx] += delta
		if p.onChunk != nil {
			p.onChunk(delta)
		}
	case "text_end":
		// Capture the complete text for final-result resolution. Don't
		// re-emit to OnChunk: the deltas already covered it. If the event
		// carries the full text (Pi's normal shape), prefer that over the
		// delta accumulator since it's authoritative.
		text := piFirstString(evt, "text", "content")
		if text == "" {
			text = p.streamText[idx]
		}
		p.completeText[idx] = text
	}
}

// finalText returns the final assistant text, preferring (in order) the
// content of the last assistant message, the text_end-completed deltas, and
// finally the in-flight stream buffer.
func (p *piParser) finalText() string {
	if text := strings.TrimSpace(textFromAssistantMessage(p.finalAssistant)); text != "" {
		return text
	}
	if text := strings.TrimSpace(joinByIndex(p.completeText)); text != "" {
		return text
	}
	return strings.TrimSpace(joinByIndex(p.streamText))
}

func textFromAssistantMessage(msg map[string]any) string {
	if msg == nil {
		return ""
	}
	switch v := msg["content"].(type) {
	case string:
		return v
	case []any:
		var b strings.Builder
		for _, block := range v {
			switch t := block.(type) {
			case string:
				b.WriteString(t)
			case map[string]any:
				if s, ok := t["text"].(string); ok {
					b.WriteString(s)
					continue
				}
				if s, ok := t["content"].(string); ok {
					b.WriteString(s)
				}
			}
		}
		return b.String()
	}
	if s, ok := msg["text"].(string); ok {
		return s
	}
	return ""
}

func joinByIndex(parts map[int]string) string {
	if len(parts) == 0 {
		return ""
	}
	max := -1
	for k := range parts {
		if k > max {
			max = k
		}
	}
	var b strings.Builder
	for i := 0; i <= max; i++ {
		b.WriteString(parts[i])
	}
	return b.String()
}

func piFirstString(m map[string]any, names ...string) string {
	for _, n := range names {
		if v, ok := m[n].(string); ok {
			return v
		}
	}
	return ""
}

func piIntField(m map[string]any, names ...string) int {
	for _, n := range names {
		switch v := m[n].(type) {
		case float64:
			return int(v)
		case int:
			return v
		case json.Number:
			if i, err := v.Int64(); err == nil {
				return int(i)
			}
		}
	}
	return 0
}

func piUsageFrom(usage map[string]any) TokenUsage {
	_, cacheCreationReported := usage["cacheWrite"]
	return TokenUsage{
		Reported:              len(usage) > 0,
		CacheCreationReported: cacheCreationReported,
		InputTokens:           piIntField(usage, "input"),
		OutputTokens:          piIntField(usage, "output"),
		CacheReadTokens:       piIntField(usage, "cacheRead"),
		CacheCreationTokens:   piIntField(usage, "cacheWrite"),
	}
}

func piUsageAdd(a, b TokenUsage) TokenUsage {
	return TokenUsage{
		Reported:              a.Reported || b.Reported,
		CacheCreationReported: a.CacheCreationReported || b.CacheCreationReported,
		InputTokens:           a.InputTokens + b.InputTokens,
		OutputTokens:          a.OutputTokens + b.OutputTokens,
		CacheReadTokens:       a.CacheReadTokens + b.CacheReadTokens,
		CacheCreationTokens:   a.CacheCreationTokens + b.CacheCreationTokens,
	}
}

func piUsageIsZero(usage TokenUsage) bool {
	return usage.InputTokens == 0 && usage.OutputTokens == 0 &&
		usage.CacheReadTokens == 0 && usage.CacheCreationTokens == 0
}

func piUsageKey(msg map[string]any) string {
	for _, name := range []string{"responseId", "id"} {
		if value, ok := msg[name].(string); ok && value != "" {
			return name + ":" + value
		}
	}
	return ""
}