package main

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"reflect"
	"slices"
	"strings"
	"testing"
	"time"
)

func TestDurationFromEnvironment(t *testing.T) {
	t.Setenv("TEST_INTERVAL", "")
	got, err := durationFromEnvironment("TEST_INTERVAL", "6h")
	if err != nil {
		t.Fatal(err)
	}
	if got != 6*time.Hour {
		t.Fatalf("durationFromEnvironment() = %s, want 6h", got)
	}

	t.Setenv("TEST_INTERVAL", "45m")
	got, err = durationFromEnvironment("TEST_INTERVAL", "6h")
	if err != nil {
		t.Fatal(err)
	}
	if got != 45*time.Minute {
		t.Fatalf("durationFromEnvironment() = %s, want 45m", got)
	}
}

func TestDurationFromEnvironmentRejectsNonPositive(t *testing.T) {
	t.Setenv("TEST_INTERVAL", "0s")
	if _, err := durationFromEnvironment("TEST_INTERVAL", "6h"); err == nil {
		t.Fatal("durationFromEnvironment() accepted a non-positive duration")
	}
}

func TestDurationFromEnvironmentRejectsInvalidDuration(t *testing.T) {
	t.Setenv("TEST_INTERVAL", "eventually")
	if _, err := durationFromEnvironment("TEST_INTERVAL", "6h"); err == nil {
		t.Fatal("durationFromEnvironment() accepted an invalid duration")
	}
}

func TestProvidersFromEnvironment(t *testing.T) {
	t.Setenv("ENABLED_PROVIDERS", "")
	providers, err := providersFromEnvironment()
	if err != nil {
		t.Fatal(err)
	}
	got := mapsKeys(providers)
	slices.Sort(got)
	if want := []string{"google-vertex", "novita"}; !reflect.DeepEqual(got, want) {
		t.Fatalf("default providers = %v, want %v", got, want)
	}

	t.Setenv("ENABLED_PROVIDERS", "openai, google-vertex")
	providers, err = providersFromEnvironment()
	if err != nil {
		t.Fatal(err)
	}
	got = mapsKeys(providers)
	slices.Sort(got)
	if want := []string{"google-vertex", "openai"}; !reflect.DeepEqual(got, want) {
		t.Fatalf("configured providers = %v, want %v", got, want)
	}
}

func TestProvidersFromEnvironmentRejectsInvalidSelections(t *testing.T) {
	for _, value := range []string{"unknown", "novita,novita", ","} {
		t.Run(value, func(t *testing.T) {
			t.Setenv("ENABLED_PROVIDERS", value)
			if _, err := providersFromEnvironment(); err == nil {
				t.Fatalf("providersFromEnvironment() accepted %q", value)
			}
		})
	}
}

func mapsKeys[T any](input map[string]T) []string {
	keys := make([]string, 0, len(input))
	for key := range input {
		keys = append(keys, key)
	}
	return keys
}

func selectProviders(names ...string) providerSelection {
	definitions := buildProviderDefinitions()
	providers := providerSelection{}
	for _, name := range names {
		providers[name] = definitions[name]
	}
	return providers
}

func number(value string) *json.Number {
	number := json.Number(value)
	return &number
}

func TestConvertModelsDev(t *testing.T) {
	provider := modelsDevProvider{Models: map[string]modelsDevModel{
		"model-a": {Cost: modelsDevCost{
			Input:     number("1.5"),
			Output:    number("7.5"),
			CacheRead: number("0.15"),
			Tiers: []modelsDevTier{{
				Input:  number("3"),
				Output: number("15"),
				Tier:   modelsDevTierSelector{Type: "context", Size: 200000},
			}},
		}},
		"unpriced": {},
	}}
	want := map[string]catalogModel{
		"model-a": {
			Rates: map[string]string{"input": "1.5", "output": "7.5", "cacheRead": "0.15"},
			Tiers: []catalogTier{{ContextOver: 200000, Rates: map[string]string{"input": "3", "output": "15"}}},
		},
	}
	if got := convertModelsDev(provider); !reflect.DeepEqual(got, want) {
		t.Fatalf("convertModelsDev() = %#v, want %#v", got, want)
	}
}

