package llm

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/yeomyeonggeori/blueclaw/internal/config"
	"github.com/yeomyeonggeori/blueprotocol/model"
)

func endpointConfiguration(url string, modelName string) config.ModelEndpointConfiguration {
	return config.ModelEndpointConfiguration{Endpoint: url, Model: modelName}
}

func configurationWithOneEndpointPerTier(url string, modelName string) config.RuntimeConfiguration {
	tiers := map[string][]config.ModelEndpointConfiguration{}
	for _, modelTier := range ModelTiers {
		tiers[modelTier] = []config.ModelEndpointConfiguration{endpointConfiguration(url, modelName)}
	}
	return config.RuntimeConfiguration{LanguageModel: config.LanguageModelConfiguration{Tiers: tiers}}
}

func TestTierProviderFactoryBuildsOneProviderPerConfiguredTier(t *testing.T) {
	tierProviderFactory, errorValue := NewTierProviderFactory(configurationWithOneEndpointPerTier("http://127.0.0.1:9/v1", "example/model"))
	if errorValue != nil {
		t.Fatal(errorValue)
	}
	for _, modelTier := range ModelTiers {
		tierProvider, tierError := tierProviderFactory(modelTier)
		if tierError != nil {
			t.Fatalf("%s: %v", modelTier, tierError)
		}
		if tierProvider.Provider == nil {
			t.Fatalf("%s reached no provider", modelTier)
		}
		if tierProvider.Reaches != "http://127.0.0.1:9/v1 example/model" {
			t.Fatalf("%s must report the endpoint it reaches, got %q", modelTier, tierProvider.Reaches)
		}
	}
}

func TestTierProviderFactoryRefusesATierItWasGivenNoEndpointFor(t *testing.T) {
	runtimeConfiguration := configurationWithOneEndpointPerTier("http://127.0.0.1:9/v1", "example/model")
	delete(runtimeConfiguration.LanguageModel.Tiers, "high")
	tierProviderFactory, errorValue := NewTierProviderFactory(runtimeConfiguration)
	if errorValue != nil {
		t.Fatal(errorValue)
	}
	if _, tierError := tierProviderFactory("high"); tierError == nil {
		t.Fatal("a tier with no endpoint must be refused rather than silently answered by another tier")
	}
}

func TestATierWithSeveralEndpointsFallsThroughThemInOrder(t *testing.T) {
	runtimeConfiguration := configurationWithOneEndpointPerTier("http://127.0.0.1:9/v1", "example/first")
	runtimeConfiguration.LanguageModel.Tiers["medium"] = []config.ModelEndpointConfiguration{
		endpointConfiguration("http://127.0.0.1:9/v1", "example/first"),
		endpointConfiguration("http://127.0.0.1:9/v1", "example/second"),
	}
	tierProviderFactory, errorValue := NewTierProviderFactory(runtimeConfiguration)
	if errorValue != nil {
		t.Fatal(errorValue)
	}
	tierProvider, errorValue := tierProviderFactory("medium")
	if errorValue != nil {
		t.Fatal(errorValue)
	}
	fallbackProvider, isFallbackProvider := tierProvider.Provider.(FallbackLanguageModelProvider)
	if !isFallbackProvider {
		t.Fatalf("a tier given two endpoints must try the second when the first fails, got %T", tierProvider.Provider)
	}
	if _, wouldChainAgain := fallbackProvider.FallbackProvider.(FallbackLanguageModelProvider); wouldChainAgain {
		t.Fatal("two endpoints make one fallback, not two")
	}
	if tierProvider.Reaches != "http://127.0.0.1:9/v1 example/first, http://127.0.0.1:9/v1 example/second" {
		t.Fatalf("the tier must report both endpoints in order, got %q", tierProvider.Reaches)
	}
}

