package aws

import (
	"context"
	"errors"
	"fmt"
	"net"

	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) collectVPC(ctx context.Context, region string, result *collector.AssetResult) error {
	svc := ec2.NewFromConfig(c.cfg, func(o *ec2.Options) {
		o.Region = region
	})

	vpcPaginator := ec2.NewDescribeVpcsPaginator(svc, &ec2.DescribeVpcsInput{})
	var errs []error
	for vpcPaginator.HasMorePages() {
		vpcResp, err := vpcPaginator.NextPage(ctx)
		if err != nil {
			return fmt.Errorf("failed to describe VPCs: %v", err)
		}
		for _, vpc := range vpcResp.Vpcs {
			if vpc.VpcId == nil {
				continue
			}

			vpcID := *vpc.VpcId
			name := getVPCName(vpc.Tags)
			if name == "" {
				name = vpcID
			}

			// Create VPC asset
			asset := c.newBaseAsset(vpcID, name, string(model.AssetTypeNetwork), "active", region)
			c.setService(&asset, "vpc")

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

			meta := map[string]interface{}{
				"cidr_block":       vpc.CidrBlock,
				"is_default":       vpc.IsDefault,
				"instance_tenancy": vpc.InstanceTenancy,
				"tags":             convertEC2Tags(vpc.Tags),
			}
			asset.Metadata = coremodel.JSONB(meta)
			result.Assets = append(result.Assets, asset)

			// Collect subnets for this VPC
			if err := c.collectSubnets(ctx, svc, vpcID, region, result); err != nil {
				errs = append(errs, fmt.Errorf("subnets for %s: %w", vpcID, err))
			}
		}
	}

	// Collect NAT Gateways
	if err := c.collectNATGateways(ctx, svc, region, result); err != nil {
		errs = append(errs, fmt.Errorf("nat gateways: %w", err))
	}

	// Collect VPC Endpoints
	if err := c.collectVPCEndpoints(ctx, svc, region, result); err != nil {
		errs = append(errs, fmt.Errorf("vpc endpoints: %w", err))
	}

	// Collect standalone ENI and EIP IP assets. These are asset-center targets even when
	// they are not attached to an EC2 instance observed through DescribeInstances.
	if err := c.collectNetworkInterfaces(ctx, svc, region, result); err != nil {
		errs = append(errs, fmt.Errorf("network interfaces: %w", err))
	}
	if err := c.collectElasticIPs(ctx, svc, region, result); err != nil {
		errs = append(errs, fmt.Errorf("elastic ips: %w", err))
	}

	return errors.Join(errs...)
}

func (c *AWSCollector) collectSubnets(ctx context.Context, svc *ec2.Client, vpcID, region string, result *collector.AssetResult) error {
	paginator := ec2.NewDescribeSubnetsPaginator(svc, &ec2.DescribeSubnetsInput{
		Filters: []types.Filter{
			{
				Name:   stringPtr("vpc-id"),
				Values: []string{vpcID},
			},
		},
	})

	for paginator.HasMorePages() {
		resp, err := paginator.NextPage(ctx)
		if err != nil {
			return err
		}
		for _, subnet := range resp.Subnets {
			if subnet.SubnetId == nil {
				continue
			}

			subnetID := *subnet.SubnetId
			name := getNameFromTags(subnet.Tags)
			if name == "" {
				name = subnetID
			}

			isPublic := subnet.MapPublicIpOnLaunch != nil && *subnet.MapPublicIpOnLaunch

			// Create Subnet asset
			asset := c.newBaseAsset(subnetID, name, string(model.AssetTypeNetwork), "active", region)
			c.setService(&asset, "vpc")

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

			meta := map[string]interface{}{
				"vpc_id":                  vpcID,
				"cidr_block":              subnet.CidrBlock,
				"availability_zone":       subnet.AvailabilityZone,
				"available_ip_count":      subnet.AvailableIpAddressCount,
				"is_public":               isPublic,
				"map_public_ip_on_launch": subnet.MapPublicIpOnLaunch,
				"tags":                    convertEC2Tags(subnet.Tags),
			}
			asset.Metadata = coremodel.JSONB(meta)
			result.Assets = append(result.Assets, asset)
		}
	}

	return nil
}

