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

func (c *AWSCollector) collectEC2(ctx context.Context, region string, result *collector.AssetResult) error {
	svc := ec2.NewFromConfig(c.cfg, func(o *ec2.Options) {
		o.Region = region
	})

	// First, collect all security groups for reference
	securityGroups, err := c.collectSecurityGroups(ctx, svc)
	if err != nil {
		return fmt.Errorf("failed to collect security groups: %v", err)
	}

	// Then collect instances
	paginator := ec2.NewDescribeInstancesPaginator(svc, &ec2.DescribeInstancesInput{})

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

		for _, reservation := range page.Reservations {
			for _, instance := range reservation.Instances {
				c.processInstance(instance, securityGroups, region, result)
			}
		}
	}
	return nil
}

func (c *AWSCollector) collectSecurityGroups(ctx context.Context, svc *ec2.Client) (map[string]types.SecurityGroup, error) {
	groups := make(map[string]types.SecurityGroup)

	paginator := ec2.NewDescribeSecurityGroupsPaginator(svc, &ec2.DescribeSecurityGroupsInput{})
	for paginator.HasMorePages() {
		page, err := paginator.NextPage(ctx)
		if err != nil {
			return nil, err
		}

		for _, sg := range page.SecurityGroups {
			if sg.GroupId != nil {
				groups[*sg.GroupId] = sg
			}
		}
	}

	return groups, nil
}

func (c *AWSCollector) processInstance(instance types.Instance, securityGroups map[string]types.SecurityGroup, region string, result *collector.AssetResult) {
	name := ""
	for _, tag := range instance.Tags {
		if *tag.Key == "Name" {
			name = *tag.Value
			break
		}
	}
	if name == "" {
		name = *instance.InstanceId
	}

	status := string(instance.State.Name)

	// Create Base Asset
	asset := c.newBaseAsset(*instance.InstanceId, name, string(model.AssetTypeCompute), status, region)
	c.setService(&asset, "ec2")

	// Extract Environment and Criticality from tags
	tags := c.convertTags(instance.Tags)
	asset.Environment = c.extractEnvironment(tags)
	asset.Criticality = c.extractCriticality(tags)

	// Enhanced Metadata with security groups
	sgIDs := make([]string, 0, len(instance.SecurityGroups))
	for _, sg := range instance.SecurityGroups {
		if sg.GroupId != nil {
			sgIDs = append(sgIDs, *sg.GroupId)
		}
	}

	meta := map[string]interface{}{
		"instance_type":   instance.InstanceType,
		"vpc_id":          instance.VpcId,
		"subnet_id":       instance.SubnetId,
		"image_id":        instance.ImageId,
		"tags":            tags,
		"security_groups": sgIDs,
	}
	asset.Metadata = coremodel.JSONB(meta)

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

	// Extract IPs for SG processing
	privateIP := ""
	if instance.PrivateIpAddress != nil {
		privateIP = *instance.PrivateIpAddress
	}
	publicIPv4s, publicIPv6s := instancePublicIPAddresses(instance)
	if privateIP != "" {
		result.IPAssets = append(result.IPAssets, model.IPAsset{
			Asset:     model.Asset{NativeID: *instance.InstanceId},
			IPAddress: privateIP,
			IsPublic:  false,
		})
	}
	publicIPs := make([]string, 0, len(publicIPv4s)+len(publicIPv6s))
	for _, ip := range publicIPv4s {
		publicIPs = appendUniqueString(publicIPs, ip)
	}
	for _, ip := range publicIPv6s {
		publicIPs = appendUniqueString(publicIPs, ip)
	}
	for _, publicIP := range publicIPs {
		if publicIP != privateIP {
			result.IPAssets = append(result.IPAssets, model.IPAsset{
				Asset:     model.Asset{NativeID: *instance.InstanceId},
				IPAddress: publicIP,
				IsPublic:  true,
			})
		}
	}

	// Process Security Group Rules -> Ports
	for _, sgRef := range instance.SecurityGroups {
		if sgRef.GroupId == nil {
			continue
		}

		sg, exists := securityGroups[*sgRef.GroupId]
		if !exists {
			continue
		}

		// Process inbound rules
		for _, rule := range sg.IpPermissions {
			c.processSecurityGroupRule(rule, *sgRef.GroupId, *instance.InstanceId, privateIP, publicIPv4s, publicIPv6s, "inbound", result)
		}

		// Optionally process outbound rules
		// for _, rule := range sg.IpPermissionsEgress {
		//     c.processSecurityGroupRule(rule, *instance.InstanceId, privateIP, publicIPv4s, publicIPv6s, "outbound", result)
		// }
	}
}

