package unity_catalog

import (
	"context"
	"fmt"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/sts"

	"github.com/seaweedfs/seaweedfs/test/testutil"
)

// TestUnityCatalogMasterRoleIntegration covers the master-role configuration
// the upstream playground notes as not-yet-working
// (https://github.com/data-engineering-helpers/mds-in-a-box/blob/main/unitycatalog-playground/etc/conf/server.properties#L45).
//
// It verifies two things end-to-end:
//
//  1. SeaweedFS' STS endpoint accepts sts:AssumeRole for the
//     UnityCatalogVendedRole, returning real STS-vended credentials. This
//     mirrors what the lakekeeper / polaris tests already exercise via the Go
//     AWS SDK and is the SeaweedFS-side prerequisite for UC's master-role
//     flow.
//
//  2. Unity Catalog OSS starts and accepts catalog/schema/EXTERNAL Delta
//     table CRUD when configured with `aws.masterRoleArn` set to that role
//     and AWS_ENDPOINT_URL_STS pointed at SeaweedFS.
//
// What it does NOT verify is the third hop -- UC's Java StsClient
// reaching SeaweedFS' STS handler. UC's
// AwsCredentialGenerator.StsAwsCredentialGenerator builds the StsClient
// without an endpointOverride and doesn't propagate aws.endpoint to it, so
// the SDK's generic AWS_ENDPOINT_URL_STS resolution doesn't kick in and the
// request always targets real AWS, returning InvalidClientTokenId. Verified
// by pointing AWS_ENDPOINT_URL_STS at port 1: same 403 either way, and a
// sniffer in front of SeaweedFS records zero traffic. UC's own tests mock
// StsClient out (BaseCRUDTestWithMockCredentials' EchoAwsStsClient,
// Mockito.mockStatic in CloudCredentialVendorTest), so this never showed
// up upstream. Fix is in unitycatalog/unitycatalog#1532; that bit is
// logged via t.Logf and not asserted.
func TestUnityCatalogMasterRoleIntegration(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping integration test in -short mode")
	}
	if !testutil.HasDocker() {
		t.Skip("docker not available")
	}

	env := newTestEnv(t)
	defer env.cleanup(t)

	t.Log(">>> starting SeaweedFS with STS-enabled IAM config...")
	env.startSeaweedFS(t, stsEnabledIAMConfig(env.accessKey, env.secretKey))
	t.Log(">>> SeaweedFS ready")

	ctx, cancel := context.WithTimeout(context.Background(), 8*time.Minute)
	defer cancel()

	s3c := env.newHostS3Client(t, ctx)
	if _, err := s3c.CreateBucket(ctx, &s3.CreateBucketInput{Bucket: aws.String(ucWarehouse)}); err != nil {
		t.Fatalf("create warehouse bucket: %v", err)
	}

	// Slice 1: prove SeaweedFS STS works for the role UC would use.
	t.Run("SeaweedFsAssumesUnityCatalogVendedRole", func(t *testing.T) {
		stsEndpoint := fmt.Sprintf("http://127.0.0.1:%d", env.s3Port)
		creds, err := assumeRoleViaSeaweedFS(ctx, stsEndpoint, env.accessKey, env.secretKey, ucVendedRoleArn)
		if err != nil {
			t.Fatalf("AssumeRole on SeaweedFS STS: %v", err)
		}
		if creds.SessionToken == "" {
			t.Fatalf("expected non-empty session_token from STS, got %+v", creds)
		}
		if creds.AccessKeyID == env.accessKey {
			t.Fatalf("STS returned the static admin access key; AssumeRole did not vend a new one")
		}

		// Confirm the vended creds actually work for an S3 round-trip.
		s3v := env.newHostS3ClientWithCreds(t, ctx, creds.AccessKeyID, creds.SecretAccessKey, creds.SessionToken)
		probeKey := "ns-probe/master-role-probe.txt"
		if _, err := s3v.PutObject(ctx, &s3.PutObjectInput{
			Bucket: aws.String(ucWarehouse),
			Key:    aws.String(probeKey),
			Body:   nil,
		}); err != nil {
			t.Fatalf("PutObject with sts-vended creds: %v", err)
		}
		_, _ = s3v.DeleteObject(ctx, &s3.DeleteObjectInput{
			Bucket: aws.String(ucWarehouse),
			Key:    aws.String(probeKey),
		})
	})

	// Slice 2: UC starts with master-role config and CRUD works.
	stsEndpoint := fmt.Sprintf("http://host.docker.internal:%d", env.s3Port)
	t.Log(">>> starting Unity Catalog server (master role)...")
	env.startUnityCatalog(t, ctx, ucServerOpts{
		MasterRoleArn: ucVendedRoleArn,
		ExtraEnv: map[string]string{
			"AWS_ENDPOINT_URL":      stsEndpoint,
			"AWS_ENDPOINT_URL_STS":  stsEndpoint,
			"AWS_ACCESS_KEY_ID":     env.accessKey,
			"AWS_SECRET_ACCESS_KEY": env.secretKey,
			"AWS_REGION":            "us-east-1",
		},
	})
	t.Log(">>> Unity Catalog ready")

	uc := newUCClient(fmt.Sprintf("http://127.0.0.1:%d", env.ucHostPort))
	suffix := time.Now().UnixNano()
	catalogName := fmt.Sprintf("seaweed_uc_role_%d", suffix)
	schemaName := fmt.Sprintf("ns_%d", suffix)
	tableName := fmt.Sprintf("events_%d", suffix)
	tableLocation := fmt.Sprintf("s3://%s/%s/%s/%s", ucWarehouse, ucWarehouseKey, schemaName, tableName)

	defer func() {
		_ = uc.deleteTable(context.Background(), catalogName+"."+schemaName+"."+tableName)
		_ = uc.deleteSchema(context.Background(), catalogName+"."+schemaName)
		_ = uc.deleteCatalog(context.Background(), catalogName)
	}()

	t.Run("UnityCatalogAcceptsMasterRoleConfig", func(t *testing.T) {
		if _, err := uc.createCatalog(ctx, ucCreateCatalog{
			Name:        catalogName,
			StorageRoot: fmt.Sprintf("s3://%s/%s", ucWarehouse, ucWarehouseKey),
		}); err != nil {
			t.Fatalf("create catalog: %v", err)
		}
		if _, err := uc.createSchema(ctx, ucCreateSchema{Name: schemaName, CatalogName: catalogName}); err != nil {
			t.Fatalf("create schema: %v", err)
		}
		created, err := uc.createTable(ctx, ucCreateTable{
			Name:             tableName,
			CatalogName:      catalogName,
			SchemaName:       schemaName,
			TableType:        "EXTERNAL",
			DataSourceFormat: "DELTA",
			Columns: []ucColumn{
				{Name: "id", TypeText: "long", TypeName: "LONG", TypeJSON: `{"name":"id","type":"long","nullable":true,"metadata":{}}`, Position: 0, Nullable: true},
			},
			StorageLocation: tableLocation,
		})
		if err != nil {
			t.Fatalf("create external delta table: %v", err)
		}
		if created.TableID == "" {
			t.Fatalf("expected table_id in CreateTable response, got %+v", created)
		}

		// The actual UC -> SeaweedFS STS handoff is the still-unsolved slice.
		// Log what happens, but don't fail the test on it: SeaweedFS' STS is
		// known-good (slice 1 above), and UC's CRUD is known-good (this slice).
		_, ucCredErr := uc.generateTemporaryTableCredentials(ctx, created.TableID, "READ_WRITE")
		if ucCredErr != nil {
			t.Logf("UC StsClient still hits real AWS (UC bug, fix in unitycatalog/unitycatalog#1532): %v", ucCredErr)
		} else {
			t.Logf("UC StsClient honors endpoint override; unitycatalog/unitycatalog#1532 (or equivalent) is in this image.")
		}
	})
}

