package aws

import (
	"context"
	"fmt"
	"slices"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	coremodel "security-center-backend/internal/model"
	"security-center-backend/internal/modules/caasm/collector"
	model "security-center-backend/internal/modules/caasm/model"

	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/s3/types"
	"github.com/aws/smithy-go"
)

const (
	s3ExposureStatePublic            = "public"
	s3ExposureStatePrivate           = "private"
	s3ExposureStateUnknown           = "unknown"
	s3ExposureStatePermissionLimited = "permission_limited"
)

type s3ExposureSignal struct {
	State   string
	Reason  string
	Details map[string]interface{}
}

type s3ExposureAssessment struct {
	State    string
	IsPublic bool
	Reason   string
}

type s3BucketObservation struct {
	Asset        model.Asset
	DomainAssets []model.DomainAsset
	Ports        []model.IPPortIndex
}

func (c *AWSCollector) collectS3(ctx context.Context, result *collector.AssetResult) error {
	sink := &route53AggregateSink{result: result}
	_, err := c.collectS3Chunks(ctx, sink)
	return err
}

func (c *AWSCollector) collectS3Chunks(ctx context.Context, sink collector.ChunkSink) (*collector.CollectionSummary, error) {
	if sink == nil {
		return nil, fmt.Errorf("s3 chunk sink is nil")
	}
	svc := s3.NewFromConfig(c.cfg)
	summary := &collector.CollectionSummary{CollectedAt: time.Now()}

	resp, err := svc.ListBuckets(ctx, &s3.ListBucketsInput{})
	if err != nil {
		return summary, fmt.Errorf("failed to list S3 buckets: %v", err)
	}

	buckets := sortS3Buckets(resp.Buckets)
	for idx, bucket := range buckets {
		if bucket.Name == nil {
			continue
		}
		observation, err := c.collectS3BucketObservation(ctx, svc, bucket)
		if err != nil {
			summary.PartialFailures = append(summary.PartialFailures, collector.CollectionIssue{
				Scope:   "global_service",
				Target:  fmt.Sprintf("s3:%s", aws.ToString(bucket.Name)),
				Message: err.Error(),
			})
			continue
		}

		checkpoint := collector.ChunkCheckpoint{}
		if idx == len(buckets)-1 {
			checkpoint.GlobalService = "s3"
			checkpoint.Eligible = len(summary.PartialFailures) == 0
		}
		chunk := &collector.AssetChunk{
			Result: collector.AssetResult{
				Assets:       []model.Asset{observation.Asset},
				DomainAssets: observation.DomainAssets,
				Ports:        observation.Ports,
				CollectedAt:  time.Now(),
			},
			Scope:    "global_service_bucket",
			Target:   fmt.Sprintf("s3/%s", observation.Asset.NativeID),
			Sequence: idx + 1,
			Metadata: map[string]interface{}{
				"global_service": "s3",
				"bucket_name":    observation.Asset.NativeID,
				"region":         observation.Asset.Region,
			},
			Checkpoint: checkpoint,
		}
		if err := sink.PersistChunk(ctx, chunk); err != nil {
			return summary, err
		}
		summary.ChunkCount++
	}

	if summary.ChunkCount == 0 {
		if err := emitS3FinalizeChunk(ctx, sink, summary, 1); err != nil {
			return summary, err
		}
		summary.ChunkCount++
	}

	return summary, nil
}

func emitS3FinalizeChunk(ctx context.Context, sink collector.ChunkSink, summary *collector.CollectionSummary, sequence int) error {
	if sink == nil || summary == nil {
		return nil
	}
	return sink.PersistChunk(ctx, &collector.AssetChunk{
		Result: collector.AssetResult{
			CollectedAt: summary.CollectedAt,
		},
		Scope:    "global_service_bucket",
		Target:   "s3/finalize",
		Sequence: sequence,
		Metadata: map[string]interface{}{
			"global_service": "s3",
			"finalize_only":  true,
		},
		Checkpoint: collector.ChunkCheckpoint{
			GlobalService: "s3",
			Eligible:      len(summary.PartialFailures) == 0,
		},
	})
}

func sortS3Buckets(buckets []types.Bucket) []types.Bucket {
	sorted := append([]types.Bucket(nil), buckets...)
	slices.SortFunc(sorted, func(a, b types.Bucket) int {
		left := ""
		right := ""
		if a.Name != nil {
			left = strings.ToLower(strings.TrimSpace(*a.Name))
		}
		if b.Name != nil {
			right = strings.ToLower(strings.TrimSpace(*b.Name))
		}
		switch {
		case left < right:
			return -1
		case left > right:
			return 1
		default:
			return 0
		}
	})
	return sorted
}

