package server

import (
	"context"
	"errors"
	"fmt"
	"io"
	"math"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"go.uber.org/zap"
	"go.uber.org/zap/zaptest/observer"

	"github.com/mem9-ai/drive9/pkg/embedding"
	"github.com/mem9-ai/drive9/pkg/meta"
	"github.com/mem9-ai/drive9/pkg/tenant"
)

func TestBuildEmbeddingConfigReplacementRejectsCounterOverflow(t *testing.T) {
	rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNative, "org-1", nil)
	enabled := false
	existing := &meta.TenantEmbeddingConfig{
		TenantID: rt.tenantID, Enabled: true, Generation: math.MaxInt64,
		ConfigRevision: math.MaxInt64,
	}
	if _, err := rt.server.buildEmbeddingConfigReplacement(context.Background(), rt.tenantID,
		&embeddingConfigPutRequest{Enabled: &enabled}, existing); err == nil {
		t.Fatal("counter overflow replacement succeeded")
	}
}

func TestBuildEmbeddingConfigReplacementGenerationTransitions(t *testing.T) {
	rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNative, "org-1", nil)
	ctx := context.Background()
	enabled := true
	disabled := false
	baseA, baseB := "https://api-a.example.com/v1", "https://api-b.example.com/v1"
	modelA, modelB := "embed-a", "embed-b"
	key := "secret"
	existing := &meta.TenantEmbeddingConfig{
		TenantID: rt.tenantID, Enabled: true, APIBase: baseA,
		APIKeyCipher: mustEncryptEmbeddingKey(t, rt, key), Model: modelA,
		Generation: 7, ConfigRevision: 11, UpdatedAt: time.Now().UTC().Add(-time.Hour),
	}

	for _, tc := range []struct {
		name       string
		req        *embeddingConfigPutRequest
		wantGen    uint64
		wantRev    uint64
		wantEnable bool
	}{
		{name: "exact no-op", req: &embeddingConfigPutRequest{Enabled: &enabled, APIBase: &baseA, APIKey: &key, Model: &modelA}, wantGen: 7, wantRev: 11, wantEnable: true},
		{name: "endpoint", req: &embeddingConfigPutRequest{Enabled: &enabled, APIBase: &baseB, APIKey: &key, Model: &modelA}, wantGen: 8, wantRev: 12, wantEnable: true},
		{name: "model", req: &embeddingConfigPutRequest{Enabled: &enabled, APIBase: &baseA, APIKey: &key, Model: &modelB}, wantGen: 8, wantRev: 12, wantEnable: true},
		{name: "disable", req: &embeddingConfigPutRequest{Enabled: &disabled}, wantGen: 8, wantRev: 12, wantEnable: false},
	} {
		t.Run(tc.name, func(t *testing.T) {
			got, err := rt.server.buildEmbeddingConfigReplacement(ctx, rt.tenantID, tc.req, existing)
			if err != nil {
				t.Fatal(err)
			}
			if got.Generation != tc.wantGen || got.ConfigRevision != tc.wantRev || got.Enabled != tc.wantEnable {
				t.Fatalf("replacement = enabled %t counters (%d,%d), want enabled %t counters (%d,%d)", got.Enabled, got.Generation, got.ConfigRevision, tc.wantEnable, tc.wantGen, tc.wantRev)
			}
			if tc.name == "exact no-op" && !got.UpdatedAt.Equal(existing.UpdatedAt) {
				t.Fatalf("no-op updated_at = %v, want %v", got.UpdatedAt, existing.UpdatedAt)
			}
		})
	}

	offMarker := &meta.TenantEmbeddingConfig{
		TenantID: rt.tenantID, Enabled: false, Generation: 8, ConfigRevision: 12,
		UpdatedAt: time.Now().UTC().Add(-time.Hour),
	}
	got, err := rt.server.buildEmbeddingConfigReplacement(ctx, rt.tenantID,
		&embeddingConfigPutRequest{Enabled: &enabled, APIBase: &baseA, APIKey: &key, Model: &modelA}, offMarker)
	if err != nil {
		t.Fatal(err)
	}
	if !got.Enabled || got.Generation != 9 || got.ConfigRevision != 13 {
		t.Fatalf("re-enable = enabled %t counters (%d,%d), want enabled true counters (9,13)", got.Enabled, got.Generation, got.ConfigRevision)
	}
}

