package service

import (
	"context"
	"io"
	"net/http"
	"strings"
	"testing"

	"github.com/go-resty/resty/v2"
	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/inference-service/internal/config"
)

type roundTripFunc func(*http.Request) (*http.Response, error)

func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) {
	return f(r)
}

func TestEscapeBedrockPathLabel(t *testing.T) {
	value := "arn:aws:bedrock:us-west-2:841162690310:application-inference-profile/5f1bs4ifgeef"
	require.Equal(t, "arn%3Aaws%3Abedrock%3Aus-west-2%3A841162690310%3Aapplication-inference-profile%2F5f1bs4ifgeef", escapeBedrockPathLabel(value))
}

func TestCohereInvoke_UsesBedrockInvokeModelIDAndEscapesPath(t *testing.T) {
	invokeModelID := "arn:aws:bedrock:us-west-2:841162690310:application-inference-profile/5f1bs4ifgeef"

	svc := &BedrockService{
		client:   resty.New(),
		apiToken: "test-token",
		region:   "us-west-2",
	}

	expectedPath := "/model/" + escapeBedrockPathLabel(invokeModelID) + "/invoke"
	svc.client.SetTransport(roundTripFunc(func(r *http.Request) (*http.Response, error) {
		require.Equal(t, "bedrock-runtime.us-west-2.amazonaws.com", r.URL.Host)
		require.Equal(t, expectedPath, r.URL.EscapedPath())

		body := `{"embeddings":[[0.1,0.2]],"usage":{"input_tokens":2}}`
		return &http.Response{
			StatusCode: http.StatusOK,
			Header: http.Header{
				"Content-Type":                     []string{"application/json"},
				"x-amzn-bedrock-input-token-count": []string{"2"},
			},
			Body:    io.NopCloser(strings.NewReader(body)),
			Request: r,
		}, nil
	}))

	params := &EmbeddingRequest{
		Model: "cohere/embed-english-v3",
		Texts: []string{"hello"},
		OtherOpts: map[string]any{
			"input_type": "search_query",
		},
	}
	modelConfig := &config.ModelConfig{
		BedrockModelID:       "cohere.embed-english-v3",
		BedrockInvokeModelID: invokeModelID,
		MicroDollarPerKToken: 100,
		AllowedOpts:          []string{"input_type", "truncate"},
	}

	embeddings, costMicroDollars, err := svc.generateCohereEmbeddings(context.Background(), params, modelConfig)
	require.NoError(t, err)
	require.Equal(t, [][]float32{{0.1, 0.2}}, embeddings)
	require.Equal(t, int64(0), costMicroDollars)
}