func (c *AWSCollector) collectS3BucketObservation(ctx context.Context, svc *s3.Client, bucket types.Bucket) (*s3BucketObservation, error) {
	if bucket.Name == nil {
		return nil, fmt.Errorf("s3 bucket missing name")
	}
	bucketName := *bucket.Name
	location, err := c.getBucketLocation(ctx, svc, bucketName)
	if err != nil {
		location = "unknown"
	}
	metadataSvc := c.bucketMetadataS3Client(location)

	publicAccessBlock := c.getBucketPublicAccess(ctx, metadataSvc, bucketName)
	websiteEndpoint := c.getBucketWebsiteEndpoint(ctx, metadataSvc, bucketName, location)
	aclSignal := c.getBucketACL(ctx, metadataSvc, bucketName)
	policySignal := c.checkBucketPolicyPublic(ctx, metadataSvc, bucketName)
	exposure := classifyS3Exposure(publicAccessBlock, aclSignal, policySignal)

	asset := c.newBaseAsset(bucketName, bucketName, string(model.AssetTypeStorage), "active", location)
	c.setService(&asset, "s3")
	tags := c.getS3Tags(ctx, metadataSvc, bucketName)
	asset.Environment = c.extractEnvironment(tags)
	asset.Criticality = c.extractCriticality(tags)
	asset.Metadata = coremodel.JSONB{
		"bucket_name":            bucketName,
		"region":                 location,
		"is_public":              exposure.IsPublic,
		"public_exposure_state":  exposure.State,
		"public_exposure_reason": exposure.Reason,
		"public_access_block":    publicAccessBlock.Details,
		"website_endpoint":       websiteEndpoint,
		"acl_details":            aclSignal.Details,
		"has_public_policy":      policySignal.State == s3ExposureStatePublic,
		"bucket_policy_status": map[string]interface{}{
			"state":  policySignal.State,
			"reason": policySignal.Reason,
		},
		"created_date": bucket.CreationDate,
	}

	return buildS3BucketObservation(asset, location, websiteEndpoint, exposure.IsPublic, time.Now()), nil
}

func buildS3BucketObservation(asset model.Asset, location, websiteEndpoint string, isPublic bool, observedAt time.Time) *s3BucketObservation {
	portExposure := "internal"
	if isPublic {
		portExposure = "public"
	}
	bucketEndpoint := fmt.Sprintf("%s.s3.%s.amazonaws.com", asset.NativeID, location)
	observation := &s3BucketObservation{
		Asset: asset,
		DomainAssets: []model.DomainAsset{
			{
				Asset:        model.Asset{NativeID: asset.NativeID},
				DomainName:   bucketEndpoint,
				DomainClass:  model.DomainClassServiceHostname,
				Registrar:    "s3",
				ZoneType:     "public",
				Status:       "active",
				LastSyncedAt: &observedAt,
				FirstSeenAt:  &observedAt,
				LastSeenAt:   &observedAt,
			},
		},
		Ports: []model.IPPortIndex{
			{
				Asset:             model.Asset{NativeID: asset.NativeID},
				IPAddress:         bucketEndpoint,
				Port:              443,
				Protocol:          "tcp",
				ObservationSource: model.PortObservationSourceInferredWeb,
				ServiceName:       "https",
				Exposure:          portExposure,
				Status:            "open",
			},
		},
	}
	if websiteEndpoint == "" {
		return observation
	}
	observation.DomainAssets = append(observation.DomainAssets, model.DomainAsset{
		Asset:        model.Asset{NativeID: asset.NativeID},
		DomainName:   websiteEndpoint,
		DomainClass:  model.DomainClassServiceHostname,
		Registrar:    "s3-website",
		ZoneType:     "public",
		Status:       "active",
		LastSyncedAt: &observedAt,
		FirstSeenAt:  &observedAt,
		LastSeenAt:   &observedAt,
	})
	observation.Ports = append(observation.Ports, model.IPPortIndex{
		Asset:             model.Asset{NativeID: asset.NativeID},
		IPAddress:         websiteEndpoint,
		Port:              80,
		Protocol:          "tcp",
		ObservationSource: model.PortObservationSourceInferredWeb,
		ServiceName:       "http",
		Exposure:          portExposure,
		Status:            "open",
	})
	return observation
}

func (c *AWSCollector) getBucketLocation(ctx context.Context, svc *s3.Client, bucketName string) (string, error) {
	resp, err := svc.GetBucketLocation(ctx, &s3.GetBucketLocationInput{
		Bucket: &bucketName,
	})
	if err != nil {
		return "", err
	}

	if resp.LocationConstraint == "" {
		return "us-east-1", nil // Default region
	}

	return normalizeS3BucketRegion(string(resp.LocationConstraint)), nil
}

func normalizeS3BucketRegion(region string) string {
	normalized := strings.ToLower(strings.TrimSpace(region))
	switch normalized {
	case "eu":
		return "eu-west-1"
	default:
		return normalized
	}
}