func TestTierProviderFactoryReadsTheKeyOutOfTheFileItWasGiven(t *testing.T) {
	apiKeyPath := filepath.Join(t.TempDir(), "model-key")
	if errorValue := os.WriteFile(apiKeyPath, []byte("  a-key\n"), 0o600); errorValue != nil {
		t.Fatal(errorValue)
	}
	runtimeConfiguration := configurationWithOneEndpointPerTier("http://127.0.0.1:9/v1", "example/model")
	runtimeConfiguration.LanguageModel.Tiers["low"] = []config.ModelEndpointConfiguration{{
		Endpoint: "http://127.0.0.1:9/v1", Model: "example/model", APIKeyPath: apiKeyPath,
	}}
	tierProviderFactory, errorValue := NewTierProviderFactory(runtimeConfiguration)
	if errorValue != nil {
		t.Fatal(errorValue)
	}
	if _, tierError := tierProviderFactory("low"); tierError != nil {
		t.Fatal(tierError)
	}
}

func TestTierProviderFactoryRefusesAKeyFileItCannotRead(t *testing.T) {
	runtimeConfiguration := configurationWithOneEndpointPerTier("http://127.0.0.1:9/v1", "example/model")
	runtimeConfiguration.LanguageModel.Tiers["low"] = []config.ModelEndpointConfiguration{{
		Endpoint: "http://127.0.0.1:9/v1", Model: "example/model", APIKeyPath: filepath.Join(t.TempDir(), "absent"),
	}}
	tierProviderFactory, errorValue := NewTierProviderFactory(runtimeConfiguration)
	if errorValue != nil {
		t.Fatal(errorValue)
	}
	if _, tierError := tierProviderFactory("low"); tierError == nil {
		t.Fatal("an endpoint whose key file is missing must be refused rather than asked without a key")
	}
}

func TestAnEndpointAsksWithWhateverKeyItsFileHoldsWhenTheRequestIsSent(t *testing.T) {
	authorizations := []string{}
	server := httptest.NewServer(http.HandlerFunc(func(responseWriter http.ResponseWriter, request *http.Request) {
		authorizations = append(authorizations, request.Header.Get("Authorization"))
		writeCompletion(responseWriter, schemaAnswer(`{"question":"done"}`), "tool_calls")
	}))
	defer server.Close()
	apiKeyPath := filepath.Join(t.TempDir(), "host-session")
	writeKeyFile(t, apiKeyPath, "first-session")
	provider, errorValue := endpointProvider(config.ModelEndpointConfiguration{Endpoint: server.URL + "/v1", Model: "example/model", APIKeyPath: apiKeyPath})
	if errorValue != nil {
		t.Fatal(errorValue)
	}
	askOnce(t, provider)
	writeKeyFile(t, apiKeyPath, "renewed-session")
	askOnce(t, provider)
	if strings.Join(authorizations, ", ") != "Bearer first-session, Bearer renewed-session" {
		t.Fatalf("a key file renewed after the provider was built must be the key the next request carries; got %v", authorizations)
	}
}

func TestADecisionEndpointAsksWithWhateverKeyItsFileHoldsWhenTheRequestIsSent(t *testing.T) {
	authorizations := []string{}
	server := httptest.NewServer(http.HandlerFunc(func(responseWriter http.ResponseWriter, request *http.Request) {
		authorizations = append(authorizations, request.Header.Get("Authorization"))
		_, _ = responseWriter.Write([]byte(`{"answers":{"route":{"choice":"start_task","probabilities":{"start_task":0.9,"answer_question":0.1}}}}`))
	}))
	defer server.Close()
	apiKeyPath := filepath.Join(t.TempDir(), "host-session")
	writeKeyFile(t, apiKeyPath, "first-session")
	decisionModel, errorValue := NewConfiguredDecisionModel(config.RuntimeConfiguration{LanguageModel: config.LanguageModelConfiguration{
		Decision: config.ModelEndpointConfiguration{Endpoint: server.URL + "/api/alpha/decisions", Model: "kev-latest", APIKeyPath: apiKeyPath},
	}})
	if errorValue != nil {
		t.Fatal(errorValue)
	}
	request := model.DecisionRequest{
		State:     "a message",
		Questions: map[string]model.DecisionQuestion{"route": model.ChoiceQuestion{Instructions: "What now?", OptionDescriptions: map[string]string{"start_task": "work", "answer_question": "words"}}.Question()},
	}
	if _, errorValue := decisionModel.Decide(context.Background(), request); errorValue != nil {
		t.Fatal(errorValue)
	}
	writeKeyFile(t, apiKeyPath, "renewed-session")
	if _, errorValue := decisionModel.Decide(context.Background(), request); errorValue != nil {
		t.Fatal(errorValue)
	}
	if strings.Join(authorizations, ", ") != "Bearer first-session, Bearer renewed-session" {
		t.Fatalf("a key file renewed after the decision model was built must be the key the next decision carries; got %v", authorizations)
	}
}

