package server

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

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

func TestInnerControlPlaneTokenIssueSourceUsesAuthenticatedSubject(t *testing.T) {
	tests := []struct {
		name     string
		subject  innerauth.ControlPlaneSubject
		wantKey  string
		wantKind string
	}{
		{
			name: "JWT user",
			subject: innerauth.ControlPlaneSubject{
				OrganizationID: 60205,
				UserID:         42,
				Kind:           innerauth.CredentialJWT,
			},
			wantKey:  "user:42",
			wantKind: "jwt",
		},
		{
			name: "API key",
			subject: innerauth.ControlPlaneSubject{
				OrganizationID: 60205,
				PublicKey:      "public-key",
				Kind:           innerauth.CredentialAPIKey,
			},
			wantKey:  "api_key:public-key",
			wantKind: "api_key",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctx := withInnerAuthSubject(context.Background(), &tt.subject)
			source := controlPlaneTokenIssueSource(ctx, &ControlPlaneScope{
				Identity: tenant.TiDBCloudAPIKeyIdentity{
					OrganizationID: "60205",
					Role:           tenant.TiDBCloudRoleOrgOwner,
				},
			})
			if source.SubjectKey != tt.wantKey {
				t.Fatalf("subject key = %q, want %q", source.SubjectKey, tt.wantKey)
			}
			if source.Provider != "tidbcloud" {
				t.Fatalf("provider = %q, want tidbcloud", source.Provider)
			}
			var metadata map[string]string
			if err := json.Unmarshal(source.MetadataJSON, &metadata); err != nil {
				t.Fatalf("decode metadata: %v", err)
			}
			if metadata["credential_kind"] != tt.wantKind {
				t.Fatalf("credential kind = %q, want %q", metadata["credential_kind"], tt.wantKind)
			}
			if _, ok := metadata["role"]; ok {
				t.Fatalf("inner issue metadata must not contain a synthesized role: %s", source.MetadataJSON)
			}
		})
	}
}

func TestControlPlaneCredentialsUseAuthenticatedExecutionCredentials(t *testing.T) {
	want := tenant.CredentialProvisionRequest{PublicKey: "server-public", PrivateKey: "server-private"}
	ctx := withInnerAuthSubject(context.Background(), &innerauth.ControlPlaneSubject{
		OrganizationID:       60205,
		Kind:                 innerauth.CredentialJWT,
		ExecutionCredentials: &want,
	})
	req := httptest.NewRequest(http.MethodGet, "/v1/admin/tenants?orgId=60205", nil).WithContext(ctx)

	got, err := controlPlaneCredentials(req, "body-public", "body-private")
	if err != nil {
		t.Fatalf("controlPlaneCredentials: %v", err)
	}
	if got != want {
		t.Fatalf("credentials = %#v, want authenticated execution credentials %#v", got, want)
	}
}

func TestControlPlaneOrganizationUsesAuthenticatedSubject(t *testing.T) {
	ctx := withInnerAuthSubject(context.Background(), &innerauth.ControlPlaneSubject{
		OrganizationID: 60205,
		Kind:           innerauth.CredentialJWT,
	})

	got, ok := controlPlaneOrganizationID(ctx)
	if !ok || got != "60205" {
		t.Fatalf("organization = %q/%t, want 60205/true", got, ok)
	}
}

func TestControlPlaneCredentialsPreserveLegacyHeaderBehavior(t *testing.T) {
	req := httptest.NewRequest(http.MethodGet, "/v1/admin/tenants", nil)
	req.Header.Set(quotaPublicKeyHeader, "header-public")
	req.Header.Set(quotaPrivateKeyHeader, "header-private")

	got, err := controlPlaneCredentials(req, "", "")
	if err != nil {
		t.Fatalf("controlPlaneCredentials: %v", err)
	}
	want := tenant.CredentialProvisionRequest{PublicKey: "header-public", PrivateKey: "header-private"}
	if got != want {
		t.Fatalf("credentials = %#v, want %#v", got, want)
	}
}

func TestResolveTiDBCloudIdentityUsesCompletedInnerAuthentication(t *testing.T) {
	prov := &fakeProvisioner{provider: tenant.ProviderTiDBCloudNative}
	srv := NewWithConfig(Config{Provisioner: prov})
	t.Cleanup(srv.Close)
	ctx := withInnerAuthSubject(context.Background(), &innerauth.ControlPlaneSubject{
		OrganizationID: 60205,
		Kind:           innerauth.CredentialJWT,
	})

	identity, err := srv.resolveTiDBCloudIdentity(ctx, tenant.CredentialProvisionRequest{
		PublicKey: "server-public", PrivateKey: "server-private",
	}, "test")
	if err != nil {
		t.Fatalf("resolveTiDBCloudIdentity: %v", err)
	}
	if identity.OrganizationID != "60205" {
		t.Fatalf("organization = %q, want 60205", identity.OrganizationID)
	}
	if identity.Role != "" {
		t.Fatalf("role = %q, want empty after inner authorization", identity.Role)
	}
	if prov.iamCalls.Load() != 0 {
		t.Fatalf("legacy IAM calls = %d, want 0 after inner auth", prov.iamCalls.Load())
	}
}