func (c *AWSCollector) collectNATGateways(ctx context.Context, svc *ec2.Client, region string, result *collector.AssetResult) error {
	paginator := ec2.NewDescribeNatGatewaysPaginator(svc, &ec2.DescribeNatGatewaysInput{})

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

		for _, nat := range page.NatGateways {
			if nat.NatGatewayId == nil {
				continue
			}

			natID := *nat.NatGatewayId
			name := getNameFromTags(nat.Tags)
			if name == "" {
				name = natID
			}

			status := string(nat.State)

			// Create NAT Gateway asset
			asset := c.newBaseAsset(natID, name, string(model.AssetTypeNetwork), status, region)
			c.setService(&asset, "vpc")

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

			meta := map[string]interface{}{
				"vpc_id":            nat.VpcId,
				"subnet_id":         nat.SubnetId,
				"connectivity_type": nat.ConnectivityType,
				"tags":              convertEC2Tags(nat.Tags),
			}
			asset.Metadata = coremodel.JSONB(meta)
			result.Assets = append(result.Assets, asset)

			appendNATGatewayIPAssets(result, asset.NativeID, nat)
		}
	}

	return nil
}

func (c *AWSCollector) collectVPCEndpoints(ctx context.Context, svc *ec2.Client, region string, result *collector.AssetResult) error {
	paginator := ec2.NewDescribeVpcEndpointsPaginator(svc, &ec2.DescribeVpcEndpointsInput{})

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

		for _, endpoint := range page.VpcEndpoints {
			if endpoint.VpcEndpointId == nil {
				continue
			}

			endpointID := *endpoint.VpcEndpointId
			name := getNameFromTags(endpoint.Tags)
			if name == "" {
				name = endpointID
			}

			status := string(endpoint.State)

			// Create VPC Endpoint asset
			asset := c.newBaseAsset(endpointID, name, string(model.AssetTypeNetwork), status, region)
			c.setService(&asset, "vpc")

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

			meta := map[string]interface{}{
				"vpc_id":                endpoint.VpcId,
				"service_name":          endpoint.ServiceName,
				"endpoint_type":         endpoint.VpcEndpointType,
				"network_interface_ids": endpoint.NetworkInterfaceIds,
				"private_dns_enabled":   endpoint.PrivateDnsEnabled,
				"tags":                  convertEC2Tags(endpoint.Tags),
			}
			asset.Metadata = coremodel.JSONB(meta)
			result.Assets = append(result.Assets, asset)

			if len(endpoint.NetworkInterfaceIds) > 0 {
				if err := c.collectNetworkInterfaceIPAssets(ctx, svc, region, asset.NativeID, endpoint.NetworkInterfaceIds, result); err != nil {
					return fmt.Errorf("vpc endpoint %s network interfaces: %w", endpointID, err)
				}
			}

			// Extract DNS entries
			for _, dns := range endpoint.DnsEntries {
				if dns.DnsName != nil {
					dnsName := *dns.DnsName
					domainAsset := model.DomainAsset{
						Asset:       model.Asset{NativeID: asset.NativeID},
						DomainName:  dnsName,
						DomainClass: model.DomainClassServiceHostname,
						Registrar:   "aws",
					}
					setServiceEndpointDomainDefaults(&domainAsset, "vpc", false)
					c.setService(&domainAsset.Asset, "vpc")
					result.DomainAssets = append(result.DomainAssets, domainAsset)
				}
			}
		}
	}

	return nil
}