// assumeRoleViaSeaweedFS calls sts:AssumeRole against SeaweedFS' STS endpoint
// using the Go AWS SDK, mirroring how lakekeeper/polaris tests prove
// SeaweedFS' STS path works.
func assumeRoleViaSeaweedFS(ctx context.Context, endpoint, accessKey, secretKey, roleArn string) (aws.Credentials, error) {
	cfg, err := config.LoadDefaultConfig(ctx,
		config.WithRegion("us-east-1"),
		config.WithCredentialsProvider(credentials.NewStaticCredentialsProvider(accessKey, secretKey, "")),
	)
	if err != nil {
		return aws.Credentials{}, fmt.Errorf("load aws config: %w", err)
	}

	// Override the STS endpoint via service-specific BaseEndpoint instead of
	// the deprecated EndpointResolverWithOptions wired through the global
	// config. This is the supported path in aws-sdk-go-v2 since v1.21.
	client := sts.NewFromConfig(cfg, func(o *sts.Options) {
		o.BaseEndpoint = aws.String(endpoint)
	})
	resp, err := client.AssumeRole(ctx, &sts.AssumeRoleInput{
		RoleArn:         aws.String(roleArn),
		RoleSessionName: aws.String("uc-master-role-test"),
	})
	if err != nil {
		return aws.Credentials{}, err
	}
	if resp.Credentials == nil {
		return aws.Credentials{}, fmt.Errorf("AssumeRole returned no credentials")
	}
	return aws.Credentials{
		AccessKeyID:     aws.ToString(resp.Credentials.AccessKeyId),
		SecretAccessKey: aws.ToString(resp.Credentials.SecretAccessKey),
		SessionToken:    aws.ToString(resp.Credentials.SessionToken),
		Source:          "seaweedfs-sts",
	}, nil
}