func TestResolveTiDBCloudAccessProfileUsesInnerOrganizationAndSkipsIAM(t *testing.T) {
	prov := &fakeProvisioner{provider: tenant.ProviderTiDBCloudNative, billingErr: errors.New("legacy Billing must not be called")}
	billing := &innerBillingFake{plan: &tenant.TiDBCloudOrganizationPlan{OrganizationID: "60205"}}
	srv := NewWithConfig(Config{Provisioner: prov, InnerBilling: billing})
	t.Cleanup(srv.Close)
	ctx := withInnerAuthSubject(context.Background(), &innerauth.ControlPlaneSubject{
		OrganizationID: 60205,
		Kind:           innerauth.CredentialJWT,
	})

	profile, err := srv.resolveTiDBCloudAccessProfile(ctx, tenant.CredentialProvisionRequest{
		PublicKey: "server-public", PrivateKey: "server-private",
	}, "test")
	if err != nil {
		t.Fatalf("resolveTiDBCloudAccessProfile: %v", err)
	}
	if profile.OrganizationID != "60205" {
		t.Fatalf("organization = %q, want 60205", profile.OrganizationID)
	}
	if prov.iamCalls.Load() != 0 {
		t.Fatalf("legacy IAM calls = %d, want 0", prov.iamCalls.Load())
	}
	if prov.billingCalls.Load() != 0 {
		t.Fatalf("legacy Billing calls = %d, want 0", prov.billingCalls.Load())
	}
	if billing.organizationID != 60205 || billing.calls != 1 {
		t.Fatalf("inner Billing calls = %d, org = %d, want 1/60205", billing.calls, billing.organizationID)
	}
	if _, err := srv.resolveTiDBCloudAccessProfile(ctx, tenant.CredentialProvisionRequest{}, "test"); err != nil {
		t.Fatalf("cached resolveTiDBCloudAccessProfile: %v", err)
	}
	if billing.calls != 1 {
		t.Fatalf("inner Billing calls after cache hit = %d, want 1", billing.calls)
	}
}

func TestResolveTiDBCloudAccessProfileUsesOpenAPIIAMAndInnerBillingWithoutInnerSubject(t *testing.T) {
	prov := &fakeProvisioner{
		provider:    tenant.ProviderTiDBCloudNative,
		identityOrg: "60205",
		billingErr:  errors.New("legacy Billing must not be called"),
	}
	billing := &innerBillingFake{plan: &tenant.TiDBCloudOrganizationPlan{OrganizationID: "60205"}}
	srv := NewWithConfig(Config{Provisioner: prov, InnerBilling: billing})
	t.Cleanup(srv.Close)

	profile, err := srv.resolveTiDBCloudAccessProfile(context.Background(), tenant.CredentialProvisionRequest{
		PublicKey: "caller-public", PrivateKey: "caller-private",
	}, "test")
	if err != nil {
		t.Fatalf("resolveTiDBCloudAccessProfile: %v", err)
	}
	if profile.OrganizationID != "60205" || profile.IsFree {
		t.Fatalf("profile = %#v, want non-free organization 60205", profile)
	}
	if prov.iamCalls.Load() != 1 {
		t.Fatalf("legacy IAM calls = %d, want 1", prov.iamCalls.Load())
	}
	if prov.billingCalls.Load() != 0 {
		t.Fatalf("legacy Billing calls = %d, want 0", prov.billingCalls.Load())
	}
	if billing.calls != 1 || billing.organizationID != 60205 {
		t.Fatalf("inner Billing calls = %d, org = %d, want 1/60205", billing.calls, billing.organizationID)
	}
}

