package workbench

import (
	"context"
	"errors"
	"strings"
	"testing"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
)

type instructionStub struct {
	text string
	err  error
}

type disabledInstructionAgentStub struct{ interfaces.CustomAgentService }

func (disabledInstructionAgentStub) GetAgentByID(ctx context.Context, id string) (*types.CustomAgent, error) {
	tenant, _ := types.TenantIDFromContext(ctx)
	return &types.CustomAgent{ID: id, TenantID: tenant, Visibility: "global", Enabled: false, Config: types.CustomAgentConfig{ModelID: "model"}}, nil
}

func TestDisabledAgentCannotGenerateWithoutInstructionResolver(t *testing.T) {
	s, ctx, _, model := generationHarness(t)
	s.agents = disabledInstructionAgentStub{}
	if _, err := s.complete(ctx, "disabled", "system", "prompt"); !errors.Is(err, ErrInvalid) {
		t.Fatalf("expected disabled Agent rejection, got %v", err)
	}
	if len(model.messages) != 0 {
		t.Fatal("disabled Agent reached provider")
	}
}

func (s instructionStub) ResolveAgentInstructions(context.Context, *types.CustomAgent) (string, error) {
	return s.text, s.err
}

func TestMountedInstructionsParticipateInGenerationAndBudget(t *testing.T) {
	s, ctx, _, model := generationHarness(t)
	s.SetInstructionResolver(instructionStub{text: "技能规则：先列判断依据，再给案例。"})
	if _, err := s.complete(ctx, "agent", "system", "prompt"); err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(model.messages[0].Content, "技能规则：先列判断依据") {
		t.Fatal("mounted instruction missing from provider request")
	}
	s.models = generationModelStub{model: model, parameters: types.ModelParameters{ContextWindow: 4096, MaxOutputTokens: 512}}
	s.SetInstructionResolver(instructionStub{text: strings.Repeat("很长的技能规则", 3000)})
	model.messages = nil
	if _, err := s.complete(ctx, "agent", "system", "prompt"); !errors.Is(err, ErrInvalid) {
		t.Fatalf("expected budget rejection, got %v", err)
	}
	if len(model.messages) != 0 {
		t.Fatal("oversized skills reached model")
	}
}

func TestUnavailableInstructionStopsGeneration(t *testing.T) {
	s, ctx, _, model := generationHarness(t)
	s.SetInstructionResolver(instructionStub{err: ErrForbidden})
	if _, err := s.complete(ctx, "agent", "system", "prompt"); !errors.Is(err, ErrForbidden) {
		t.Fatalf("got %v", err)
	}
	if len(model.messages) != 0 {
		t.Fatal("generation silently ignored unavailable instruction")
	}
}