package runtime

import (
	"context"
	"testing"
	"time"
)

// These use a real question, never a greeting: a conversational message takes
// the smalltalk path and is capped well below its tier (see smalltalk.go), so
// "hi" here would assert the greeting cap while appearing to assert the tier's.
//
// A per-call cap is the difference between "the run may spend 120k tokens" and
// "any one answer may be 16k tokens long, eight times over". These assert the
// cap actually reaches the gateway, because it is inert until it does.

func TestOutputCapReachesGatewayPerTier(t *testing.T) {
	for _, tc := range []struct {
		reasoning string
		want      int64
	}{
		{"fast", 1_500},
		{"balanced", 3_000},
		{"deep", 4_000},
		{"nonsense-tier", 3_000}, // normalised to balanced, still capped
	} {
		t.Run(tc.reasoning, func(t *testing.T) {
			gw := &fakeGateway{text: "done"}
			exec := NewModelExecutor(gw, &MemorySink{}, nil)

			agent := testAgent()
			agent.Reasoning = tc.reasoning

			if _, err := exec.ExecuteAgent(context.Background(), agent, testInput("what happened today?")); err != nil {
				t.Fatal(err)
			}
			if got := gw.lastReq.MaxOutputTokens; got != tc.want {
				t.Errorf("MaxOutputTokens = %d, want %d", got, tc.want)
			}
		})
	}
}

// The cap comes from the budget, not from the agent's own tier. That is what
// makes a subagent inherit the parent's ceiling along with the parent's budget
// instead of reading its own tier and buying a longer answer.
func TestOutputCapComesFromTheBudgetNotTheTier(t *testing.T) {
	gw := &fakeGateway{text: "done"}
	exec := NewModelExecutor(gw, &MemorySink{}, nil)

	agent := testAgent()
	agent.Reasoning = "deep" // would be 4_000 if the tier decided

	limits := LimitsForTier("deep")
	limits.MaxOutputTokens = 777

	if _, err := exec.executeWithLimits(context.Background(), agent, testInput("what happened today?"), limits); err != nil {
		t.Fatal(err)
	}
	if got := gw.lastReq.MaxOutputTokens; got != 777 {
		t.Errorf("MaxOutputTokens = %d, want 777 from the supplied limits", got)
	}
}

// blockingSink is a database that has stopped answering. It returns only when
// its context ends, and records what ended it.
type blockingSink struct {
	ctxErr  error
	elapsed time.Duration
}

func (b *blockingSink) Save(ctx context.Context, _ *Trajectory) error {
	start := time.Now()
	<-ctx.Done()
	b.elapsed = time.Since(start)
	b.ctxErr = ctx.Err()
	return ctx.Err()
}

// A run that answered must not then wait on the sink indefinitely: to the person
// watching, a slow write is a slow run. I3 bounds the run; this bounds its tail.
func TestTrajectoryPersistenceIsBounded(t *testing.T) {
	restore := trajectoryPersistTimeout
	trajectoryPersistTimeout = 50 * time.Millisecond
	t.Cleanup(func() { trajectoryPersistTimeout = restore })

	sink := &blockingSink{}
	exec := NewModelExecutor(&fakeGateway{text: "answered"}, sink, nil)

	start := time.Now()
	res, err := exec.ExecuteAgent(context.Background(), testAgent(), testInput("hi"))
	total := time.Since(start)

	// The answer survives the sink failing — §6 keeps the two separate.
	if err != nil {
		t.Fatalf("a failed save must not fail the run: %v", err)
	}
	if res.Output != "answered" {
		t.Errorf("Output = %q, want the model's answer", res.Output)
	}
	if len(res.Unsaved) != 1 {
		t.Errorf("Unsaved = %v, want the one trajectory that could not be written", res.Unsaved)
	}

	if sink.ctxErr != context.DeadlineExceeded {
		t.Errorf("sink ctx ended with %v, want DeadlineExceeded — the write was not bounded", sink.ctxErr)
	}
	// Generous slack: the assertion is "bounded", not "fast".
	if total > 2*time.Second {
		t.Errorf("run took %v with a hung sink, want the persist timeout to cut it", total)
	}
}

// The bound must not become a cancellation: a run terminated by its own deadline
// is the one whose record matters most, so the write starts from a live context
// even when the caller's is already dead.
func TestTrajectoryPersistenceOutlivesACancelledCaller(t *testing.T) {
	sink := &MemorySink{}
	exec := NewModelExecutor(&fakeGateway{text: "answered", delay: 50 * time.Millisecond}, sink, nil)

	ctx, cancel := context.WithCancel(context.Background())
	// Cancelled while the model call is in flight: the run ends unhappily and
	// must still be recorded.
	go func() {
		time.Sleep(10 * time.Millisecond)
		cancel()
	}()

	res, _ := exec.ExecuteAgent(ctx, testAgent(), testInput("hi"))
	if res == nil {
		t.Fatal("no result")
	}
	if len(res.Unsaved) != 0 {
		t.Errorf("Unsaved = %v, want the trajectory written despite the cancelled caller", res.Unsaved)
	}
	if got := len(sink.Runs); got != 1 {
		t.Errorf("sink holds %d trajectories, want 1", got)
	}
}