func TestConvertNovita(t *testing.T) {
	models, err := convertNovita(novitaCatalog{Data: []novitaModel{{
		ID:          "model-a",
		InputPrice:  json.Number("3900"),
		OutputPrice: json.Number("12500"),
	}}})
	if err != nil {
		t.Fatal(err)
	}
	want := map[string]catalogModel{
		"model-a": {Rates: map[string]string{"input": "0.39", "output": "1.25"}},
	}
	if !reflect.DeepEqual(models, want) {
		t.Fatalf("convertNovita() = %#v, want %#v", models, want)
	}
}

func TestConvertNovitaSkipsIncompleteModels(t *testing.T) {
	models, err := convertNovita(novitaCatalog{Data: []novitaModel{
		{ID: "", InputPrice: json.Number("10000"), OutputPrice: json.Number("20000")},
		{ID: "missing-input", OutputPrice: json.Number("20000")},
		{ID: "missing-output", InputPrice: json.Number("10000")},
	}})
	if err != nil {
		t.Fatal(err)
	}
	if len(models) != 0 {
		t.Fatalf("convertNovita() = %#v, want no models", models)
	}
}

func TestBuildCatalogUsesAgentgatewayProviderIDs(t *testing.T) {
	input := modelsDevCatalog{
		"openai": {Models: map[string]modelsDevModel{"openai-model": {Cost: modelsDevCost{Input: number("1")}}}},
		"google-vertex": {Models: map[string]modelsDevModel{
			"vertex-model": {Cost: modelsDevCost{Input: number("2")}},
		}},
	}
	output, err := buildCatalog(input, novitaCatalog{Data: []novitaModel{{
		ID: "novita-model", InputPrice: json.Number("10000"), OutputPrice: json.Number("20000"),
	}}}, selectProviders("openai", "google-vertex", "novita"))
	if err != nil {
		t.Fatal(err)
	}
	want := []string{"custom", "gcp.vertex_ai", "openai"}
	got := make([]string, 0, len(output.Providers))
	for provider := range output.Providers {
		got = append(got, provider)
	}
	slices.Sort(got)
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("provider IDs = %v, want %v", got, want)
	}
}

func TestBuildCatalogRejectsMissingOrInvalidSources(t *testing.T) {
	valid := modelsDevCatalog{
		"google-vertex": {Models: map[string]modelsDevModel{
			"vertex-model": {Cost: modelsDevCost{Input: number("2")}},
		}},
	}
	tests := []struct {
		name   string
		input  modelsDevCatalog
		novita novitaCatalog
	}{
		{
			name:  "missing models.dev provider",
			input: modelsDevCatalog{},
			novita: novitaCatalog{Data: []novitaModel{{
				ID: "novita-model", InputPrice: json.Number("10000"), OutputPrice: json.Number("20000"),
			}}},
		},
		{
			name:  "missing Novita models",
			input: valid,
		},
		{
			name:  "invalid Novita price",
			input: valid,
			novita: novitaCatalog{Data: []novitaModel{{
				ID: "novita-model", InputPrice: json.Number("-1"), OutputPrice: json.Number("20000"),
			}}},
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			if _, err := buildCatalog(
				test.input,
				test.novita,
				selectProviders("google-vertex", "novita"),
			); err == nil {
				t.Fatal("buildCatalog() unexpectedly succeeded")
			}
		})
	}
}

func TestTenThousandths(t *testing.T) {
	tests := []struct {
		input string
		want  string
	}{
		{input: "0", want: "0"},
		{input: "1", want: "0.0001"},
		{input: "3900", want: "0.39"},
		{input: "10000", want: "1"},
		{input: "12500", want: "1.25"},
		{input: "00012500", want: "1.25"},
	}
	for _, test := range tests {
		t.Run(test.input, func(t *testing.T) {
			got, err := tenThousandths(json.Number(test.input))
			if err != nil {
				t.Fatal(err)
			}
			if got != test.want {
				t.Fatalf("tenThousandths(%q) = %q, want %q", test.input, got, test.want)
			}
		})
	}
}

func TestTenThousandthsRejectsInvalidNumbers(t *testing.T) {
	for _, input := range []string{"0.39", "-1", "1e4"} {
		t.Run(input, func(t *testing.T) {
			if _, err := tenThousandths(json.Number(input)); err == nil {
				t.Fatalf("tenThousandths(%q) unexpectedly succeeded", input)
			}
		})
	}
}