func mustEncryptEmbeddingKey(t *testing.T, rt *extractConfigRuntime, key string) []byte {
	t.Helper()
	cipher, err := rt.pool.Encrypt(context.Background(), []byte(key))
	if err != nil {
		t.Fatal(err)
	}
	return cipher
}

func setNativeEmbeddingMode(t *testing.T, rt *extractConfigRuntime, mode string) {
	t.Helper()
	now := time.Now().UTC()
	if err := rt.meta.UpsertTenantAutoEmbeddingProfile(context.Background(), &meta.TenantAutoEmbeddingProfile{
		TenantID: rt.tenantID, EmbeddingMode: mode, Dimensions: 1024,
		CreatedAt: now, UpdatedAt: now,
	}); err != nil {
		t.Fatal(err)
	}
}

func doEmbeddingConfigRequest(t *testing.T, ts *httptest.Server, method, tenantID, body string, headers map[string]string) *http.Response {
	t.Helper()
	var reader io.Reader
	if body != "" {
		reader = strings.NewReader(body)
	}
	req, err := http.NewRequest(method, ts.URL+"/v1/admin/tenants/"+tenantID+"/embedding-config", reader)
	if err != nil {
		t.Fatal(err)
	}
	for key, value := range headers {
		req.Header.Set(key, value)
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	return resp
}

func TestInnerControlPlaneRouteEmbeddingConfigMappings(t *testing.T) {
	for _, tc := range []struct {
		method string
		want   string
	}{
		{method: http.MethodGet, want: "/drive9/orgs/{org_id}/admin/tenants/view"},
		{method: http.MethodPut, want: "/drive9/orgs/{org_id}/admin/tenants/modify"},
	} {
		got, ok := matchInnerControlPlaneRoute(tc.method, "/v1/admin/tenants/tenant-1/embedding-config")
		if !ok || got != tc.want {
			t.Fatalf("route %s = (%q, %t), want %q", tc.method, got, ok, tc.want)
		}
	}
}

func TestAdminEmbeddingConfigRejectsNullUnknownAndTrailingFields(t *testing.T) {
	rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNative, "org-1", nil)
	setNativeEmbeddingMode(t, rt, meta.TenantEmbeddingModeFTSOnly)
	rt.server.embeddingProviderValidator = func(context.Context, *embeddingConfigPutRequest) error { return nil }
	ts := httptest.NewServer(rt.server)
	defer ts.Close()

	for _, body := range []string{
		`{"enabled":null}`,
		`{"enabled":false,"api_base":null}`,
		`{"enabled":true,"api_base":"https://api.example.com","api_key":null,"model":"m"}`,
		`{"enabled":false,"extra":1}`,
		`{"enabled":false}{"enabled":false}`,
	} {
		resp := doEmbeddingConfigRequest(t, ts, http.MethodPut, rt.tenantID, body, legacyExtractHeaders())
		if resp.StatusCode != http.StatusBadRequest {
			out := decodeExtractConfigBody(t, resp)
			t.Fatalf("body %s status=%d response=%v, want 400", body, resp.StatusCode, out)
		}
		_ = decodeExtractConfigBody(t, resp)
	}
	if _, err := rt.meta.GetTenantEmbeddingConfig(context.Background(), rt.tenantID); !errors.Is(err, meta.ErrNotFound) {
		t.Fatalf("invalid requests persisted config: %v", err)
	}
}

func TestAdminEmbeddingConfigRejectsOversizedProviderFieldsBeforeValidation(t *testing.T) {
	rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNative, "org-1", nil)
	setNativeEmbeddingMode(t, rt, meta.TenantEmbeddingModeFTSOnly)
	validations := 0
	rt.server.embeddingProviderValidator = func(context.Context, *embeddingConfigPutRequest) error {
		validations++
		return nil
	}
	ts := httptest.NewServer(rt.server)
	defer ts.Close()

	longBase := "https://api.example.com/" + strings.Repeat("a", meta.TenantEmbeddingAPIBaseMaxChars)
	longModel := strings.Repeat("m", meta.TenantEmbeddingModelMaxChars+1)
	for _, body := range []string{
		fmt.Sprintf(`{"enabled":true,"api_base":%q,"api_key":"secret","model":"model"}`, longBase),
		fmt.Sprintf(`{"enabled":true,"api_base":"https://api.example.com","api_key":"secret","model":%q}`, longModel),
	} {
		resp := doEmbeddingConfigRequest(t, ts, http.MethodPut, rt.tenantID, body, legacyExtractHeaders())
		if resp.StatusCode != http.StatusBadRequest {
			t.Fatalf("status=%d response=%v, want 400", resp.StatusCode, decodeExtractConfigBody(t, resp))
		}
		_ = decodeExtractConfigBody(t, resp)
	}
	if validations != 0 {
		t.Fatalf("provider validations=%d, want 0", validations)
	}
}

