package aws

import (
	"context"
	"errors"
	"fmt"
	"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/cloudfront"
	cftypes "github.com/aws/aws-sdk-go-v2/service/cloudfront/types"
)

func (c *AWSCollector) collectCloudFront(ctx context.Context, result *collector.AssetResult) error {
	if result == nil {
		return errors.New("cloudfront collection result is nil")
	}
	sink := &route53AggregateSink{result: result}
	_, err := c.collectCloudFrontChunks(ctx, sink)
	return err
}

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

func (c *AWSCollector) collectCloudFrontChunks(ctx context.Context, sink collector.ChunkSink) (*collector.CollectionSummary, error) {
	if sink == nil {
		return nil, errors.New("cloudfront chunk sink is nil")
	}

	svc := cloudfront.NewFromConfig(c.cfg)
	summary := &collector.CollectionSummary{CollectedAt: time.Now()}
	paginator := cloudfront.NewListDistributionsPaginator(svc, &cloudfront.ListDistributionsInput{})
	var observations []*cloudFrontDistributionObservation

	for paginator.HasMorePages() {
		page, err := paginator.NextPage(ctx)
		if err != nil {
			return summary, fmt.Errorf("failed to list distributions: %v", err)
		}
		if page.DistributionList == nil || len(page.DistributionList.Items) == 0 {
			continue
		}

		for _, item := range page.DistributionList.Items {
			if err := c.checkExecutionGate(ctx); err != nil {
				return summary, err
			}
			if item.Id == nil {
				continue
			}
			observation, err := c.collectCloudFrontDistributionObservation(ctx, svc, item, summary.CollectedAt)
			if err != nil {
				summary.PartialFailures = append(summary.PartialFailures, collector.CollectionIssue{
					Scope:   "global_service",
					Target:  fmt.Sprintf("cloudfront:%s", aws.ToString(item.Id)),
					Message: err.Error(),
				})
				continue
			}
			observations = append(observations, observation)
		}
	}

	emittedChunks, err := emitCloudFrontObservationChunks(ctx, sink, summary, observations)
	if err != nil {
		return summary, err
	}
	summary.ChunkCount += emittedChunks
	return summary, nil
}

func emitCloudFrontObservationChunks(ctx context.Context, sink collector.ChunkSink, summary *collector.CollectionSummary, observations []*cloudFrontDistributionObservation) (int, error) {
	if sink == nil || summary == nil {
		return 0, nil
	}
	sequence := 0
	for _, observation := range observations {
		if observation == nil {
			continue
		}
		sequence++
		chunk := &collector.AssetChunk{
			Result: collector.AssetResult{
				Assets:       []model.Asset{observation.Asset},
				DomainAssets: observation.DomainAssets,
				Ports:        observation.Ports,
				CollectedAt:  summary.CollectedAt,
			},
			Scope:    "global_service_distribution",
			Target:   fmt.Sprintf("cloudfront/%s", observation.Asset.NativeID),
			Sequence: sequence,
			Metadata: map[string]interface{}{
				"global_service":  "cloudfront",
				"distribution_id": observation.Asset.NativeID,
			},
			Checkpoint: collector.ChunkCheckpoint{},
		}
		if err := sink.PersistChunk(ctx, chunk); err != nil {
			return sequence - 1, err
		}
	}
	checkpointChunk := &collector.AssetChunk{
		Result: collector.AssetResult{
			CollectedAt: summary.CollectedAt,
		},
		Scope:    "global_service_distribution",
		Target:   "cloudfront/finalize",
		Sequence: sequence + 1,
		Metadata: map[string]interface{}{
			"global_service": "cloudfront",
			"finalize_only":  true,
		},
		Checkpoint: collector.ChunkCheckpoint{
			GlobalService: "cloudfront",
			Eligible:      len(summary.PartialFailures) == 0,
		},
	}
	if err := sink.PersistChunk(ctx, checkpointChunk); err != nil {
		return sequence, err
	}
	return sequence + 1, nil
}