func writeKeyFile(t *testing.T, path string, key string) {
	t.Helper()
	if errorValue := os.WriteFile(path, []byte(key+"\n"), 0o600); errorValue != nil {
		t.Fatal(errorValue)
	}
}

func askOnce(t *testing.T, provider LanguageModelProvider) {
	t.Helper()
	if _, errorValue := provider.GenerateResponse(context.Background(), "hello"); errorValue != nil {
		t.Fatal(errorValue)
	}
}

func capabilityLadderConfiguration(modelName string) config.RuntimeConfiguration {
	return config.RuntimeConfiguration{LanguageModel: config.LanguageModelConfiguration{
		Capability: config.LanguageModelCapabilityConfiguration{
			ExecutionMode: "auto",
			XLowModel:     modelName,
			LowModel:      modelName,
			MediumModel:   modelName,
			HighModel:     modelName,
			XHighModel:    modelName,
			MaxModel:      modelName,
		},
	}}
}

func TestACapabilityLadderReachesEveryTierThroughTheCapabilityRoute(t *testing.T) {
	tierProviderFactory, errorValue := NewTierProviderFactory(capabilityLadderConfiguration("example/model"))
	if errorValue != nil {
		t.Fatal(errorValue)
	}
	for _, modelTier := range ModelTiers {
		tierProvider, tierError := tierProviderFactory(modelTier)
		if tierError != nil {
			t.Fatalf("%s: %v", modelTier, tierError)
		}
		capabilityClient, isCapabilityClient := tierProvider.Provider.(CapabilityLLMClient)
		if !isCapabilityClient {
			t.Fatalf("%s must reach the capability route, got %T", modelTier, tierProvider.Provider)
		}
		if capabilityClient.ModelName != "example/model" || capabilityClient.ExecutionMode != "auto" {
			t.Fatalf("%s lost the model or the execution mode it was configured with: %+v", modelTier, capabilityClient)
		}
		if capabilityClient.ModelTier != modelTier {
			t.Fatalf("%s must tell capabilityd which tier is asking, got %q", modelTier, capabilityClient.ModelTier)
		}
	}
}

func TestACapabilityLadderMissingATierIsRefused(t *testing.T) {
	runtimeConfiguration := capabilityLadderConfiguration("example/model")
	runtimeConfiguration.LanguageModel.Capability.MaxModel = ""
	tierProviderFactory, errorValue := NewTierProviderFactory(runtimeConfiguration)
	if errorValue != nil {
		t.Fatal(errorValue)
	}
	if _, tierError := tierProviderFactory("max"); tierError == nil {
		t.Fatal("a capability tier with no model must be refused")
	}
}

func TestAConfigurationNamingNeitherShapeIsRefused(t *testing.T) {
	if _, errorValue := NewTierProviderFactory(config.RuntimeConfiguration{}); errorValue == nil {
		t.Fatal("a document naming no endpoints and no capability route must be refused")
	}
}

func TestAConfigurationNamingBothShapesIsRefused(t *testing.T) {
	runtimeConfiguration := configurationWithOneEndpointPerTier("http://127.0.0.1:9/v1", "example/model")
	runtimeConfiguration.LanguageModel.Capability = capabilityLadderConfiguration("example/model").LanguageModel.Capability
	if _, errorValue := NewTierProviderFactory(runtimeConfiguration); errorValue == nil {
		t.Fatal("a document naming both endpoints and the capability route is ambiguous and must be refused")
	}
}