func TestAdminEmbeddingConfigRejectsAPIKeyWhoseCiphertextExceedsSchemaLimit(t *testing.T) {
	rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNative, "org-1", nil)
	setNativeEmbeddingMode(t, rt, meta.TenantEmbeddingModeFTSOnly)
	rt.server.embeddingProviderValidator = func(context.Context, *embeddingConfigPutRequest) error { return nil }
	ts := httptest.NewServer(rt.server)
	defer ts.Close()

	body := fmt.Sprintf(`{"enabled":true,"api_base":"https://api.example.com","api_key":%q,"model":"model"}`,
		strings.Repeat("k", meta.TenantEmbeddingAPIKeyCipherMaxBytes+1))
	resp := doEmbeddingConfigRequest(t, ts, http.MethodPut, rt.tenantID, body, legacyExtractHeaders())
	if resp.StatusCode != http.StatusBadRequest {
		t.Fatalf("status=%d response=%v, want 400", resp.StatusCode, decodeExtractConfigBody(t, resp))
	}
	_ = decodeExtractConfigBody(t, resp)
	if _, err := rt.meta.GetTenantEmbeddingConfig(context.Background(), rt.tenantID); !errors.Is(err, meta.ErrNotFound) {
		t.Fatalf("oversized encrypted key persisted config: %v", err)
	}
}

func TestAdminEmbeddingConfigNativeFTSOnlyLifecycleDoesNotOpenTenantDB(t *testing.T) {
	rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNative, "org-1", nil)
	setNativeEmbeddingMode(t, rt, meta.TenantEmbeddingModeFTSOnly)
	validations := 0
	rt.server.embeddingProviderValidator = func(context.Context, *embeddingConfigPutRequest) error {
		validations++
		return nil
	}
	ts := httptest.NewServer(rt.server)
	defer ts.Close()

	resp := doEmbeddingConfigRequest(t, ts, http.MethodGet, rt.tenantID, "", legacyExtractHeaders())
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("initial GET status=%d", resp.StatusCode)
	}
	out := decodeExtractConfigBody(t, resp)
	if out["enabled"] != false || out["source"] != "none" {
		t.Fatalf("initial response=%v", out)
	}

	resp = doEmbeddingConfigRequest(t, ts, http.MethodPut, rt.tenantID,
		`{"enabled":true,"api_base":"https://api.example.com/v1","api_key":"sk-first-secret","model":"embed-a"}`,
		legacyExtractHeaders())
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("enable status=%d response=%v", resp.StatusCode, decodeExtractConfigBody(t, resp))
	}
	out = decodeExtractConfigBody(t, resp)
	if out["source"] != "custom" || out["enabled"] != true || out["api_key"] != maskAPIKey("sk-first-secret") || out["generation"] != float64(1) {
		t.Fatalf("enable response=%v", out)
	}
	row, err := rt.meta.GetTenantEmbeddingConfig(context.Background(), rt.tenantID)
	if err != nil {
		t.Fatal(err)
	}
	if row.Generation != 1 || row.ConfigRevision != 1 {
		t.Fatalf("initial counters=(%d,%d), want (1,1)", row.Generation, row.ConfigRevision)
	}

	resp = doEmbeddingConfigRequest(t, ts, http.MethodPut, rt.tenantID,
		`{"enabled":true,"api_base":"https://api.example.com/v1","api_key":"sk-replaced-secret","model":"embed-a"}`,
		legacyExtractHeaders())
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("key replacement status=%d response=%v", resp.StatusCode, decodeExtractConfigBody(t, resp))
	}
	_ = decodeExtractConfigBody(t, resp)
	row, err = rt.meta.GetTenantEmbeddingConfig(context.Background(), rt.tenantID)
	if err != nil {
		t.Fatal(err)
	}
	if row.Generation != 1 || row.ConfigRevision != 2 {
		t.Fatalf("key replacement counters=(%d,%d), want (1,2)", row.Generation, row.ConfigRevision)
	}

	resp = doEmbeddingConfigRequest(t, ts, http.MethodPut, rt.tenantID, `{"enabled":false}`, legacyExtractHeaders())
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("disable status=%d response=%v", resp.StatusCode, decodeExtractConfigBody(t, resp))
	}
	out = decodeExtractConfigBody(t, resp)
	if out["enabled"] != false || out["source"] != "custom" {
		t.Fatalf("disable response=%v", out)
	}
	if validations != 2 {
		t.Fatalf("provider validations=%d, want 2", validations)
	}
	tenantRow, err := rt.meta.GetTenant(context.Background(), rt.tenantID)
	if err != nil {
		t.Fatal(err)
	}
	if _, release, ok := rt.pool.AcquireCached(tenantRow); ok {
		release()
		t.Fatal("embedding config API opened a cold tenant database")
	}
}