func (c *AWSCollector) collectNetworkInterfaces(ctx context.Context, svc *ec2.Client, region string, result *collector.AssetResult) error {
	paginator := ec2.NewDescribeNetworkInterfacesPaginator(svc, &ec2.DescribeNetworkInterfacesInput{})
	for paginator.HasMorePages() {
		page, err := paginator.NextPage(ctx)
		if err != nil {
			return err
		}
		for _, eni := range page.NetworkInterfaces {
			if eni.NetworkInterfaceId == nil {
				continue
			}
			eniID := *eni.NetworkInterfaceId
			name := getNameFromTags(eni.TagSet)
			if name == "" {
				name = eniID
			}
			status := string(eni.Status)
			asset := c.newBaseAsset(eniID, name, string(model.AssetTypeNetwork), status, region)
			c.setService(&asset, "ec2")
			tags := convertEC2Tags(eni.TagSet)
			asset.Environment = c.extractEnvironment(tags)
			asset.Criticality = c.extractCriticality(tags)
			meta := map[string]interface{}{
				"vpc_id":                eni.VpcId,
				"subnet_id":             eni.SubnetId,
				"description":           eni.Description,
				"interface_type":        eni.InterfaceType,
				"availability_zone":     eni.AvailabilityZone,
				"private_ip_address":    eni.PrivateIpAddress,
				"association_public_ip": "",
				"tags":                  tags,
			}
			if eni.Association != nil && eni.Association.PublicIp != nil {
				meta["association_public_ip"] = *eni.Association.PublicIp
			}
			if eni.Attachment != nil {
				meta["attachment_instance_id"] = eni.Attachment.InstanceId
			}
			asset.Metadata = coremodel.JSONB(meta)
			result.Assets = append(result.Assets, asset)
			appendNetworkInterfaceIPAssets(result, asset.NativeID, eni, "")
		}
	}
	return nil
}

func (c *AWSCollector) collectNetworkInterfaceIPAssets(ctx context.Context, svc *ec2.Client, region string, ownerNativeID string, eniIDs []string, result *collector.AssetResult) error {
	if ownerNativeID == "" || len(eniIDs) == 0 {
		return nil
	}
	for start := 0; start < len(eniIDs); start += 100 {
		end := start + 100
		if end > len(eniIDs) {
			end = len(eniIDs)
		}
		resp, err := svc.DescribeNetworkInterfaces(ctx, &ec2.DescribeNetworkInterfacesInput{NetworkInterfaceIds: eniIDs[start:end]})
		if err != nil {
			return err
		}
		for _, eni := range resp.NetworkInterfaces {
			appendNetworkInterfaceIPAssets(result, ownerNativeID, eni, "")
		}
	}
	_ = region
	return nil
}

func (c *AWSCollector) collectElasticIPs(ctx context.Context, svc *ec2.Client, region string, result *collector.AssetResult) error {
	page, err := svc.DescribeAddresses(ctx, &ec2.DescribeAddressesInput{})
	if err != nil {
		return err
	}
	for _, address := range page.Addresses {
		nativeID := ""
		if address.AllocationId != nil {
			nativeID = *address.AllocationId
		} else if address.PublicIp != nil {
			nativeID = "eip:" + *address.PublicIp
		}
		if nativeID == "" {
			continue
		}
		name := nativeID
		asset := c.newBaseAsset(nativeID, name, string(model.AssetTypeNetwork), "active", region)
		c.setService(&asset, "ec2")
		meta := map[string]interface{}{
			"allocation_id":        address.AllocationId,
			"association_id":       address.AssociationId,
			"domain":               address.Domain,
			"instance_id":          address.InstanceId,
			"network_interface_id": address.NetworkInterfaceId,
			"private_ip_address":   address.PrivateIpAddress,
			"public_ip":            address.PublicIp,
			"public_ipv4_pool":     address.PublicIpv4Pool,
			"network_border_group": address.NetworkBorderGroup,
		}
		asset.Metadata = coremodel.JSONB(meta)
		result.Assets = append(result.Assets, asset)
		if address.PublicIp != nil {
			appendAWSIPAsset(result, asset.NativeID, *address.PublicIp, "", true)
		}
		if address.PrivateIpAddress != nil {
			appendAWSIPAsset(result, asset.NativeID, *address.PrivateIpAddress, "", false)
		}
		if address.NetworkInterfaceId != nil {
			if address.PublicIp != nil {
				appendAWSIPAsset(result, *address.NetworkInterfaceId, *address.PublicIp, "", true)
			}
			if address.PrivateIpAddress != nil {
				appendAWSIPAsset(result, *address.NetworkInterfaceId, *address.PrivateIpAddress, "", false)
			}
		}
	}
	return nil
}