func TestEmbeddingProviderFollowsTheEndpointWhenOneIsNamed(t *testing.T) {
	runtimeConfiguration := config.RuntimeConfiguration{LanguageModel: config.LanguageModelConfiguration{
		Embedding: config.ModelEndpointConfiguration{Endpoint: "http://127.0.0.1:9/v1", Model: "example/embedding"},
	}}
	embeddingProvider, errorValue := NewConfiguredEmbeddingProvider(runtimeConfiguration)
	if errorValue != nil {
		t.Fatal(errorValue)
	}
	if _, isCapabilityClient := embeddingProvider.(CapabilityEmbeddingClient); isCapabilityClient {
		t.Fatal("an embedding endpoint is reached directly, not through the capability route")
	}
}

func TestEmbeddingProviderTakesTheCapabilityRouteWhenOnlyAModelIsNamed(t *testing.T) {
	runtimeConfiguration := config.RuntimeConfiguration{LanguageModel: config.LanguageModelConfiguration{
		Embedding:  config.ModelEndpointConfiguration{Model: "example/embedding"},
		Capability: config.LanguageModelCapabilityConfiguration{ExecutionMode: "auto"},
	}}
	embeddingProvider, errorValue := NewConfiguredEmbeddingProvider(runtimeConfiguration)
	if errorValue != nil {
		t.Fatal(errorValue)
	}
	capabilityClient, isCapabilityClient := embeddingProvider.(CapabilityEmbeddingClient)
	if !isCapabilityClient {
		t.Fatalf("an embedding model with no endpoint is asked for through the capability route, got %T", embeddingProvider)
	}
	if capabilityClient.ModelName != "example/embedding" {
		t.Fatalf("the embedding model must be the one configured, got %q", capabilityClient.ModelName)
	}
	if capabilityClient.OutputDimensions != DefaultEmbeddingDimensions {
		t.Fatalf("skill retrieval must ask for the same width as memory, got %d", capabilityClient.OutputDimensions)
	}
}

func TestEmbeddingModelFallsBackToTheMemoryModelThenTheDefault(t *testing.T) {
	if name := ConfiguredEmbeddingModelName(config.RuntimeConfiguration{}); name != DefaultEmbeddingModelName {
		t.Fatalf("expected the default, got %q", name)
	}
	memoryOnly := config.RuntimeConfiguration{Memory: config.MemoryConfiguration{EmbeddingModel: "example/memory"}}
	if name := ConfiguredEmbeddingModelName(memoryOnly); name != "example/memory" {
		t.Fatalf("expected the memory model, got %q", name)
	}
	memoryOnly.LanguageModel.Embedding.Model = "example/language"
	if name := ConfiguredEmbeddingModelName(memoryOnly); name != "example/language" {
		t.Fatalf("expected the language model entry to win, got %q", name)
	}
}

func TestStandaloneEmbedderAsksTheEndpointWithTemplatesAndFailsOnAnotherWidth(t *testing.T) {
	inputs := []string{}
	vectorWidth := 3
	server := httptest.NewServer(http.HandlerFunc(func(responseWriter http.ResponseWriter, request *http.Request) {
		var received map[string]any
		_ = json.NewDecoder(request.Body).Decode(&received)
		inputs = append(inputs, received["input"].(string))
		_ = json.NewEncoder(responseWriter).Encode(map[string]any{"data": []map[string]any{{"embedding": make([]float64, vectorWidth)}}})
	}))
	defer server.Close()
	runtimeConfiguration := config.RuntimeConfiguration{
		LanguageModel: config.LanguageModelConfiguration{Embedding: config.ModelEndpointConfiguration{Endpoint: server.URL + "/v1", Model: "google/embeddinggemma-2"}},
		Memory:        config.MemoryConfiguration{EmbeddingDimensions: 3},
	}
	embedder, errorValue := NewConfiguredEmbeddingProvider(runtimeConfiguration)
	if errorValue != nil {
		t.Fatal(errorValue)
	}

	if _, errorValue := embedder.EmbedQuery(context.Background(), "slides"); errorValue != nil {
		t.Fatal(errorValue)
	}
	if _, errorValue := embedder.EmbedDocuments(context.Background(), []string{"notes"}); errorValue != nil {
		t.Fatal(errorValue)
	}
	if embedder.EmbeddingModelName() != "google/embeddinggemma-2" {
		t.Fatalf("rows must record the configured model, got %q", embedder.EmbeddingModelName())
	}
	if len(inputs) != 2 || inputs[0] != "task: search result | query: slides" || inputs[1] != "title: none | text: notes" {
		t.Fatalf("expected the EmbeddingGemma templates, got %v", inputs)
	}
	vectorWidth = 4
	if _, errorValue := embedder.EmbedDocuments(context.Background(), []string{"notes"}); errorValue == nil {
		t.Fatal("a vector of another width than configured must be refused, not stored")
	}
}

