package adapterimpl

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/onsi/gomega"

	globalcommonv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/common/v1"
	"github.com/tidbcloud/tidb-management-service/serverless-global-service/domain/adapter"
	"github.com/tidbcloud/tidb-management-service/serverless-global-service/domain/entity"
)

type fakeSys9SecretProvider struct {
	profile *adapter.Sys9OEMProfile
	err     error
}

func (f *fakeSys9SecretProvider) ResolveProfile(_ context.Context, _ string) (*adapter.Sys9OEMProfile, error) {
	if f.err != nil {
		return nil, f.err
	}
	return f.profile, nil
}

func TestSys9AuthManagerGetSessionHappyPathAndCache(t *testing.T) {
	g := gomega.NewWithT(t)

	var challengeCalls atomic.Int32
	var exchangeCalls atomic.Int32
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/api/v1/federation/bootstrap-challenges":
			challengeCalls.Add(1)
			var body map[string]any
			g.Expect(json.NewDecoder(r.Body).Decode(&body)).ShouldNot(gomega.HaveOccurred())
			g.Expect(body["temporaryDpopPublicKey"]).Should(gomega.Equal("tmp-pub"))
			w.Header().Set("Content-Type", "application/json")
			_, _ = w.Write([]byte(`{"challengeId":"challenge-1"}`))
		case "/api/v1/tokens:exchange":
			exchangeCalls.Add(1)
			g.Expect(r.Header.Get("X-Request-Context-ID")).ShouldNot(gomega.BeEmpty())
			g.Expect(r.Header.Get("Idempotency-Key")).Should(gomega.BeEmpty())
			var body map[string]any
			g.Expect(json.NewDecoder(r.Body).Decode(&body)).ShouldNot(gomega.HaveOccurred())
			_, hasRequestContextID := body["requestContextId"]
			_, hasIdempotencyKey := body["idempotencyKey"]
			g.Expect(hasRequestContextID).Should(gomega.BeFalse())
			g.Expect(hasIdempotencyKey).Should(gomega.BeFalse())
			_, hasTenantID := body["tenantId"]
			_, hasProjectID := body["projectId"]
			_, hasClusterID := body["clusterId"]
			g.Expect(hasTenantID).Should(gomega.BeFalse())
			g.Expect(hasProjectID).Should(gomega.BeFalse())
			g.Expect(hasClusterID).Should(gomega.BeFalse())
			w.Header().Set("Content-Type", "application/json")
			_, _ = w.Write([]byte(`{"sessionToken":"sess-1","accessToken":"acc-1","expiresInSec":120}`))
		default:
			w.WriteHeader(http.StatusNotFound)
		}
	}))
	defer server.Close()

	mgr := NewSys9AuthManager(&fakeSys9SecretProvider{profile: &adapter.Sys9OEMProfile{
		BrandID:                "brand",
		AccountID:              "account",
		Issuer:                 "issuer",
		Audience:               "audience",
		Assertion:              "assertion",
		SubjectAnchor:          "sa",
		SubjectAnchorVersion:   1,
		EntityAnchor:           "ea",
		EntityAnchorVersion:    1,
		TemporaryDPoPPublicKey: "tmp-pub",
		DPoPPublicKey:          "pub",
		DPoPProof:              "proof",
		TLSExporterDigest:      "digest",
		EdgeBinding:            adapter.Sys9EdgeBinding{"binding": "edge"},
	}})

	region := &entity.Region{
		Id:              globalcommonv1.Region_AWS_US_WEST_2,
		CloudProviderId: globalcommonv1.CloudProvider_AWS,
		ServicePlan:     globalcommonv1.ServicePlan_Starter,
		EngineType:      globalcommonv1.Region_POSTGRESQL_COMPATIBLE,
		ServiceAddress:  server.URL,
		Sys9Config: &globalcommonv1.RegionSys9Config{
			OemProfileRef:           "profile/ref",
			TokenTtlSec:             120,
			SessionRefreshBeforeSec: 30,
		},
	}

	s1, err := mgr.GetSession(t.Context(), region)
	g.Expect(err).ShouldNot(gomega.HaveOccurred())
	g.Expect(s1.AccessToken).Should(gomega.Equal("acc-1"))

	s2, err := mgr.GetSession(t.Context(), region)
	g.Expect(err).ShouldNot(gomega.HaveOccurred())
	g.Expect(s2.AccessToken).Should(gomega.Equal("acc-1"))
	g.Expect(challengeCalls.Load()).Should(gomega.Equal(int32(1)))
	g.Expect(exchangeCalls.Load()).Should(gomega.Equal(int32(1)))
}

