package aws

import (
	"context"
	"fmt"

	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/elasticloadbalancingv2"
	elbv2types "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2/types"
)

func (c *AWSCollector) collectELB(ctx context.Context, region string, result *collector.AssetResult) error {
	svc := elasticloadbalancingv2.NewFromConfig(c.cfg, func(o *elasticloadbalancingv2.Options) {
		o.Region = region
	})
	paginator := elasticloadbalancingv2.NewDescribeLoadBalancersPaginator(svc, &elasticloadbalancingv2.DescribeLoadBalancersInput{})

	for paginator.HasMorePages() {
		page, err := paginator.NextPage(ctx)
		if err != nil {
			return err
		}

		for _, lb := range page.LoadBalancers {
			name := *lb.LoadBalancerName
			status := string(lb.State.Code)

			// Base Asset
			asset := c.newBaseAsset(*lb.LoadBalancerArn, name, string(model.AssetTypeNetwork), status, region)
			c.setService(&asset, "elb")

			// Extract tags and set Environment/Criticality
			tags := c.getELBTags(ctx, svc, *lb.LoadBalancerArn)
			asset.Environment = c.extractEnvironment(tags)
			asset.Criticality = c.extractCriticality(tags)

			// Collect listeners for this load balancer
			listeners, err := c.collectListeners(ctx, svc, *lb.LoadBalancerArn)
			if err != nil {
				appendAWSPartialFailure(
					result,
					"service",
					fmt.Sprintf("%s:elb:%s:listeners", region, name),
					fmt.Sprintf("failed to collect ELB listeners for %s: %v", name, err),
				)
				c.report("warning", fmt.Sprintf("Failed to collect ELB listeners for %s: %v", name, err), map[string]interface{}{
					"region":            region,
					"load_balancer_arn": *lb.LoadBalancerArn,
					"error":             err.Error(),
				})
				listeners = []string{}
			}

			// Check if internet-facing
			isPublic := lb.Scheme == "internet-facing"

			meta := map[string]interface{}{
				"scheme":             lb.Scheme,
				"type":               lb.Type,
				"vpc_id":             lb.VpcId,
				"dns_name":           lb.DNSName,
				"availability_zones": lb.AvailabilityZones,
				"is_public":          isPublic,
				"listeners":          listeners,
			}
			asset.Metadata = coremodel.JSONB(meta)

			result.Assets = append(result.Assets, asset)

			// Domain Asset (DNS Name)
			if lb.DNSName != nil {
				domainAsset := model.DomainAsset{
					Asset:       model.Asset{NativeID: asset.NativeID},
					DomainName:  *lb.DNSName,
					DomainClass: model.DomainClassServiceHostname,
					Registrar:   "aws",
				}
				result.DomainAssets = append(result.DomainAssets, model.DomainAsset{
					Asset:       domainAsset.Asset,
					DomainName:  domainAsset.DomainName,
					DomainClass: domainAsset.DomainClass,
					Registrar:   domainAsset.Registrar,
				})
			}

			appendELBIPAssets(asset.NativeID, lb, isPublic, result)

			// Process listeners to extract ports
			dnsName := ""
			if lb.DNSName != nil {
				dnsName = *lb.DNSName
			}
			if err := c.processListeners(ctx, svc, *lb.LoadBalancerArn, asset.NativeID, dnsName, isPublic, result); err != nil {
				appendAWSPartialFailure(
					result,
					"service",
					fmt.Sprintf("%s:elb:%s:listener_ports", region, name),
					fmt.Sprintf("failed to process ELB listener ports for %s: %v", name, err),
				)
				c.report("warning", fmt.Sprintf("Failed to process ELB listener ports for %s: %v", name, err), map[string]interface{}{
					"region":            region,
					"load_balancer_arn": *lb.LoadBalancerArn,
					"error":             err.Error(),
				})
			}
		}
	}
	return nil
}

func (c *AWSCollector) collectListeners(ctx context.Context, svc *elasticloadbalancingv2.Client, lbArn string) ([]string, error) {
	paginator := elasticloadbalancingv2.NewDescribeListenersPaginator(svc, &elasticloadbalancingv2.DescribeListenersInput{
		LoadBalancerArn: &lbArn,
	})

	var listeners []string
	for paginator.HasMorePages() {
		resp, err := paginator.NextPage(ctx)
		if err != nil {
			return nil, err
		}
		for _, listener := range resp.Listeners {
			if listener.Port != nil {
				listeners = append(listeners, fmt.Sprintf("%s:%d", listener.Protocol, *listener.Port))
			}
		}
	}

	return listeners, nil
}