func TestDecisionModelAsksTheEndpointWhenOneIsNamed(t *testing.T) {
	var authorization, modelName string
	decisionServer := httptest.NewServer(http.HandlerFunc(func(responseWriter http.ResponseWriter, request *http.Request) {
		var requestDocument struct {
			Model string `json:"model"`
		}
		_ = json.NewDecoder(request.Body).Decode(&requestDocument)
		authorization, modelName = request.Header.Get("Authorization"), requestDocument.Model
		_, _ = responseWriter.Write([]byte(`{"answers":{"route":{"choice":"start_task","probabilities":{"start_task":0.9,"answer_question":0.1}}}}`))
	}))
	defer decisionServer.Close()
	apiKeyPath := filepath.Join(t.TempDir(), "decision-key")
	if errorValue := os.WriteFile(apiKeyPath, []byte("decision-key\n"), 0o600); errorValue != nil {
		t.Fatal(errorValue)
	}
	runtimeConfiguration := config.RuntimeConfiguration{LanguageModel: config.LanguageModelConfiguration{
		Decision: config.ModelEndpointConfiguration{Endpoint: decisionServer.URL + "/v1/systemone", Model: "kev-latest", APIKeyPath: apiKeyPath},
	}}
	decisionModel, errorValue := NewConfiguredDecisionModel(runtimeConfiguration)
	if errorValue != nil {
		t.Fatal(errorValue)
	}
	response, errorValue := decisionModel.Decide(context.Background(), model.DecisionRequest{
		State:     "a message",
		Questions: map[string]model.DecisionQuestion{"route": model.ChoiceQuestion{Instructions: "What now?", OptionDescriptions: map[string]string{"start_task": "work", "answer_question": "words"}}.Question()},
	})
	if errorValue != nil {
		t.Fatal(errorValue)
	}
	if authorization != "Bearer decision-key" || modelName != "kev-latest" || response.Answers["route"].Choice != "start_task" {
		t.Fatalf("expected the configured endpoint, key and model to answer; got authorization %q, model %q, answers %+v", authorization, modelName, response.Answers)
	}
}

func TestDecisionModelTakesTheCapabilityRouteWhenNoEndpointIsNamed(t *testing.T) {
	runtimeConfiguration := config.RuntimeConfiguration{LanguageModel: config.LanguageModelConfiguration{
		Capability: config.LanguageModelCapabilityConfiguration{DecisionModel: "example/decision", ExecutionMode: "auto"},
	}}
	decisionModel, errorValue := NewConfiguredDecisionModel(runtimeConfiguration)
	if errorValue != nil {
		t.Fatal(errorValue)
	}
	if _, isCapabilityClient := decisionModel.(CapabilityDecisionClient); !isCapabilityClient {
		t.Fatalf("a decision model with no endpoint is asked for through the capability route, got %T", decisionModel)
	}
}

func TestDecisionEndpointWithoutAModelIsRefused(t *testing.T) {
	runtimeConfiguration := config.RuntimeConfiguration{LanguageModel: config.LanguageModelConfiguration{
		Decision: config.ModelEndpointConfiguration{Endpoint: "http://127.0.0.1:9/v1/systemone"},
	}}
	if _, errorValue := NewConfiguredDecisionModel(runtimeConfiguration); errorValue == nil {
		t.Fatal("a decision endpoint that names no model must be refused")
	}
}