func (c *AWSCollector) bucketMetadataS3Config(location string) aws.Config {
	cfg := c.cfg
	region := normalizeS3BucketRegion(location)
	if region == "" || region == "unknown" {
		return cfg
	}
	cfg.Region = region
	return cfg
}

func (c *AWSCollector) bucketMetadataS3Client(location string) *s3.Client {
	return s3.NewFromConfig(c.bucketMetadataS3Config(location))
}

func (c *AWSCollector) getBucketPublicAccess(ctx context.Context, svc *s3.Client, bucketName string) s3ExposureSignal {
	resp, err := svc.GetPublicAccessBlock(ctx, &s3.GetPublicAccessBlockInput{
		Bucket: &bucketName,
	})

	details := map[string]interface{}{
		"state":                   s3ExposureStateUnknown,
		"block_public_acls":       false,
		"ignore_public_acls":      false,
		"block_public_policy":     false,
		"restrict_public_buckets": false,
	}

	if err != nil {
		state := classifyS3ReadError(err)
		details["state"] = state
		details["error"] = err.Error()
		return s3ExposureSignal{
			State:   state,
			Reason:  "get_public_access_block_failed",
			Details: details,
		}
	}

	if resp.PublicAccessBlockConfiguration == nil {
		details["state"] = s3ExposureStateUnknown
		return s3ExposureSignal{
			State:   s3ExposureStateUnknown,
			Reason:  "public_access_block_not_configured",
			Details: details,
		}
	}

	config := resp.PublicAccessBlockConfiguration
	details["block_public_acls"] = config.BlockPublicAcls != nil && *config.BlockPublicAcls
	details["ignore_public_acls"] = config.IgnorePublicAcls != nil && *config.IgnorePublicAcls
	details["block_public_policy"] = config.BlockPublicPolicy != nil && *config.BlockPublicPolicy
	details["restrict_public_buckets"] = config.RestrictPublicBuckets != nil && *config.RestrictPublicBuckets

	// Public access block is only one signal. We record whether it fully blocks public access,
	// but we do not treat a missing or partially open config as confirmed public on its own.
	allBlocked := details["block_public_acls"].(bool) &&
		details["ignore_public_acls"].(bool) &&
		details["block_public_policy"].(bool) &&
		details["restrict_public_buckets"].(bool)

	state := s3ExposureStateUnknown
	reason := "public_access_block_partial_or_open"
	if allBlocked {
		state = s3ExposureStatePrivate
		reason = "public_access_block_all_enabled"
	}
	details["state"] = state
	return s3ExposureSignal{
		State:   state,
		Reason:  reason,
		Details: details,
	}
}

func (c *AWSCollector) getBucketWebsiteEndpoint(ctx context.Context, svc *s3.Client, bucketName, region string) string {
	_, err := svc.GetBucketWebsite(ctx, &s3.GetBucketWebsiteInput{
		Bucket: &bucketName,
	})

	if err != nil {
		// Website hosting not enabled
		return ""
	}

	// Website endpoint format: bucket-name.s3-website-region.amazonaws.com
	// or bucket-name.s3-website.region.amazonaws.com (depending on region)
	if region == "us-east-1" {
		return fmt.Sprintf("%s.s3-website-us-east-1.amazonaws.com", bucketName)
	}
	return fmt.Sprintf("%s.s3-website.%s.amazonaws.com", bucketName, region)
}

func (c *AWSCollector) getBucketACL(ctx context.Context, svc *s3.Client, bucketName string) s3ExposureSignal {
	resp, err := svc.GetBucketAcl(ctx, &s3.GetBucketAclInput{
		Bucket: &bucketName,
	})

	details := map[string]interface{}{
		"state":             s3ExposureStateUnknown,
		"has_public_grants": false,
		"grant_count":       0,
	}

	if err != nil {
		state := classifyS3ReadError(err)
		details["state"] = state
		details["error"] = err.Error()
		return s3ExposureSignal{
			State:   state,
			Reason:  "get_bucket_acl_failed",
			Details: details,
		}
	}

	details["grant_count"] = len(resp.Grants)

	// Check for public grants
	for _, grant := range resp.Grants {
		if grant.Grantee != nil && grant.Grantee.Type == types.TypeGroup {
			if grant.Grantee.URI != nil {
				uri := *grant.Grantee.URI
				// Check for public URIs
				if uri == "http://acs.amazonaws.com/groups/global/AllUsers" ||
					uri == "http://acs.amazonaws.com/groups/global/AuthenticatedUsers" {
					details["has_public_grants"] = true
					break
				}
			}
		}
	}

	state := s3ExposureStatePrivate
	reason := "acl_not_public"
	if details["has_public_grants"].(bool) {
		state = s3ExposureStatePublic
		reason = "acl_public_grant"
	}
	details["state"] = state
	return s3ExposureSignal{
		State:   state,
		Reason:  reason,
		Details: details,
	}
}