func (c *AWSCollector) processSecurityGroupRule(rule types.IpPermission, securityGroupID, instanceID, privateIP string, publicIPv4s, publicIPv6s []string, direction string, result *collector.AssetResult) {
	// Determine if rule allows public access
	isPublicIPv4 := false
	for _, ipRange := range rule.IpRanges {
		if ipRange.CidrIp != nil && *ipRange.CidrIp == "0.0.0.0/0" {
			isPublicIPv4 = true
			break
		}
	}
	isPublicIPv6 := false
	for _, ipv6Range := range rule.Ipv6Ranges {
		if ipv6Range.CidrIpv6 != nil && *ipv6Range.CidrIpv6 == "::/0" {
			isPublicIPv6 = true
			break
		}
	}

	ipAddresses := securityGroupRuleIPAddresses(privateIP, publicIPv4s, publicIPv6s, isPublicIPv4, isPublicIPv6)
	if len(ipAddresses) == 0 {
		return // Skip if no IP available
	}

	// Determine exposure level
	exposure := "private"
	if (isPublicIPv4 && len(publicIPv4s) > 0) || (isPublicIPv6 && len(publicIPv6s) > 0) {
		exposure = "public"
	}

	// Extract port information
	protocol := "tcp"
	if rule.IpProtocol != nil {
		protocol = *rule.IpProtocol
	}
	if protocol == "-1" {
		protocol = "all"
	}

	ports := securityGroupRulePorts(protocol, rule.FromPort, rule.ToPort)
	for _, ipAddress := range ipAddresses {
		for _, port := range ports {
			ruleEvidence := securityGroupRuleEvidence(rule, securityGroupID, direction, protocol, exposure)
			result.Ports = append(result.Ports, model.IPPortIndex{
				Asset:             model.Asset{NativeID: instanceID},
				IPAddress:         ipAddress,
				Port:              port.port,
				Protocol:          protocol,
				ObservationSource: model.PortObservationSourceSecurityGroup,
				ServiceName:       port.serviceName,
				Product:           port.product,
				Exposure:          exposure,
				Status:            "open",
				Evidence:          coremodel.JSONB(ruleEvidence),
				Confidence: coremodel.JSONB{
					"level": "medium",
					"sources": map[string]interface{}{
						"security_group_rule": true,
					},
				},
			})
		}
	}
}

func securityGroupRuleEvidence(rule types.IpPermission, securityGroupID, direction, protocol, exposure string) map[string]interface{} {
	fromPort := int32(0)
	toPort := int32(0)
	if rule.FromPort != nil {
		fromPort = *rule.FromPort
	}
	if rule.ToPort != nil {
		toPort = *rule.ToPort
	}
	cidrs := make([]string, 0, len(rule.IpRanges)+len(rule.Ipv6Ranges))
	for _, ipRange := range rule.IpRanges {
		if ipRange.CidrIp != nil {
			cidrs = appendUniqueString(cidrs, *ipRange.CidrIp)
		}
	}
	for _, ipv6Range := range rule.Ipv6Ranges {
		if ipv6Range.CidrIpv6 != nil {
			cidrs = appendUniqueString(cidrs, *ipv6Range.CidrIpv6)
		}
	}
	sourceGroups := make([]string, 0, len(rule.UserIdGroupPairs))
	for _, pair := range rule.UserIdGroupPairs {
		if pair.GroupId != nil {
			sourceGroups = appendUniqueString(sourceGroups, *pair.GroupId)
		}
	}
	prefixLists := make([]string, 0, len(rule.PrefixListIds))
	for _, prefixList := range rule.PrefixListIds {
		if prefixList.PrefixListId != nil {
			prefixLists = appendUniqueString(prefixLists, *prefixList.PrefixListId)
		}
	}
	source := ""
	switch {
	case len(cidrs) > 0:
		source = cidrs[0]
	case len(sourceGroups) > 0:
		source = sourceGroups[0]
	case len(prefixLists) > 0:
		source = prefixLists[0]
	}
	return map[string]interface{}{
		"sources": map[string]interface{}{
			"security_group_rule": true,
		},
		"security_group_rule": fmt.Sprintf("%s %s %d-%d from %s", direction, protocol, fromPort, toPort, source),
		"security_group_id":   securityGroupID,
		"direction":           direction,
		"protocol":            protocol,
		"from_port":           fromPort,
		"to_port":             toPort,
		"cidr":                source,
		"source":              source,
		"source_cidrs":        cidrs,
		"source_groups":       sourceGroups,
		"prefix_lists":        prefixLists,
		"exposure":            exposure,
	}
}