func TestSys9AuthManagerConcurrentExpiredRefreshSingleflight(t *testing.T) {
	g := gomega.NewWithT(t)

	var challengeCalls atomic.Int32
	var exchangeCalls atomic.Int32
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/api/v1/federation/bootstrap-challenges":
			challengeCalls.Add(1)
			var body map[string]any
			g.Expect(json.NewDecoder(r.Body).Decode(&body)).ShouldNot(gomega.HaveOccurred())
			g.Expect(body["temporaryDpopPublicKey"]).Should(gomega.Equal("tmp-pub"))
			w.Header().Set("Content-Type", "application/json")
			_, _ = w.Write([]byte(`{"challengeId":"challenge-2"}`))
		case "/api/v1/tokens:exchange":
			exchangeCalls.Add(1)
			g.Expect(r.Header.Get("X-Request-Context-ID")).ShouldNot(gomega.BeEmpty())
			g.Expect(r.Header.Get("Idempotency-Key")).Should(gomega.BeEmpty())
			var body map[string]any
			g.Expect(json.NewDecoder(r.Body).Decode(&body)).ShouldNot(gomega.HaveOccurred())
			_, hasRequestContextID := body["requestContextId"]
			_, hasIdempotencyKey := body["idempotencyKey"]
			g.Expect(hasRequestContextID).Should(gomega.BeFalse())
			g.Expect(hasIdempotencyKey).Should(gomega.BeFalse())
			_, hasTenantID := body["tenantId"]
			_, hasProjectID := body["projectId"]
			_, hasClusterID := body["clusterId"]
			g.Expect(hasTenantID).Should(gomega.BeFalse())
			g.Expect(hasProjectID).Should(gomega.BeFalse())
			g.Expect(hasClusterID).Should(gomega.BeFalse())
			w.Header().Set("Content-Type", "application/json")
			_, _ = w.Write([]byte(`{"sessionToken":"sess-2","accessToken":"acc-2","expiresInSec":120}`))
		default:
			w.WriteHeader(http.StatusNotFound)
		}
	}))
	defer server.Close()

	mgrRaw := NewSys9AuthManager(&fakeSys9SecretProvider{profile: &adapter.Sys9OEMProfile{
		BrandID:                "brand",
		AccountID:              "account",
		Issuer:                 "issuer",
		Audience:               "audience",
		Assertion:              "assertion",
		SubjectAnchor:          "sa",
		SubjectAnchorVersion:   1,
		EntityAnchor:           "ea",
		EntityAnchorVersion:    1,
		TemporaryDPoPPublicKey: "tmp-pub",
		DPoPPublicKey:          "pub",
		DPoPProof:              "proof",
		TLSExporterDigest:      "digest",
		EdgeBinding:            adapter.Sys9EdgeBinding{"binding": "edge"},
	}})
	mgr, ok := mgrRaw.(*sys9AuthManager)
	g.Expect(ok).Should(gomega.BeTrue())

	region := &entity.Region{
		Id:              globalcommonv1.Region_AWS_US_WEST_2,
		CloudProviderId: globalcommonv1.CloudProvider_AWS,
		ServicePlan:     globalcommonv1.ServicePlan_Starter,
		EngineType:      globalcommonv1.Region_POSTGRESQL_COMPATIBLE,
		ServiceAddress:  server.URL,
		Sys9Config: &globalcommonv1.RegionSys9Config{
			OemProfileRef:           "profile/ref",
			TokenTtlSec:             120,
			SessionRefreshBeforeSec: 30,
		},
	}

	key := mgr.cacheKey(region)
	mgr.mu.Lock()
	mgr.cache[key] = &adapter.Sys9Session{
		AccessToken:  "expired",
		SessionToken: "expired",
		ExpiresAt:    time.Now().Add(-time.Second),
		RefreshAfter: time.Now().Add(-2 * time.Second),
	}
	mgr.mu.Unlock()

	var wg sync.WaitGroup
	errCh := make(chan error, 20)
	for range 20 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			s, err := mgr.GetSession(context.Background(), region)
			if err != nil {
				errCh <- err
				return
			}
			if s.AccessToken != "acc-2" {
				errCh <- errors.New("unexpected access token")
			}
		}()
	}
	wg.Wait()
	close(errCh)
	for err := range errCh {
		g.Expect(err).ShouldNot(gomega.HaveOccurred())
	}

	g.Expect(challengeCalls.Load()).Should(gomega.Equal(int32(1)))
	g.Expect(exchangeCalls.Load()).Should(gomega.Equal(int32(1)))
}