func TestRenderMetrics(t *testing.T) {
	input := catalog{Providers: map[string]catalogProvider{
		"custom": {Models: map[string]catalogModel{
			"vendor/model": {Rates: map[string]string{"output": "1.25", "cacheRead": "0.1"}},
		}},
	}}
	want := "# HELP agentgateway_model_token_price_usd_per_million Model token price in USD per million tokens.\n" +
		"# TYPE agentgateway_model_token_price_usd_per_million gauge\n" +
		"agentgateway_model_token_price_usd_per_million{gen_ai_system=\"custom\",gen_ai_request_model=\"vendor/model\",gen_ai_token_type=\"cache_read\",source_provider=\"novita\"} 0.1\n" +
		"agentgateway_model_token_price_usd_per_million{gen_ai_system=\"custom\",gen_ai_request_model=\"vendor/model\",gen_ai_token_type=\"output\",source_provider=\"novita\"} 1.25\n"
	if got := string(renderMetrics(input)); got != want {
		t.Fatalf("renderMetrics() = %q, want %q", got, want)
	}
}

func TestRenderModelList(t *testing.T) {
	input := catalog{Providers: map[string]catalogProvider{
		"gcp.vertex_ai": {Models: map[string]catalogModel{"gemini-z": {}, "gemini-a": {}}},
		"custom":        {Models: map[string]catalogModel{"vendor/model": {}}},
	}}
	want := modelList{Object: "list", Data: []modelListItem{
		{ID: "novita/vendor/model", Object: "model", Created: 0, OwnedBy: "novita"},
		{ID: "vertex_ai/gemini-a", Object: "model", Created: 0, OwnedBy: "google-vertex"},
		{ID: "vertex_ai/gemini-z", Object: "model", Created: 0, OwnedBy: "google-vertex"},
	}}
	if got := renderModelList(input); !reflect.DeepEqual(got, want) {
		t.Fatalf("renderModelList() = %#v, want %#v", got, want)
	}
}

func TestReadCatalog(t *testing.T) {
	path := filepath.Join(t.TempDir(), "catalog.json")
	want := catalog{Providers: map[string]catalogProvider{
		"gcp.vertex_ai": {Models: map[string]catalogModel{"gemini": {Rates: map[string]string{"input": "1"}}}},
	}}
	contents, err := json.Marshal(want)
	if err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(path, contents, 0o600); err != nil {
		t.Fatal(err)
	}
	got, err := readCatalog(path)
	if err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("readCatalog() = %#v, want %#v", got, want)
	}

	if err := os.WriteFile(path, []byte("{"), 0o600); err != nil {
		t.Fatal(err)
	}
	if _, err := readCatalog(path); err == nil {
		t.Fatal("readCatalog() accepted malformed JSON")
	}
	if _, err := readCatalog(filepath.Join(t.TempDir(), "missing")); err == nil {
		t.Fatal("readCatalog() accepted a missing file")
	}
}

func TestCatalogHandler(t *testing.T) {
	path := filepath.Join(t.TempDir(), "catalog.json")
	contents := `{"providers":{
		"gcp.vertex_ai":{"models":{"gemini":{"rates":{"input":"2"}}}},
		"custom":{"models":{"novita/model":{"rates":{"input":"3"}}}}
	}}`
	if err := os.WriteFile(path, []byte(contents), 0o600); err != nil {
		t.Fatal(err)
	}
	handler := catalogHandler(path, selectProviders("google-vertex", "novita"))
	tests := []struct {
		name        string
		method      string
		path        string
		status      int
		contentType string
		body        string
	}{
		{name: "health", method: http.MethodGet, path: "/healthz", status: http.StatusNoContent},
		{name: "ready", method: http.MethodGet, path: "/readyz", status: http.StatusNoContent},
		{
			name:        "metrics",
			method:      http.MethodGet,
			path:        "/metrics",
			status:      http.StatusOK,
			contentType: "text/plain; version=0.0.4",
			body:        `gen_ai_request_model="gemini"`,
		},
		{
			name:        "models",
			method:      http.MethodGet,
			path:        "/v1/models",
			status:      http.StatusOK,
			contentType: "application/json",
			body:        `"id":"vertex_ai/gemini"`,
		},
		{name: "models method", method: http.MethodPost, path: "/v1/models", status: http.StatusMethodNotAllowed},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			request := httptest.NewRequest(test.method, test.path, nil)
			response := httptest.NewRecorder()
			handler.ServeHTTP(response, request)
			if response.Code != test.status {
				t.Fatalf("status = %d, want %d", response.Code, test.status)
			}
			if test.contentType != "" && response.Header().Get("Content-Type") != test.contentType {
				t.Fatalf("Content-Type = %q, want %q", response.Header().Get("Content-Type"), test.contentType)
			}
			if test.body != "" && !strings.Contains(response.Body.String(), test.body) {
				t.Fatalf("body = %q, want substring %q", response.Body.String(), test.body)
			}
		})
	}
}