func TestAnEndpointReadsItsKeyFromTheNamedEnvironmentVariable(t *testing.T) {
	t.Setenv("EXAMPLE_MODEL_KEY", "a-key")
	apiKey, errorValue := endpointAPIKey(config.ModelEndpointConfiguration{Endpoint: "https://example.com/v1", APIKeyEnvironment: "EXAMPLE_MODEL_KEY"})
	if errorValue != nil || apiKey != "a-key" {
		t.Fatalf("expected the key from EXAMPLE_MODEL_KEY, got %q, %v", apiKey, errorValue)
	}
}

func TestAnEndpointRefusesAnUnsetKeyVariableAndASecondKeySource(t *testing.T) {
	t.Setenv("EXAMPLE_MODEL_KEY", "a-key")
	if _, errorValue := endpointAPIKey(config.ModelEndpointConfiguration{Endpoint: "https://example.com/v1", APIKeyEnvironment: "EXAMPLE_UNSET_KEY"}); errorValue == nil || !strings.Contains(errorValue.Error(), "EXAMPLE_UNSET_KEY") {
		t.Fatalf("an unset key variable must be refused by name, got %v", errorValue)
	}
	if _, errorValue := endpointAPIKey(config.ModelEndpointConfiguration{Endpoint: "https://example.com/v1", APIKeyEnvironment: "EXAMPLE_MODEL_KEY", APIKeyPath: "/tmp/key"}); errorValue == nil {
		t.Fatal("an endpoint naming both a key file and a key variable must be refused")
	}
}

func TestVisualReviewModelAsksTheModelItsEndpointNames(t *testing.T) {
	var modelName string
	reviewServer := httptest.NewServer(http.HandlerFunc(func(responseWriter http.ResponseWriter, request *http.Request) {
		var requestDocument struct {
			Model string `json:"model"`
		}
		_ = json.NewDecoder(request.Body).Decode(&requestDocument)
		modelName = requestDocument.Model
		_, _ = responseWriter.Write([]byte(`{"answers":{"visual_defect":{"choice":"none","probabilities":{"none":0.9,"crowded":0.1}}}}`))
	}))
	defer reviewServer.Close()
	apiKeyPath := filepath.Join(t.TempDir(), "review-key")
	if errorValue := os.WriteFile(apiKeyPath, []byte("review-key\n"), 0o600); errorValue != nil {
		t.Fatal(errorValue)
	}
	runtimeConfiguration := config.RuntimeConfiguration{LanguageModel: config.LanguageModelConfiguration{
		Decision:     config.ModelEndpointConfiguration{Endpoint: reviewServer.URL + "/decisions", Model: "decision-model", APIKeyPath: apiKeyPath},
		VisualReview: config.ModelEndpointConfiguration{Endpoint: reviewServer.URL + "/decisions", Model: "vision-decision-model", APIKeyPath: apiKeyPath},
	}}
	reviewModel, errorValue := NewConfiguredVisualReviewModel(runtimeConfiguration)
	if errorValue != nil {
		t.Fatal(errorValue)
	}
	if _, errorValue := reviewModel.Decide(context.Background(), model.DecisionRequest{
		State:     "a slide",
		Questions: map[string]model.DecisionQuestion{"visual_defect": model.ChoiceQuestion{Instructions: "Which defect?", OptionDescriptions: map[string]string{"none": "clean", "crowded": "packed"}}.Question()},
	}); errorValue != nil {
		t.Fatal(errorValue)
	}
	if modelName != "vision-decision-model" {
		t.Fatalf("the visual review asks the model its own endpoint names, got %q", modelName)
	}
}

func TestNoVisualReviewEndpointMeansNoVisualReview(t *testing.T) {
	reviewModel, errorValue := NewConfiguredVisualReviewModel(config.RuntimeConfiguration{})
	if errorValue != nil || reviewModel != nil {
		t.Fatalf("an unconfigured visual review is off, got %v, %v", reviewModel, errorValue)
	}
}