func TestAdminEmbeddingConfigNativeAutoIsReadOnlyDatabaseManaged(t *testing.T) {
	rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNative, "org-1", nil)
	setNativeEmbeddingMode(t, rt, meta.TenantEmbeddingModeAuto)
	validated := false
	rt.server.embeddingProviderValidator = func(context.Context, *embeddingConfigPutRequest) error {
		validated = true
		return nil
	}
	ts := httptest.NewServer(rt.server)
	defer ts.Close()

	resp := doEmbeddingConfigRequest(t, ts, http.MethodGet, rt.tenantID, "", legacyExtractHeaders())
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("GET status=%d", resp.StatusCode)
	}
	out := decodeExtractConfigBody(t, resp)
	if out["enabled"] != false || out["source"] != "database_auto" {
		t.Fatalf("GET response=%v", out)
	}
	resp = doEmbeddingConfigRequest(t, ts, http.MethodPut, rt.tenantID,
		`{"enabled":true,"api_base":"https://api.example.com","api_key":"secret","model":"m"}`,
		legacyExtractHeaders())
	if resp.StatusCode != http.StatusConflict {
		t.Fatalf("PUT status=%d response=%v, want 409", resp.StatusCode, decodeExtractConfigBody(t, resp))
	}
	_ = decodeExtractConfigBody(t, resp)
	if validated {
		t.Fatal("native auto PUT called provider validator")
	}
	if _, err := rt.meta.GetTenantEmbeddingConfig(context.Background(), rt.tenantID); !errors.Is(err, meta.ErrNotFound) {
		t.Fatalf("native auto PUT persisted config: %v", err)
	}
}

func TestAdminEmbeddingConfigProviderFailureDoesNotPersist(t *testing.T) {
	rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNative, "org-1", nil)
	setNativeEmbeddingMode(t, rt, meta.TenantEmbeddingModeFTSOnly)
	rt.server.embeddingProviderValidator = func(context.Context, *embeddingConfigPutRequest) error {
		return errors.New("provider leaked secret detail")
	}
	ts := httptest.NewServer(rt.server)
	defer ts.Close()

	resp := doEmbeddingConfigRequest(t, ts, http.MethodPut, rt.tenantID,
		`{"enabled":true,"api_base":"https://api.example.com","api_key":"secret","model":"m"}`,
		legacyExtractHeaders())
	if resp.StatusCode != http.StatusBadGateway {
		t.Fatalf("status=%d response=%v, want 502", resp.StatusCode, decodeExtractConfigBody(t, resp))
	}
	out := decodeExtractConfigBody(t, resp)
	if out["error"] != "embedding provider validation failed: provider request failed" {
		t.Fatalf("response=%v", out)
	}
	if _, err := rt.meta.GetTenantEmbeddingConfig(context.Background(), rt.tenantID); !errors.Is(err, meta.ErrNotFound) {
		t.Fatalf("failed provider validation persisted config: %v", err)
	}
}