func TestCatalogHandlerRejectsInvalidCatalog(t *testing.T) {
	missing := filepath.Join(t.TempDir(), "missing")
	providers := selectProviders("google-vertex", "novita")
	handler := catalogHandler(missing, providers)
	for _, path := range []string{"/readyz", "/metrics", "/v1/models"} {
		t.Run(path, func(t *testing.T) {
			request := httptest.NewRequest(http.MethodGet, path, nil)
			response := httptest.NewRecorder()
			handler.ServeHTTP(response, request)
			if response.Code != http.StatusServiceUnavailable {
				t.Fatalf("status = %d, want %d", response.Code, http.StatusServiceUnavailable)
			}
		})
	}

	empty := filepath.Join(t.TempDir(), "catalog.json")
	if err := os.WriteFile(empty, []byte(`{"providers":{}}`), 0o600); err != nil {
		t.Fatal(err)
	}
	request := httptest.NewRequest(http.MethodGet, "/readyz", nil)
	response := httptest.NewRecorder()
	catalogHandler(empty, providers).ServeHTTP(response, request)
	if response.Code != http.StatusServiceUnavailable {
		t.Fatalf("empty catalog status = %d, want %d", response.Code, http.StatusServiceUnavailable)
	}
}

func TestFetchJSON(t *testing.T) {
	client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
		if request.Header.Get("Accept") != "application/json" {
			t.Errorf("Accept = %q", request.Header.Get("Accept"))
		}
		if request.Header.Get("Authorization") != "Bearer test" {
			t.Errorf("Authorization = %q", request.Header.Get("Authorization"))
		}
		return response(http.StatusOK, `{"value":12345678901234567890}`), nil
	})}

	var target struct {
		Value json.Number `json:"value"`
	}
	if err := fetchJSON(context.Background(), client, "https://models.invalid", "Bearer test", &target); err != nil {
		t.Fatal(err)
	}
	if target.Value.String() != "12345678901234567890" {
		t.Fatalf("value = %q", target.Value)
	}
}

func TestFetchJSONRejectsHTTPAndJSONErrors(t *testing.T) {
	tests := []struct {
		name   string
		status int
		body   string
	}{
		{name: "status", status: http.StatusBadGateway, body: `{"error":"upstream"}`},
		{name: "JSON", status: http.StatusOK, body: "{"},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
				return response(test.status, test.body), nil
			})}
			if err := fetchJSON(context.Background(), client, "https://models.invalid", "", &map[string]any{}); err == nil {
				t.Fatal("fetchJSON() unexpectedly succeeded")
			}
		})
	}
}

func TestWriteConfigMapPatchesExistingConfigMap(t *testing.T) {
	client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
		if request.Method != http.MethodPatch {
			t.Errorf("method = %q", request.Method)
		}
		if request.URL.Path != "/configmaps/agentgateway-model-costs" {
			t.Errorf("path = %q", request.URL.Path)
		}
		if request.Header.Get("Authorization") != "Bearer token" {
			t.Errorf("Authorization = %q", request.Header.Get("Authorization"))
		}
		if request.Header.Get("Content-Type") != "application/merge-patch+json" {
			t.Errorf("Content-Type = %q", request.Header.Get("Content-Type"))
		}
		var patch configMapPatch
		if err := json.NewDecoder(request.Body).Decode(&patch); err != nil {
			t.Error(err)
		}
		if patch.Data["catalog.json"] != `{"providers":{}}` {
			t.Errorf("catalog.json = %q", patch.Data["catalog.json"])
		}
		return response(http.StatusOK, `{}`), nil
	})}

	err := writeConfigMap(
		context.Background(),
		client,
		"https://kubernetes.invalid/configmaps",
		[]byte("token"),
		[]byte(`{"providers":{}}`),
	)
	if err != nil {
		t.Fatal(err)
	}
}