func (c *AWSCollector) checkBucketPolicyPublic(ctx context.Context, svc *s3.Client, bucketName string) s3ExposureSignal {
	// Get bucket policy status (simpler than parsing the policy)
	resp, err := svc.GetBucketPolicyStatus(ctx, &s3.GetBucketPolicyStatusInput{
		Bucket: &bucketName,
	})

	if err != nil {
		state := classifyS3ReadError(err)
		if state == s3ExposureStateUnknown {
			var apiErr smithy.APIError
			if ok := errorAs(err, &apiErr); ok {
				code := strings.ToLower(apiErr.ErrorCode())
				if code == "nosuchbucketpolicy" || code == "nosuchbucketpolicystatus" {
					state = s3ExposureStatePrivate
				}
			}
		}
		return s3ExposureSignal{
			State:  state,
			Reason: "get_bucket_policy_status_failed",
			Details: map[string]interface{}{
				"state": state,
				"error": err.Error(),
			},
		}
	}

	if resp.PolicyStatus == nil || resp.PolicyStatus.IsPublic == nil {
		return s3ExposureSignal{
			State:  s3ExposureStateUnknown,
			Reason: "bucket_policy_status_missing",
			Details: map[string]interface{}{
				"state": s3ExposureStateUnknown,
			},
		}
	}

	state := s3ExposureStatePrivate
	reason := "bucket_policy_not_public"
	if *resp.PolicyStatus.IsPublic {
		state = s3ExposureStatePublic
		reason = "bucket_policy_public"
	}
	return s3ExposureSignal{
		State:  state,
		Reason: reason,
		Details: map[string]interface{}{
			"state":     state,
			"is_public": *resp.PolicyStatus.IsPublic,
		},
	}
}

func classifyS3Exposure(publicAccessBlock, aclSignal, policySignal s3ExposureSignal) s3ExposureAssessment {
	signals := []s3ExposureSignal{publicAccessBlock, aclSignal, policySignal}
	for _, signal := range signals {
		if signal.State == s3ExposureStatePublic {
			return s3ExposureAssessment{
				State:    s3ExposureStatePublic,
				IsPublic: true,
				Reason:   signal.Reason,
			}
		}
	}
	for _, signal := range signals {
		if signal.State == s3ExposureStatePermissionLimited {
			return s3ExposureAssessment{
				State:    s3ExposureStatePermissionLimited,
				IsPublic: false,
				Reason:   signal.Reason,
			}
		}
	}
	for _, signal := range signals {
		if signal.State == s3ExposureStateUnknown {
			return s3ExposureAssessment{
				State:    s3ExposureStateUnknown,
				IsPublic: false,
				Reason:   signal.Reason,
			}
		}
	}
	return s3ExposureAssessment{
		State:    s3ExposureStatePrivate,
		IsPublic: false,
		Reason:   "no_public_signal_detected",
	}
}

func classifyS3ReadError(err error) string {
	var apiErr smithy.APIError
	if !errorAs(err, &apiErr) {
		return s3ExposureStateUnknown
	}

	code := strings.ToLower(apiErr.ErrorCode())
	if strings.Contains(code, "accessdenied") || strings.Contains(code, "forbidden") || strings.Contains(code, "unauthorized") {
		return s3ExposureStatePermissionLimited
	}
	return s3ExposureStateUnknown
}

func errorAs(err error, target interface{}) bool {
	switch t := target.(type) {
	case *smithy.APIError:
		apiErr, ok := err.(smithy.APIError)
		if !ok {
			return false
		}
		*t = apiErr
		return true
	default:
		return false
	}
}

// getS3Tags fetches tags for an S3 bucket
func (c *AWSCollector) getS3Tags(ctx context.Context, svc *s3.Client, bucketName string) map[string]string {
	tags := make(map[string]string)
	resp, err := svc.GetBucketTagging(ctx, &s3.GetBucketTaggingInput{
		Bucket: &bucketName,
	})
	if err != nil {
		// Check if it's "NoSuchTagSet" error (bucket has no tags) - this is not an error
		if !strings.Contains(err.Error(), "NoSuchTagSet") {
			c.report("warning", fmt.Sprintf("Failed to retrieve tags for S3 bucket %s: %v", bucketName, err), map[string]interface{}{
				"bucket_name": bucketName,
				"error":       err.Error(),
			})
		}
		return tags
	}

	for _, tag := range resp.TagSet {
		if tag.Key != nil && tag.Value != nil {
			tags[*tag.Key] = *tag.Value
		}
	}
	return tags
}