package agent

import (
	"context"
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

func TestPiRPCSpawnResumePreflightFailureIsFailClosedAndReaped(t *testing.T) {
	stateDir, sessionID := piResumedSessionFixture(t, "11111111-1111-1111-1111-111111111111")
	rejected := false
	binary, _ := writeFakePiRPCModeBinary(t, []piFakeScript{
		// Source-faithful pi-mono response: success=false + error, with no data.
		{Cmd: "get_state", Resp: &rejected, Error: "session preflight rejected"},
	})
	lifecycleDir := t.TempDir()
	cfg := piTestConfig(t, binary)
	cfg.Env["AGENTTEAM_TEST_PI_LIFECYCLE_DIR"] = lifecycleDir
	cfg.RuntimeStateDir = stateDir
	cfg.SessionID = sessionID
	setTestLaunchSurface(&cfg, "FreshPrompt", "must not be replayed")

	result, err := (&PiRPCDriver{}).Spawn(context.Background(), cfg)
	require.ErrorContains(t, err, "session preflight")
	require.Nil(t, result.FreshSessionFallback, "Pi has no pinned-protocol fresh-session authority")
	assertPiFailedPreflightWasReaped(t, lifecycleDir)
}

func TestPiRPCSpawnResumePreflightRejectsOrdinaryProviderError(t *testing.T) {
	stateDir, sessionID := piResumedSessionFixture(t, "22222222-2222-2222-2222-222222222222")
	rejected := false
	binary, _ := writeFakePiRPCModeBinary(t, []piFakeScript{
		{Cmd: "get_state", Resp: &rejected, Error: "model rejected"},
	})
	lifecycleDir := t.TempDir()
	cfg := piTestConfig(t, binary)
	cfg.Env["AGENTTEAM_TEST_PI_LIFECYCLE_DIR"] = lifecycleDir
	cfg.RuntimeStateDir = stateDir
	cfg.SessionID = sessionID
	setTestLaunchSurface(&cfg, "FreshPrompt", "must not be replayed")

	result, err := (&PiRPCDriver{}).Spawn(context.Background(), cfg)
	require.ErrorContains(t, err, "model rejected")
	require.Nil(t, result.FreshSessionFallback)
	assertPiFailedPreflightWasReaped(t, lifecycleDir)
}

func TestPiRPCSpawnResumePreflightRejectsMissingSessionState(t *testing.T) {
	stateDir, sessionID := piResumedSessionFixture(t, "33333333-3333-3333-3333-333333333333")
	success := true
	binary, _ := writeFakePiRPCModeBinary(t, []piFakeScript{
		{Cmd: "get_state", Resp: &success},
	})
	lifecycleDir := t.TempDir()
	cfg := piTestConfig(t, binary)
	cfg.Env["AGENTTEAM_TEST_PI_LIFECYCLE_DIR"] = lifecycleDir
	cfg.RuntimeStateDir = stateDir
	cfg.SessionID = sessionID
	setTestLaunchSurface(&cfg, "FreshPrompt", "must not be replayed")

	result, err := (&PiRPCDriver{}).Spawn(context.Background(), cfg)
	require.ErrorContains(t, err, "no session state")
	require.Nil(t, result.FreshSessionFallback)
	assertPiFailedPreflightWasReaped(t, lifecycleDir)
}

func TestPiRPCSpawnResumePreflightRejectsMalformedSessionState(t *testing.T) {
	stateDir, sessionID := piResumedSessionFixture(t, "77777777-7777-7777-7777-777777777777")
	success := true
	binary, _ := writeFakePiRPCModeBinary(t, []piFakeScript{
		{Cmd: "get_state", Resp: &success, RawData: json.RawMessage(`[]`)},
	})
	lifecycleDir := t.TempDir()
	cfg := piTestConfig(t, binary)
	cfg.Env["AGENTTEAM_TEST_PI_LIFECYCLE_DIR"] = lifecycleDir
	cfg.RuntimeStateDir = stateDir
	cfg.SessionID = sessionID
	setTestLaunchSurface(&cfg, "FreshPrompt", "must not be replayed")

	result, err := (&PiRPCDriver{}).Spawn(context.Background(), cfg)
	require.ErrorContains(t, err, "decode pi get_state response")
	require.Nil(t, result.FreshSessionFallback)
	assertPiFailedPreflightWasReaped(t, lifecycleDir)
}

func TestPiRPCSpawnResumePreflightRejectsSessionIDMismatch(t *testing.T) {
	stateDir, sessionID := piResumedSessionFixture(t, "44444444-4444-4444-4444-444444444444")
	success := true
	binary, _ := writeFakePiRPCModeBinary(t, []piFakeScript{
		{Cmd: "get_state", Resp: &success, Data: map[string]any{"sessionId": "55555555-5555-5555-5555-555555555555"}},
	})
	lifecycleDir := t.TempDir()
	cfg := piTestConfig(t, binary)
	cfg.Env["AGENTTEAM_TEST_PI_LIFECYCLE_DIR"] = lifecycleDir
	cfg.RuntimeStateDir = stateDir
	cfg.SessionID = sessionID
	setTestLaunchSurface(&cfg, "FreshPrompt", "must not be replayed")

	result, err := (&PiRPCDriver{}).Spawn(context.Background(), cfg)
	require.ErrorContains(t, err, "returned session")
	require.Nil(t, result.FreshSessionFallback)
	assertPiFailedPreflightWasReaped(t, lifecycleDir)
}

func TestPiRPCSpawnResumePreflightTransportFailureDoesNotFreshFallback(t *testing.T) {
	stateDir, sessionID := piResumedSessionFixture(t, "66666666-6666-6666-6666-666666666666")
	binary, _ := writeFakePiRPCModeBinary(t, []piFakeScript{{Cmd: "get_state"}})
	lifecycleDir := t.TempDir()
	cfg := piTestConfig(t, binary)
	cfg.Env["AGENTTEAM_TEST_PI_LIFECYCLE_DIR"] = lifecycleDir
	cfg.RuntimeStateDir = stateDir
	cfg.SessionID = sessionID
	setTestLaunchSurface(&cfg, "FreshPrompt", "must not be replayed")
	ctx, cancel := context.WithCancel(context.Background())
	outcomes := make(chan struct {
		err        error
		noFallback bool
	}, 1)
	go func() {
		result, err := (&PiRPCDriver{}).Spawn(ctx, cfg)
		outcomes <- struct {
			err        error
			noFallback bool
		}{err: err, noFallback: result.FreshSessionFallback == nil}
	}()
	require.Eventually(t, func() bool {
		commands, err := os.ReadFile(filepath.Join(lifecycleDir, "commands"))
		return err == nil && strings.Contains(string(commands), "get_state")
	}, 3*time.Second, 10*time.Millisecond, "the fake Pi process must receive the preflight command before cancellation")
	cancel()
	outcome := <-outcomes
	require.Error(t, outcome.err)
	require.NotErrorIs(t, outcome.err, errRuntimeInputMutationNotAttempted)
	require.True(t, outcome.noFallback)
	assertPiFailedPreflightWasReaped(t, lifecycleDir)
}

func piResumedSessionFixture(t *testing.T, sessionID string) (string, string) {
	t.Helper()
	stateDir := t.TempDir()
	sessionDir := filepath.Join(stateDir, ".pi-sessions")
	require.NoError(t, os.MkdirAll(sessionDir, 0o700))
	require.NoError(t, os.WriteFile(
		filepath.Join(sessionDir, "2026_"+sessionID+".jsonl"),
		[]byte(`{"type":"session","id":"`+sessionID+`"}`+"\n"),
		0o600,
	))
	return stateDir, sessionID
}

func assertPiFailedPreflightWasReaped(t *testing.T, lifecycleDir string) {
	t.Helper()
	starts, err := os.ReadFile(filepath.Join(lifecycleDir, "starts"))
	require.NoError(t, err)
	require.Equal(t, 1, len(nonEmptyLines(string(starts))), "a failed preflight launches exactly one Pi process")

	commands, err := os.ReadFile(filepath.Join(lifecycleDir, "commands"))
	require.NoError(t, err)
	require.Equal(t, []string{"get_state"}, nonEmptyLines(string(commands)), "preflight failure must not write the initial prompt")

	active, err := filepath.Glob(filepath.Join(lifecycleDir, "active-*"))
	require.NoError(t, err)
	require.Empty(t, active, "failed preflight must reap the Pi process before Spawn returns")
	require.NoFileExists(t, filepath.Join(lifecycleDir, "overlap"))
}

func nonEmptyLines(value string) []string {
	var lines []string
	for _, line := range strings.Split(value, "\n") {
		if strings.TrimSpace(line) != "" {
			lines = append(lines, strings.TrimSpace(line))
		}
	}
	return lines
}