func (c *AWSCollector) collectCloudFrontDistributionObservation(ctx context.Context, svc *cloudfront.Client, item cftypes.DistributionSummary, observedAt time.Time) (*cloudFrontDistributionObservation, error) {
	if item.Id == nil {
		return nil, fmt.Errorf("cloudfront distribution missing id")
	}

	distID := aws.ToString(item.Id)
	domainName := aws.ToString(item.DomainName)
	status := "unknown"
	if item.Status != nil {
		status = aws.ToString(item.Status)
	}

	asset := c.newBaseAsset(distID, distID, string(model.AssetTypeNetwork), status, "global")
	c.setService(&asset, "cloudfront")

	tags := c.getCloudFrontTags(ctx, svc, distID)
	asset.Environment = c.extractEnvironment(tags)
	asset.Criticality = c.extractCriticality(tags)

	aliases := make([]string, 0)
	if item.Aliases != nil && item.Aliases.Items != nil {
		aliases = append(aliases, item.Aliases.Items...)
	}

	enabled := false
	if item.Enabled != nil {
		enabled = *item.Enabled
	}

	origins := make([]string, 0)
	if item.Origins != nil && item.Origins.Items != nil {
		for _, origin := range item.Origins.Items {
			if origin.DomainName != nil {
				origins = append(origins, aws.ToString(origin.DomainName))
			}
		}
	}

	asset.Metadata = coremodel.JSONB{
		"domain_name":     domainName,
		"aliases":         aliases,
		"enabled":         enabled,
		"origins":         origins,
		"price_class":     item.PriceClass,
		"http_version":    item.HttpVersion,
		"is_ipv6_enabled": item.IsIPV6Enabled,
	}

	observation := &cloudFrontDistributionObservation{Asset: asset}
	if domainName != "" {
		observation.DomainAssets = append(observation.DomainAssets, newCloudFrontDomainAsset(asset.NativeID, domainName, model.DomainClassServiceHostname, "cloudfront", observedAt))
		observation.Ports = append(observation.Ports, newCloudFrontHTTPPorts(asset.NativeID, domainName)...)
	}
	for _, alias := range aliases {
		if alias == "" {
			continue
		}
		observation.DomainAssets = append(observation.DomainAssets, newCloudFrontDomainAsset(asset.NativeID, alias, model.DomainClassMappedCustom, "custom", observedAt))
		observation.Ports = append(observation.Ports, newCloudFrontHTTPPorts(asset.NativeID, alias)...)
	}
	return observation, nil
}

func newCloudFrontDomainAsset(nativeID, domainName string, class string, registrar string, observedAt time.Time) model.DomainAsset {
	asset := model.Asset{NativeID: nativeID}
	asset.Service = "cloudfront"
	return model.DomainAsset{
		Asset:        asset,
		DomainName:   domainName,
		DomainClass:  class,
		Registrar:    registrar,
		ZoneType:     "public",
		Status:       "active",
		LastSyncedAt: &observedAt,
		FirstSeenAt:  &observedAt,
		LastSeenAt:   &observedAt,
	}
}

func newCloudFrontHTTPPorts(nativeID, hostname string) []model.IPPortIndex {
	return []model.IPPortIndex{
		{
			Asset:             model.Asset{NativeID: nativeID},
			IPAddress:         hostname,
			Port:              443,
			Protocol:          "tcp",
			ObservationSource: model.PortObservationSourceInferredWeb,
			ServiceName:       "https",
			Exposure:          "public",
			Status:            "open",
		},
		{
			Asset:             model.Asset{NativeID: nativeID},
			IPAddress:         hostname,
			Port:              80,
			Protocol:          "tcp",
			ObservationSource: model.PortObservationSourceInferredWeb,
			ServiceName:       "http",
			Exposure:          "public",
			Status:            "open",
		},
	}
}

// getCloudFrontTags fetches tags for a CloudFront distribution
func (c *AWSCollector) getCloudFrontTags(ctx context.Context, svc *cloudfront.Client, distributionID string) map[string]string {
	tags := make(map[string]string)

	// Get distribution config to retrieve ARN
	distConfig, err := svc.GetDistribution(ctx, &cloudfront.GetDistributionInput{
		Id: &distributionID,
	})
	if err != nil {
		eventType := "warning"
		message := fmt.Sprintf("Failed to get CloudFront distribution details: %v", err)
		if downgradedType, downgradedMessage, ok := classifyAWSHelperEvent("cloudfront", "GetDistribution", err); ok {
			eventType = downgradedType
			message = downgradedMessage
		}
		c.report(eventType, message, map[string]interface{}{
			"distribution_id": distributionID,
			"error":           err.Error(),
		})
		return tags
	}

	if distConfig.Distribution == nil || distConfig.Distribution.ARN == nil {
		return tags
	}

	// List tags using ARN
	tagsOutput, err := svc.ListTagsForResource(ctx, &cloudfront.ListTagsForResourceInput{
		Resource: distConfig.Distribution.ARN,
	})
	if err != nil {
		eventType := "warning"
		message := fmt.Sprintf("Failed to retrieve tags for CloudFront distribution: %v", err)
		if downgradedType, downgradedMessage, ok := classifyAWSHelperEvent("cloudfront", "ListTagsForResource", err); ok {
			eventType = downgradedType
			message = downgradedMessage
		}
		c.report(eventType, message, map[string]interface{}{
			"distribution_id": distributionID,
			"arn":             *distConfig.Distribution.ARN,
			"error":           err.Error(),
		})
		return tags
	}

	if tagsOutput.Tags != nil && tagsOutput.Tags.Items != nil {
		for _, tag := range tagsOutput.Tags.Items {
			if tag.Key != nil && tag.Value != nil {
				tags[*tag.Key] = *tag.Value
			}
		}
	}

	return tags
}