func TestSys9AuthManagerMapKnownError(t *testing.T) {
	g := gomega.NewWithT(t)
	err := mapSys9Error(http.StatusUnauthorized, []byte(`{"code":"INVALID_PROOF","message":"invalid"}`))
	g.Expect(err).Should(gomega.HaveOccurred())
	g.Expect(err.Error()).Should(gomega.ContainSubstring("sys9 auth failed"))
}

func TestSys9AuthManagerTokenExchangeSessionIDFallback(t *testing.T) {
	g := gomega.NewWithT(t)

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/api/v1/federation/bootstrap-challenges":
			w.Header().Set("Content-Type", "application/json")
			_, _ = w.Write([]byte(`{"challengeId":"challenge-3"}`))
		case "/api/v1/tokens:exchange":
			w.Header().Set("Content-Type", "application/json")
			_, _ = w.Write([]byte(`{"sessionId":"sess-id-only","expiresInSec":120}`))
		default:
			w.WriteHeader(http.StatusNotFound)
		}
	}))
	defer server.Close()

	mgr := NewSys9AuthManager(&fakeSys9SecretProvider{profile: &adapter.Sys9OEMProfile{
		BrandID:                "brand",
		AccountID:              "account",
		Issuer:                 "issuer",
		Assertion:              "assertion",
		SubjectAnchor:          "sa",
		SubjectAnchorVersion:   1,
		EntityAnchor:           "ea",
		EntityAnchorVersion:    1,
		TemporaryDPoPPublicKey: "tmp-pub",
		DPoPProof:              "proof",
		TLSExporterDigest:      "digest",
		EdgeBinding:            adapter.Sys9EdgeBinding{"binding": "edge"},
	}})

	region := &entity.Region{
		Id:              globalcommonv1.Region_AWS_US_WEST_2,
		CloudProviderId: globalcommonv1.CloudProvider_AWS,
		ServicePlan:     globalcommonv1.ServicePlan_Starter,
		EngineType:      globalcommonv1.Region_POSTGRESQL_COMPATIBLE,
		ServiceAddress:  server.URL,
		Sys9Config: &globalcommonv1.RegionSys9Config{
			OemProfileRef:           "profile/ref",
			TokenTtlSec:             120,
			SessionRefreshBeforeSec: 30,
		},
	}

	s, err := mgr.GetSession(t.Context(), region)
	g.Expect(err).ShouldNot(gomega.HaveOccurred())
	g.Expect(s.SessionToken).Should(gomega.Equal("sess-id-only"))
	g.Expect(s.AccessToken).Should(gomega.Equal("sess-id-only"))
}