func appendELBIPAssets(assetNativeID string, lb elbv2types.LoadBalancer, isPublic bool, result *collector.AssetResult) {
	if result == nil || assetNativeID == "" {
		return
	}

	hostname := ""
	if lb.DNSName != nil {
		hostname = *lb.DNSName
	}
	seen := make(map[string]struct{})
	appendIP := func(ipAddress string, public bool) {
		if ipAddress == "" {
			return
		}
		if _, ok := seen[ipAddress]; ok {
			return
		}
		seen[ipAddress] = struct{}{}
		appendAWSIPAsset(result, assetNativeID, ipAddress, hostname, public)
	}

	for _, zone := range lb.AvailabilityZones {
		for _, address := range zone.LoadBalancerAddresses {
			if address.IpAddress != nil {
				appendIP(*address.IpAddress, isPublic)
			}
			if address.PrivateIPv4Address != nil {
				appendIP(*address.PrivateIPv4Address, false)
			}
			if address.IPv6Address != nil {
				appendIP(*address.IPv6Address, isPublic)
			}
		}
	}
}

func (c *AWSCollector) processListeners(ctx context.Context, svc *elasticloadbalancingv2.Client, lbArn, assetNativeID, dnsName string, isPublic bool, result *collector.AssetResult) error {
	if dnsName == "" {
		return nil
	}
	paginator := elasticloadbalancingv2.NewDescribeListenersPaginator(svc, &elasticloadbalancingv2.DescribeListenersInput{
		LoadBalancerArn: &lbArn,
	})

	exposure := "private"
	if isPublic {
		exposure = "public"
	}

	for paginator.HasMorePages() {
		resp, err := paginator.NextPage(ctx)
		if err != nil {
			return err
		}
		for _, listener := range resp.Listeners {
			if listener.Port == nil {
				continue
			}

			port := uint32(*listener.Port)
			protocol := string(listener.Protocol)

			// Map protocol to lowercase
			protocolLower := "tcp"
			if protocol == "HTTP" || protocol == "HTTPS" {
				protocolLower = "tcp"
			} else if protocol == "UDP" {
				protocolLower = "udp"
			}

			result.Ports = append(result.Ports, model.IPPortIndex{
				Asset:             model.Asset{NativeID: assetNativeID},
				Endpoint:          dnsName,
				IPAddress:         dnsName,
				Port:              port,
				Protocol:          protocolLower,
				ObservationSource: model.PortObservationSourceListener,
				ServiceName:       getServiceName(port),
				Exposure:          exposure,
				Status:            "open",
				Evidence:          directPublicPortEvidence("elasticloadbalancing", assetNativeID, "elbv2_listener", "account_resource"),
				Confidence:        directPublicConfidence("low", "ELB listener on service-managed DNS endpoint"),
			})
		}
	}

	return nil
}

// getELBTags fetches tags for an ELB load balancer
func (c *AWSCollector) getELBTags(ctx context.Context, svc *elasticloadbalancingv2.Client, arn string) map[string]string {
	tags := make(map[string]string)
	out, err := svc.DescribeTags(ctx, &elasticloadbalancingv2.DescribeTagsInput{
		ResourceArns: []string{arn},
	})
	if err != nil {
		eventType := "warning"
		message := fmt.Sprintf("Failed to retrieve tags for ELB: %v", err)
		if downgradedType, downgradedMessage, ok := classifyAWSHelperEvent("elasticloadbalancing", "DescribeTags", err); ok {
			eventType = downgradedType
			message = downgradedMessage
		}
		c.report(eventType, message, map[string]interface{}{
			"arn":   arn,
			"error": err.Error(),
		})
		return tags
	}

	for _, tagDesc := range out.TagDescriptions {
		for _, tag := range tagDesc.Tags {
			if tag.Key != nil && tag.Value != nil {
				tags[*tag.Key] = *tag.Value
			}
		}
	}
	return tags
}