func TestAdminEmbeddingConfigProviderFailureLogIsSanitized(t *testing.T) {
	const providerSecret = "provider-response-secret-detail"
	rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNative, "org-1", nil)
	setNativeEmbeddingMode(t, rt, meta.TenantEmbeddingModeFTSOnly)
	provider := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusUnauthorized)
		_, _ = io.WriteString(w, `{"error":{"message":"`+providerSecret+`"}}`)
	}))
	defer provider.Close()
	providerClient, err := embedding.NewOpenAIClient(embedding.OpenAIClientConfig{
		BaseURL: provider.URL, APIKey: "secret", Model: "model", Client: provider.Client(),
	})
	if err != nil {
		t.Fatal(err)
	}
	rt.server.embeddingProviderValidator = func(context.Context, *embeddingConfigPutRequest) error {
		_, err := providerClient.EmbedText(context.Background(), "probe")
		return err
	}
	core, recorded := observer.New(zap.WarnLevel)
	rt.server.logger = zap.New(core)
	rt.server.embeddingProviderValidator = func(context.Context, *embeddingConfigPutRequest) error {
		_, err := providerClient.EmbedText(context.Background(), "probe")
		return err
	}
	ts := httptest.NewServer(rt.server)
	defer ts.Close()

	resp := doEmbeddingConfigRequest(t, ts, http.MethodPut, rt.tenantID,
		`{"enabled":true,"api_base":"https://api.example.com","api_key":"secret","model":"m"}`,
		legacyExtractHeaders())
	if resp.StatusCode != http.StatusBadRequest {
		t.Fatalf("status=%d response=%v, want 400", resp.StatusCode, decodeExtractConfigBody(t, resp))
	}
	_ = decodeExtractConfigBody(t, resp)
	entries := recorded.FilterMessage("admin_tenant_embedding_provider_validation_failed").All()
	if len(entries) != 1 {
		t.Fatalf("provider failure logs=%d, want 1", len(entries))
	}
	if strings.Contains(entries[0].Message+fmt.Sprint(entries[0].ContextMap()), providerSecret) {
		t.Fatalf("provider response detail leaked in log: %+v", entries[0])
	}
}

func TestEmbeddingProviderValidationResponseUsesTypedSafeFailure(t *testing.T) {
	provider := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusUnauthorized)
		_, _ = io.WriteString(w, `{"error":{"message":"credential detail must stay internal"}}`)
	}))
	defer provider.Close()
	client, err := embedding.NewOpenAIClient(embedding.OpenAIClientConfig{
		BaseURL: provider.URL, APIKey: "secret", Model: "model", Client: provider.Client(),
	})
	if err != nil {
		t.Fatal(err)
	}
	_, err = client.EmbedText(context.Background(), "probe")
	if err == nil {
		t.Fatal("provider call succeeded")
	}
	if got := safeEmbeddingProviderValidationError(err); got != "provider rejected the credentials" {
		t.Fatalf("safe error=%q", got)
	}
	if got := embeddingProviderValidationHTTPStatus(err); got != http.StatusBadRequest {
		t.Fatalf("status=%d, want 400", got)
	}
}

func TestAdminEmbeddingConfigSharedAndCrossOrg(t *testing.T) {
	rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNativeShared, "", nil)
	bindSharedExtractTenant(t, rt, "org-shared")
	rt.server.embeddingProviderValidator = func(context.Context, *embeddingConfigPutRequest) error { return nil }
	rt.prov.iamIdentities = []*tenant.TiDBCloudAPIKeyIdentity{
		{OrganizationID: "org-shared", Role: tenant.TiDBCloudRoleOrgOwner},
	}
	ts := httptest.NewServer(rt.server)
	defer ts.Close()

	resp := doEmbeddingConfigRequest(t, ts, http.MethodPut, rt.tenantID,
		`{"enabled":true,"api_base":"https://api.example.com","api_key":"shared-secret","model":"m"}`,
		legacyExtractHeaders())
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("shared PUT status=%d response=%v", resp.StatusCode, decodeExtractConfigBody(t, resp))
	}
	_ = decodeExtractConfigBody(t, resp)

	rt.prov.iamIdentities = []*tenant.TiDBCloudAPIKeyIdentity{{OrganizationID: "org-other", Role: tenant.TiDBCloudRoleOrgOwner}}
	headers := map[string]string{quotaPublicKeyHeader: "other-public", quotaPrivateKeyHeader: "other-private"}
	for _, method := range []string{http.MethodGet, http.MethodPut} {
		body := ""
		if method == http.MethodPut {
			body = `{"enabled":false}`
		}
		resp = doEmbeddingConfigRequest(t, ts, method, rt.tenantID, body, headers)
		if resp.StatusCode != http.StatusForbidden {
			t.Fatalf("cross-org %s status=%d response=%v, want 403", method, resp.StatusCode, decodeExtractConfigBody(t, resp))
		}
		_ = decodeExtractConfigBody(t, resp)
	}
}