type securityGroupRulePort struct {
	port        uint32
	serviceName string
	product     string
}

func securityGroupRulePorts(protocol string, fromPort, toPort *int32) []securityGroupRulePort {
	if protocol == "all" {
		return []securityGroupRulePort{{
			port:        0,
			serviceName: "all-protocol",
			product:     "all-protocol",
		}}
	}

	if fromPort == nil || toPort == nil || *fromPort < 0 || *toPort < 0 {
		return []securityGroupRulePort{{
			port:        0,
			serviceName: fmt.Sprintf("all-%s-ports", protocol),
			product:     fmt.Sprintf("%s-all-ports", protocol),
		}}
	}

	from := *fromPort
	to := *toPort
	if to < from {
		from, to = to, from
	}

	fromPortUint := uint32(from)
	toPortUint := uint32(to)
	if toPortUint-fromPortUint <= 10 {
		ports := make([]securityGroupRulePort, 0, toPortUint-fromPortUint+1)
		for port := fromPortUint; port <= toPortUint; port++ {
			ports = append(ports, securityGroupRulePort{
				port:        port,
				serviceName: getServiceName(port),
			})
		}
		return ports
	}

	return []securityGroupRulePort{{
		port:        fromPortUint,
		serviceName: fmt.Sprintf("%s-port-range-%d-%d", protocol, fromPortUint, toPortUint),
		product:     fmt.Sprintf("port_range:%d-%d", fromPortUint, toPortUint),
	}}
}

func securityGroupRuleIPAddresses(privateIP string, publicIPv4s, publicIPv6s []string, isPublicIPv4, isPublicIPv6 bool) []string {
	ips := make([]string, 0, len(publicIPv4s)+len(publicIPv6s)+1)
	if isPublicIPv4 {
		for _, ip := range publicIPv4s {
			ips = appendUniqueString(ips, ip)
		}
	}
	if isPublicIPv6 {
		for _, ip := range publicIPv6s {
			ips = appendUniqueString(ips, ip)
		}
	}
	if len(ips) == 0 && privateIP != "" {
		ips = append(ips, privateIP)
	}
	return ips
}

func instancePublicIPAddresses(instance types.Instance) ([]string, []string) {
	var ipv4s []string
	var ipv6s []string
	if instance.PublicIpAddress != nil {
		ipv4s = appendUniqueString(ipv4s, *instance.PublicIpAddress)
	}
	if instance.Ipv6Address != nil {
		ipv6s = appendUniqueString(ipv6s, *instance.Ipv6Address)
	}

	for _, ni := range instance.NetworkInterfaces {
		if ni.Association != nil && ni.Association.PublicIp != nil {
			ipv4s = appendUniqueString(ipv4s, *ni.Association.PublicIp)
		}
		for _, ipv6 := range ni.Ipv6Addresses {
			if ipv6.Ipv6Address != nil {
				ipv6s = appendUniqueString(ipv6s, *ipv6.Ipv6Address)
			}
		}
	}

	return ipv4s, ipv6s
}

func (c *AWSCollector) convertTags(tags []types.Tag) map[string]string {
	out := make(map[string]string)
	for _, t := range tags {
		if t.Key != nil && t.Value != nil {
			out[*t.Key] = *t.Value
		}
	}
	return out
}

// getServiceName returns common service names for well-known ports
func getServiceName(port uint32) string {
	services := map[uint32]string{
		20:    "ftp-data",
		21:    "ftp",
		22:    "ssh",
		23:    "telnet",
		25:    "smtp",
		53:    "dns",
		80:    "http",
		110:   "pop3",
		143:   "imap",
		443:   "https",
		445:   "smb",
		3306:  "mysql",
		3389:  "rdp",
		5432:  "postgresql",
		5984:  "couchdb",
		6379:  "redis",
		8080:  "http-alt",
		8443:  "https-alt",
		27017: "mongodb",
	}

	if name, ok := services[port]; ok {
		return name
	}
	return fmt.Sprintf("port-%d", port)
}