func appendNATGatewayIPAssets(result *collector.AssetResult, assetNativeID string, nat types.NatGateway) {
	for _, addr := range nat.NatGatewayAddresses {
		if addr.PublicIp != nil {
			appendAWSIPAsset(result, assetNativeID, *addr.PublicIp, "", true)
		}
		if addr.PrivateIp != nil {
			appendAWSIPAsset(result, assetNativeID, *addr.PrivateIp, "", false)
		}
		if addr.NetworkInterfaceId != nil {
			if addr.PublicIp != nil {
				appendAWSIPAsset(result, *addr.NetworkInterfaceId, *addr.PublicIp, "", true)
			}
			if addr.PrivateIp != nil {
				appendAWSIPAsset(result, *addr.NetworkInterfaceId, *addr.PrivateIp, "", false)
			}
		}
	}
}

func appendNetworkInterfaceIPAssets(result *collector.AssetResult, assetNativeID string, eni types.NetworkInterface, hostname string) {
	if eni.PrivateIpAddress != nil {
		appendAWSIPAsset(result, assetNativeID, *eni.PrivateIpAddress, hostname, false)
	}
	if eni.Association != nil && eni.Association.PublicIp != nil {
		appendAWSIPAsset(result, assetNativeID, *eni.Association.PublicIp, hostname, true)
	}
	for _, private := range eni.PrivateIpAddresses {
		if private.PrivateIpAddress != nil {
			appendAWSIPAsset(result, assetNativeID, *private.PrivateIpAddress, hostname, false)
		}
		if private.Association != nil && private.Association.PublicIp != nil {
			appendAWSIPAsset(result, assetNativeID, *private.Association.PublicIp, hostname, true)
		}
	}
	for _, ipv6 := range eni.Ipv6Addresses {
		if ipv6.Ipv6Address != nil {
			appendAWSIPAsset(result, assetNativeID, *ipv6.Ipv6Address, hostname, true)
		}
	}
}

func appendAWSIPAsset(result *collector.AssetResult, assetNativeID, ipAddress, hostname string, isPublic bool) {
	if result == nil || assetNativeID == "" || net.ParseIP(ipAddress) == nil {
		return
	}
	result.IPAssets = append(result.IPAssets, model.IPAsset{
		Asset:     model.Asset{NativeID: assetNativeID},
		IPAddress: ipAddress,
		Hostname:  hostname,
		IsPublic:  isPublic,
		GeoLocation: coremodel.JSONB{
			"source_service":     inferAWSIPAssetSource(assetNativeID),
			"resource_id":        assetNativeID,
			"owning_resource_id": assetNativeID,
			"ownership_status":   "attributed",
		},
	})
}

// Helper functions

func getVPCName(tags []types.Tag) string {
	return getNameFromTags(tags)
}

func getNameFromTags(tags []types.Tag) string {
	for _, tag := range tags {
		if tag.Key != nil && *tag.Key == "Name" && tag.Value != nil {
			return *tag.Value
		}
	}
	return ""
}

func convertEC2Tags(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
}

func stringPtr(s string) *string {
	return &s
}