func TestWriteConfigMapCreatesMissingConfigMap(t *testing.T) {
	requests := 0
	client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
		requests++
		if requests == 1 {
			if request.Method != http.MethodPatch {
				t.Errorf("first method = %q", request.Method)
			}
			return response(http.StatusNotFound, ""), nil
		}
		if request.Method != http.MethodPost {
			t.Errorf("second method = %q", request.Method)
		}
		if request.URL.Path != "/configmaps" {
			t.Errorf("path = %q", request.URL.Path)
		}
		var resource configMapResource
		if err := json.NewDecoder(request.Body).Decode(&resource); err != nil {
			t.Error(err)
		}
		if resource.APIVersion != "v1" || resource.Kind != "ConfigMap" {
			t.Errorf("resource = %#v", resource)
		}
		if resource.Metadata.Name != configMapName {
			t.Errorf("name = %q", resource.Metadata.Name)
		}
		if resource.Metadata.Annotations["kustomize.toolkit.fluxcd.io/prune"] != "disabled" {
			t.Errorf("annotations = %#v", resource.Metadata.Annotations)
		}
		if resource.Metadata.Labels["app.kubernetes.io/instance"] != "agentgateway" ||
			resource.Metadata.Labels["app.kubernetes.io/name"] != "model-catalog" ||
			resource.Metadata.Labels["app.kubernetes.io/part-of"] != "agentgateway" {
			t.Errorf("labels = %#v", resource.Metadata.Labels)
		}
		if resource.Data["catalog.json"] != `{"providers":{}}` {
			t.Errorf("catalog.json = %q", resource.Data["catalog.json"])
		}
		return response(http.StatusCreated, `{}`), nil
	})}

	if err := writeConfigMap(
		context.Background(),
		client,
		"https://kubernetes.invalid/configmaps",
		[]byte("token"),
		[]byte(`{"providers":{}}`),
	); err != nil {
		t.Fatal(err)
	}
	if requests != 2 {
		t.Fatalf("requests = %d, want 2", requests)
	}
}

func TestWriteConfigMapRetriesPatchAfterCreateConflict(t *testing.T) {
	requests := 0
	client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
		requests++
		switch requests {
		case 1:
			return response(http.StatusNotFound, ""), nil
		case 2:
			return response(http.StatusConflict, ""), nil
		default:
			if request.Method != http.MethodPatch {
				t.Errorf("third method = %q", request.Method)
			}
			return response(http.StatusOK, `{}`), nil
		}
	})}

	if err := writeConfigMap(
		context.Background(),
		client,
		"https://kubernetes.invalid/configmaps",
		[]byte("token"),
		[]byte(`{"providers":{}}`),
	); err != nil {
		t.Fatal(err)
	}
	if requests != 3 {
		t.Fatalf("requests = %d, want 3", requests)
	}
}

func TestWriteConfigMapRejectsHTTPAndTransportErrors(t *testing.T) {
	client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
		return response(http.StatusForbidden, ""), nil
	})}
	if err := writeConfigMap(context.Background(), client, "https://kubernetes.invalid", []byte("token"), []byte("{}")); err == nil {
		t.Fatal("writeConfigMap() accepted an HTTP error")
	}

	client = &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
		return nil, errors.New("transport failed")
	})}
	if err := writeConfigMap(context.Background(), client, "https://kubernetes.invalid", []byte("token"), []byte("{}")); err == nil {
		t.Fatal("writeConfigMap() accepted a transport error")
	}
}

func TestKubernetesRequestDecodesResponse(t *testing.T) {
	client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
		if request.Header.Get("Content-Type") != "" {
			t.Errorf("Content-Type = %q, want empty", request.Header.Get("Content-Type"))
		}
		return response(http.StatusOK, `{"value":"decoded"}`), nil
	})}
	target := map[string]string{}
	status, err := kubernetesRequest(
		context.Background(),
		client,
		http.MethodGet,
		"https://kubernetes.invalid",
		[]byte("token"),
		nil,
		&target,
	)
	if err != nil {
		t.Fatal(err)
	}
	if status != http.StatusOK || target["value"] != "decoded" {
		t.Fatalf("status = %d, target = %#v", status, target)
	}
}

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

func (function roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) {
	return function(request)
}

func response(status int, body string) *http.Response {
	return &http.Response{
		StatusCode: status,
		Status:     http.StatusText(status),
		Header:     make(http.Header),
		Body:       io.NopCloser(strings.NewReader(body)),
	}
}