package server

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
	"time"

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

func TestAdminExtractConfigConcurrentPartialPuts(t *testing.T) {
	rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNative, "org-1", nil)
	seedExtractConfigRow(t, rt, true, "https://a.example.com", "key-one-1", "model-0", "prompt-0")
	ts := httptest.NewServer(rt.server)
	defer ts.Close()

	// Concurrent full-trio and prompt-only PUTs updating disjoint fields must both win: the
	// read-modify-write transaction holds SELECT ... FOR UPDATE on the row.
	// Rounds are serialized so the final state is deterministic.
	const rounds = 4
	for i := 1; i <= rounds; i++ {
		var wg sync.WaitGroup
		errs := make(chan string, 2)
		wg.Add(2)
		go func(i int) {
			defer wg.Done()
			resp, err := doExtractConfigRequestE(ts, http.MethodPut, rt.tenantID, "image",
				fmt.Sprintf(`{"api_base":"https://provider-%d.example.com","api_key":"key-%d-value","model":"model-%d"}`, i, i, i), legacyExtractHeaders())
			if err != nil {
				errs <- fmt.Sprintf("provider PUT %d: %v", i, err)
				return
			}
			defer func() { _ = resp.Body.Close() }()
			if resp.StatusCode != http.StatusOK {
				errs <- fmt.Sprintf("provider PUT %d: status %d", i, resp.StatusCode)
			}
		}(i)
		go func(i int) {
			defer wg.Done()
			resp, err := doExtractConfigRequestE(ts, http.MethodPut, rt.tenantID, "image",
				fmt.Sprintf(`{"prompt":"prompt-%d"}`, i), legacyExtractHeaders())
			if err != nil {
				errs <- fmt.Sprintf("prompt PUT %d: %v", i, err)
				return
			}
			defer func() { _ = resp.Body.Close() }()
			if resp.StatusCode != http.StatusOK {
				errs <- fmt.Sprintf("prompt PUT %d: status %d", i, resp.StatusCode)
			}
		}(i)
		wg.Wait()
		close(errs)
		for msg := range errs {
			t.Fatal(msg)
		}
		row := getPersistedExtractConfig(t, rt)
		wantBase := fmt.Sprintf("https://provider-%d.example.com", i)
		wantModel, wantPrompt, wantKey := fmt.Sprintf("model-%d", i), fmt.Sprintf("prompt-%d", i), fmt.Sprintf("key-%d-value", i)
		if row.APIBase != wantBase || row.Model != wantModel || row.Prompt != wantPrompt || decryptExtractKey(t, rt, row.APIKeyCipher) != wantKey {
			t.Fatalf("round %d: lost update, base = %q model = %q prompt = %q, want %q/%q/%q",
				i, row.APIBase, row.Model, row.Prompt, wantBase, wantModel, wantPrompt)
		}
	}

	row := getPersistedExtractConfig(t, rt)
	if !row.Enabled || row.APIBase != "https://provider-4.example.com" {
		t.Fatalf("row = %+v", row)
	}
	if got := decryptExtractKey(t, rt, row.APIKeyCipher); got != "key-4-value" {
		t.Fatalf("decrypted key = %q, want latest key", got)
	}
}

// A PUT against a tenant that is already deleting must fail closed with 404
// (admin anti-enumeration) and must not create a config row.
func TestAdminExtractConfigPutDeletingTenantNotFound(t *testing.T) {
	rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNative, "org-1", nil)
	ts := httptest.NewServer(rt.server)
	defer ts.Close()
	ctx := context.Background()

	if err := rt.meta.UpdateTenantStatus(ctx, rt.tenantID, meta.TenantDeleting); err != nil {
		t.Fatal(err)
	}
	resp := doExtractConfigRequest(t, ts, http.MethodPut, rt.tenantID, "image",
		`{"enabled":true,"api_base":"https://api.example.com","api_key":"sk-live-key","model":"model-1"}`, legacyExtractHeaders())
	if resp.StatusCode != http.StatusNotFound {
		t.Fatalf("PUT status = %d, want 404", resp.StatusCode)
	}
	_ = decodeExtractConfigBody(t, resp)
	if _, err := rt.meta.GetTenantExtractConfig(ctx, rt.tenantID, meta.ExtractMediaTypeImage); !errors.Is(err, meta.ErrNotFound) {
		t.Fatalf("config row after PUT on deleting tenant = %v, want %v", err, meta.ErrNotFound)
	}
}

// A successful PUT evicts the tenant's cached resolution on the local pod so
// the new config (and its plaintext key) replaces the old one immediately
// instead of after a TTL window.
func TestAdminExtractConfigPutEvictsResolverCache(t *testing.T) {
	rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNative, "org-1", nil)
	resolver := newExtractConfigResolver(rt.meta, rt.pool, defaultTestRuntime(), time.Minute)
	t.Cleanup(resolver.Close)
	rt.server.SetExtractConfigResolver(resolver)
	ts := httptest.NewServer(rt.server)
	defer ts.Close()
	ctx := context.Background()

	// Populate the local cache: the absent row resolves to the process
	// default.
	res, err := resolver.ResolveImage(ctx, rt.tenantID)
	if err != nil {
		t.Fatalf("ResolveImage: %v", err)
	}
	if res.state != extractResolveReady {
		t.Fatalf("state = %v, want ready (process default)", res.state)
	}
	if !cacheHas(resolver, rt.tenantID) {
		t.Fatal("ResolveImage did not populate the cache")
	}

	resp := doExtractConfigRequest(t, ts, http.MethodPut, rt.tenantID, "image",
		`{"enabled":true,"api_base":"https://api.example.com","api_key":"sk-live-key","model":"model-1"}`, legacyExtractHeaders())
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("PUT status = %d, want 200", resp.StatusCode)
	}
	_ = decodeExtractConfigBody(t, resp)
	if cacheHas(resolver, rt.tenantID) {
		t.Fatal("PUT did not evict the cached resolution on the local pod")
	}

	// The next resolve re-reads the row and applies the custom config.
	res, err = resolver.ResolveImage(ctx, rt.tenantID)
	if err != nil {
		t.Fatalf("ResolveImage after PUT: %v", err)
	}
	if res.state != extractResolveReady {
		t.Fatalf("state after PUT = %v, want ready (tenant custom)", res.state)
	}
}

// The local tenant-teardown seam used by every deleting/deleted lifecycle
// transition also drops the tenant's cached extract resolution (and its
// plaintext provider key) on this pod.
func TestClearLocalTenantMetricsEvictsExtractConfigCache(t *testing.T) {
	rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNative, "org-1", nil)
	resolver := newExtractConfigResolver(rt.meta, rt.pool, defaultTestRuntime(), time.Minute)
	t.Cleanup(resolver.Close)
	rt.server.SetExtractConfigResolver(resolver)
	ctx := context.Background()

	if _, err := resolver.ResolveImage(ctx, rt.tenantID); err != nil {
		t.Fatalf("ResolveImage: %v", err)
	}
	if !cacheHas(resolver, rt.tenantID) {
		t.Fatal("ResolveImage did not populate the cache")
	}
	rt.server.clearLocalTenantMetrics(rt.tenantID)
	if cacheHas(resolver, rt.tenantID) {
		t.Fatal("tenant lifecycle cleanup did not evict the cached resolution")
	}
}