package retry

import (
	"testing"
)

func TestParseOverflow(t *testing.T) {
	tests := []struct {
		name    string
		msg     string
		wantNil bool
		wantIn  int
		wantMax int
		wantLim int
	}{
		{
			"standard format",
			"input length and `max_tokens` exceed context limit: 188059 + 20000 > 200000",
			false, 188059, 20000, 200000,
		},
		{
			"tight spacing",
			"input length and `max_tokens` exceed context limit: 100+50>150",
			false, 100, 50, 150,
		},
		{
			"not overflow",
			"prompt is too long: 137500 tokens > 135000 maximum",
			true, 0, 0, 0,
		},
		{
			"empty",
			"",
			true, 0, 0, 0,
		},
		{
			"partial match",
			"input length and `max_tokens` exceed context limit: abc + 20000 > 200000",
			true, 0, 0, 0,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			info := ParseOverflow(tt.msg)
			if tt.wantNil {
				if info != nil {
					t.Errorf("expected nil, got %+v", info)
				}
				return
			}
			if info == nil {
				t.Fatal("expected non-nil")
			}
			if info.InputTokens != tt.wantIn {
				t.Errorf("InputTokens = %d, want %d", info.InputTokens, tt.wantIn)
			}
			if info.MaxTokens != tt.wantMax {
				t.Errorf("MaxTokens = %d, want %d", info.MaxTokens, tt.wantMax)
			}
			if info.ContextLimit != tt.wantLim {
				t.Errorf("ContextLimit = %d, want %d", info.ContextLimit, tt.wantLim)
			}
		})
	}
}

func TestContextOverflowHandler_Adjust(t *testing.T) {
	h := &ContextOverflowHandler{
		FloorOutputTokens: 3000,
		SafetyBuffer:      1000,
	}

	tests := []struct {
		name    string
		info    *OverflowInfo
		wantOK  bool
		wantMin int
		wantMax int
	}{
		{
			"enough space",
			&OverflowInfo{InputTokens: 180000, MaxTokens: 20000, ContextLimit: 200000},
			true, 19000, 19000, // 200000 - 180000 - 1000 = 19000
		},
		{
			"tight space but above floor",
			&OverflowInfo{InputTokens: 195000, MaxTokens: 20000, ContextLimit: 200000},
			true, 4000, 4000, // 200000 - 195000 - 1000 = 4000
		},
		{
			"below floor",
			&OverflowInfo{InputTokens: 198000, MaxTokens: 20000, ContextLimit: 200000},
			false, 0, 0, // 200000 - 198000 - 1000 = 1000 < 3000
		},
		{
			"no space",
			&OverflowInfo{InputTokens: 200000, MaxTokens: 20000, ContextLimit: 200000},
			false, 0, 0,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			adj, ok := h.Adjust(tt.info)
			if ok != tt.wantOK {
				t.Errorf("ok = %v, want %v", ok, tt.wantOK)
			}
			if ok && (adj < tt.wantMin || adj > tt.wantMax) {
				t.Errorf("adjusted = %d, want %d-%d", adj, tt.wantMin, tt.wantMax)
			}
		})
	}
}

func TestContextOverflowHandler_ThinkingBudget(t *testing.T) {
	h := &ContextOverflowHandler{
		FloorOutputTokens: 3000,
		SafetyBuffer:      1000,
		ThinkingBudget:    50000,
	}

	// Available = 200000 - 180000 - 1000 = 19000
	// But thinking needs 50001 → not enough
	info := &OverflowInfo{InputTokens: 180000, MaxTokens: 20000, ContextLimit: 200000}
	_, ok := h.Adjust(info)
	if ok {
		t.Error("should fail when thinking budget exceeds available space")
	}

	// With less input, enough for thinking
	info2 := &OverflowInfo{InputTokens: 140000, MaxTokens: 20000, ContextLimit: 200000}
	adj, ok := h.Adjust(info2)
	if !ok {
		t.Fatal("should succeed with enough space for thinking")
	}
	// available = 200000 - 140000 - 1000 = 59000
	if adj != 59000 {
		t.Errorf("adjusted = %d, want 59000", adj)
	}
}