func TestResolveTiDBCloudAccessProfileRejectsMismatchedInnerBillingPlan(t *testing.T) {
	billing := &innerBillingFake{plan: &tenant.TiDBCloudOrganizationPlan{OrganizationID: "7"}}
	srv := NewWithConfig(Config{Provisioner: &fakeProvisioner{provider: tenant.ProviderTiDBCloudNative}, InnerBilling: billing})
	t.Cleanup(srv.Close)
	ctx := withInnerAuthSubject(context.Background(), &innerauth.ControlPlaneSubject{OrganizationID: 60205, Kind: innerauth.CredentialJWT})

	_, err := srv.resolveTiDBCloudAccessProfile(ctx, tenant.CredentialProvisionRequest{}, "test")
	if !errors.Is(err, tenant.ErrTiDBCloudBillingResponseInvalid) {
		t.Fatalf("error = %v, want errors.Is(..., %v)", err, tenant.ErrTiDBCloudBillingResponseInvalid)
	}
}

func TestResolveTiDBCloudAccessProfileUsesLegacyBillingWhenInnerBillingIsDisabled(t *testing.T) {
	prov := &fakeProvisioner{provider: tenant.ProviderTiDBCloudNative}
	srv := NewWithConfig(Config{Provisioner: prov})
	t.Cleanup(srv.Close)
	ctx := withInnerAuthSubject(context.Background(), &innerauth.ControlPlaneSubject{OrganizationID: 60205, Kind: innerauth.CredentialJWT})

	profile, err := srv.resolveTiDBCloudAccessProfile(ctx, tenant.CredentialProvisionRequest{}, "test")
	if err != nil {
		t.Fatalf("resolveTiDBCloudAccessProfile: %v", err)
	}
	if profile.OrganizationID != "60205" {
		t.Fatalf("organization = %q, want 60205", profile.OrganizationID)
	}
	if prov.billingCalls.Load() != 1 {
		t.Fatalf("legacy Billing calls = %d, want 1", prov.billingCalls.Load())
	}
}

func TestResolveTiDBCloudAccessProfileUsesLegacyBillingForTypedNilInnerClient(t *testing.T) {
	prov := &fakeProvisioner{
		provider:    tenant.ProviderTiDBCloudNative,
		identityOrg: "60205",
	}
	var billing *innerauth.Client
	srv := NewWithConfig(Config{Provisioner: prov, InnerBilling: billing})
	t.Cleanup(srv.Close)
	ctx := withInnerAuthSubject(context.Background(), &innerauth.ControlPlaneSubject{OrganizationID: 60205, Kind: innerauth.CredentialJWT})

	access, err := srv.resolveTiDBCloudAccessProfile(ctx, tenant.CredentialProvisionRequest{}, "test")
	if err != nil {
		t.Fatalf("resolveTiDBCloudAccessProfile: %v", err)
	}
	if access.OrganizationID != "60205" || access.IsFree || prov.billingCalls.Load() != 1 {
		t.Fatalf("access/legacy calls = %#v/%d, want non-free organization 60205 and one call", access, prov.billingCalls.Load())
	}
}

func TestMapInnerBillingError(t *testing.T) {
	tests := []struct {
		name       string
		err        error
		wantStatus int
		wantErr    error
	}{
		{name: "bad request", err: innerauth.ErrBadRequest, wantStatus: http.StatusBadRequest},
		{name: "unauthenticated", err: innerauth.ErrUnauthenticated, wantErr: tenant.ErrTiDBCloudBillingUnavailable},
		{name: "forbidden", err: innerauth.ErrForbidden, wantErr: tenant.ErrTiDBCloudBillingUnavailable},
		{name: "invalid response", err: innerauth.ErrInvalidResponse, wantErr: tenant.ErrTiDBCloudBillingResponseInvalid},
		{name: "invalid plan", err: tenant.ErrTiDBCloudBillingResponseInvalid, wantErr: tenant.ErrTiDBCloudBillingResponseInvalid},
		{name: "upstream unavailable", err: innerauth.ErrUpstreamUnavailable, wantErr: tenant.ErrTiDBCloudBillingUnavailable},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			err := mapInnerBillingError(tt.err)
			if tt.wantErr != nil {
				if !errors.Is(err, tt.wantErr) {
					t.Fatalf("error = %v, want errors.Is(..., %v)", err, tt.wantErr)
				}
				return
			}
			var apiErr *tenant.TiDBCloudAPIError
			if !errors.As(err, &apiErr) || apiErr.StatusCode != tt.wantStatus || apiErr.Service != tenant.TiDBCloudAPIServiceBilling {
				t.Fatalf("error = %#v, want Billing status %d", err, tt.wantStatus)
			}
		})
	}
}

type innerBillingFake struct {
	plan           *tenant.TiDBCloudOrganizationPlan
	err            error
	calls          int
	organizationID uint64
}

func (f *innerBillingFake) BillingPlan(_ context.Context, organizationID uint64) (*tenant.TiDBCloudOrganizationPlan, error) {
	f.calls++
	f.organizationID = organizationID
	return f.plan, f.err
}