package aws

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

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/lightsail"
	lstypes "github.com/aws/aws-sdk-go-v2/service/lightsail/types"

	coremodel "security-center-backend/internal/model"
	"security-center-backend/internal/modules/caasm/collector"
	model "security-center-backend/internal/modules/caasm/model"
)

type lightsailAPI interface {
	GetInstances(context.Context, *lightsail.GetInstancesInput, ...func(*lightsail.Options)) (*lightsail.GetInstancesOutput, error)
	GetStaticIps(context.Context, *lightsail.GetStaticIpsInput, ...func(*lightsail.Options)) (*lightsail.GetStaticIpsOutput, error)
	GetLoadBalancers(context.Context, *lightsail.GetLoadBalancersInput, ...func(*lightsail.Options)) (*lightsail.GetLoadBalancersOutput, error)
}

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

func (c *AWSCollector) collectLightsailWithClient(ctx context.Context, region string, svc lightsailAPI, result *collector.AssetResult) error {
	if result == nil || svc == nil {
		return nil
	}
	if err := c.collectLightsailInstances(ctx, region, svc, result); err != nil {
		appendAWSPartialFailure(result, "service", fmt.Sprintf("%s:lightsail:instances", region), fmt.Sprintf("failed to collect Lightsail instances: %v", err))
	}
	if err := c.collectLightsailStaticIPs(ctx, region, svc, result); err != nil {
		appendAWSPartialFailure(result, "service", fmt.Sprintf("%s:lightsail:static_ips", region), fmt.Sprintf("failed to collect Lightsail static IPs: %v", err))
	}
	if err := c.collectLightsailLoadBalancers(ctx, region, svc, result); err != nil {
		appendAWSPartialFailure(result, "service", fmt.Sprintf("%s:lightsail:load_balancers", region), fmt.Sprintf("failed to collect Lightsail load balancers: %v", err))
	}
	return nil
}

func (c *AWSCollector) collectLightsailInstances(ctx context.Context, region string, svc lightsailAPI, result *collector.AssetResult) error {
	var pageToken *string
	for {
		page, err := svc.GetInstances(ctx, &lightsail.GetInstancesInput{PageToken: pageToken})
		if err != nil {
			return err
		}
		for _, instance := range page.Instances {
			if err := c.checkExecutionGate(ctx); err != nil {
				return err
			}
			c.appendLightsailInstance(region, instance, result)
		}
		if page.NextPageToken == nil || aws.ToString(page.NextPageToken) == "" {
			break
		}
		pageToken = page.NextPageToken
	}
	return nil
}

func (c *AWSCollector) appendLightsailInstance(region string, instance lstypes.Instance, result *collector.AssetResult) {
	nativeID := aws.ToString(instance.Arn)
	if nativeID == "" {
		nativeID = aws.ToString(instance.Name)
	}
	if nativeID == "" {
		return
	}
	name := aws.ToString(instance.Name)
	if name == "" {
		name = nativeID
	}
	status := "unknown"
	if instance.State != nil && instance.State.Name != nil {
		status = aws.ToString(instance.State.Name)
	}

	asset := c.newBaseAsset(nativeID, name, string(model.AssetTypeCompute), status, lightsailRegion(region, instance.Location))
	c.setService(&asset, "lightsail")
	tags := convertLightsailTags(instance.Tags)
	asset.Environment = c.extractEnvironment(tags)
	asset.Criticality = c.extractCriticality(tags)
	asset.Metadata = coremodel.JSONB{
		"blueprint_id":       instance.BlueprintId,
		"bundle_id":          instance.BundleId,
		"ip_address_type":    instance.IpAddressType,
		"is_static_ip":       instance.IsStaticIp,
		"private_ip_address": instance.PrivateIpAddress,
		"public_ip_address":  instance.PublicIpAddress,
		"ipv6_addresses":     instance.Ipv6Addresses,
		"availability_zone":  lightsailAvailabilityZone(instance.Location),
		"resource_type":      instance.ResourceType,
		"networking_ports":   lightsailInstancePortMetadata(instance),
		"tags":               tags,
	}
	result.Assets = append(result.Assets, asset)
	if instance.PrivateIpAddress != nil {
		appendAWSIPAsset(result, asset.NativeID, aws.ToString(instance.PrivateIpAddress), "", false)
	}
	if instance.PublicIpAddress != nil {
		appendAWSIPAsset(result, asset.NativeID, aws.ToString(instance.PublicIpAddress), "", true)
	}
	for _, ip := range instance.Ipv6Addresses {
		appendAWSIPAsset(result, asset.NativeID, ip, "", true)
	}
	c.appendLightsailInstancePorts(asset.NativeID, aws.ToString(instance.PublicIpAddress), instance.Ipv6Addresses, instance.Networking, result)
}

func (c *AWSCollector) appendLightsailInstancePorts(nativeID, publicIP string, ipv6Addresses []string, networking *lstypes.InstanceNetworking, result *collector.AssetResult) {
	if networking == nil {
		return
	}
	endpoints := make([]string, 0, 1+len(ipv6Addresses))
	if publicIP != "" {
		endpoints = append(endpoints, publicIP)
	}
	endpoints = append(endpoints, ipv6Addresses...)
	if len(endpoints) == 0 {
		return
	}
	for _, portInfo := range networking.Ports {
		if portInfo.AccessType != lstypes.PortAccessTypePublic {
			continue
		}
		protocol := strings.ToLower(string(portInfo.Protocol))
		for _, port := range lightsailPortInfoPorts(portInfo) {
			for _, endpoint := range endpoints {
				record := directPublicEndpointPort(nativeID, endpoint, port, protocol, model.PortObservationSourceServiceEndpoint, "public", getServiceName(port), directPublicPortEvidence("lightsail", nativeID, "lightsail_instance_networking", "account_resource"))
				record.IPAddress = endpoint
				record.Confidence = directPublicConfidence("medium", "Lightsail public networking port on account-owned instance IP")
				result.Ports = append(result.Ports, record)
			}
		}
	}
}

func (c *AWSCollector) collectLightsailStaticIPs(ctx context.Context, region string, svc lightsailAPI, result *collector.AssetResult) error {
	var pageToken *string
	for {
		page, err := svc.GetStaticIps(ctx, &lightsail.GetStaticIpsInput{PageToken: pageToken})
		if err != nil {
			return err
		}
		for _, staticIP := range page.StaticIps {
			if err := c.checkExecutionGate(ctx); err != nil {
				return err
			}
			c.appendLightsailStaticIP(region, staticIP, result)
		}
		if page.NextPageToken == nil || aws.ToString(page.NextPageToken) == "" {
			break
		}
		pageToken = page.NextPageToken
	}
	return nil
}

func (c *AWSCollector) appendLightsailStaticIP(region string, staticIP lstypes.StaticIp, result *collector.AssetResult) {
	nativeID := aws.ToString(staticIP.Arn)
	if nativeID == "" {
		nativeID = aws.ToString(staticIP.Name)
	}
	if nativeID == "" {
		return
	}
	name := aws.ToString(staticIP.Name)
	if name == "" {
		name = nativeID
	}
	status := "available"
	if aws.ToBool(staticIP.IsAttached) {
		status = "attached"
	}
	asset := c.newBaseAsset(nativeID, name, string(model.AssetTypeNetwork), status, lightsailRegion(region, staticIP.Location))
	c.setService(&asset, "lightsail_static_ip")
	asset.Metadata = coremodel.JSONB{
		"attached_to":       staticIP.AttachedTo,
		"is_attached":       staticIP.IsAttached,
		"ip_address":        staticIP.IpAddress,
		"availability_zone": lightsailAvailabilityZone(staticIP.Location),
		"resource_type":     staticIP.ResourceType,
	}
	result.Assets = append(result.Assets, asset)
	if staticIP.IpAddress != nil {
		appendAWSIPAsset(result, asset.NativeID, aws.ToString(staticIP.IpAddress), "", true)
	}
}

func (c *AWSCollector) collectLightsailLoadBalancers(ctx context.Context, region string, svc lightsailAPI, result *collector.AssetResult) error {
	var pageToken *string
	for {
		page, err := svc.GetLoadBalancers(ctx, &lightsail.GetLoadBalancersInput{PageToken: pageToken})
		if err != nil {
			return err
		}
		for _, lb := range page.LoadBalancers {
			if err := c.checkExecutionGate(ctx); err != nil {
				return err
			}
			c.appendLightsailLoadBalancer(region, lb, result)
		}
		if page.NextPageToken == nil || aws.ToString(page.NextPageToken) == "" {
			break
		}
		pageToken = page.NextPageToken
	}
	return nil
}

func (c *AWSCollector) appendLightsailLoadBalancer(region string, lb lstypes.LoadBalancer, result *collector.AssetResult) {
	nativeID := aws.ToString(lb.Arn)
	if nativeID == "" {
		nativeID = aws.ToString(lb.Name)
	}
	if nativeID == "" {
		return
	}
	name := aws.ToString(lb.Name)
	if name == "" {
		name = nativeID
	}
	status := string(lb.State)
	if status == "" {
		status = "unknown"
	}
	asset := c.newBaseAsset(nativeID, name, string(model.AssetTypeNetwork), status, lightsailRegion(region, lb.Location))
	c.setService(&asset, "lightsail_lb")
	tags := convertLightsailTags(lb.Tags)
	asset.Environment = c.extractEnvironment(tags)
	asset.Criticality = c.extractCriticality(tags)
	asset.Metadata = coremodel.JSONB{
		"dns_name":                  lb.DnsName,
		"ip_address_type":           lb.IpAddressType,
		"protocol":                  lb.Protocol,
		"public_ports":              lb.PublicPorts,
		"instance_port":             lb.InstancePort,
		"health_check_path":         lb.HealthCheckPath,
		"https_redirection_enabled": lb.HttpsRedirectionEnabled,
		"availability_zone":         lightsailAvailabilityZone(lb.Location),
		"configuration_options":     lb.ConfigurationOptions,
		"tags":                      tags,
	}
	result.Assets = append(result.Assets, asset)
	dnsName := aws.ToString(lb.DnsName)
	if dnsName == "" {
		return
	}
	result.DomainAssets = append(result.DomainAssets, directPublicDomainAsset(asset.NativeID, dnsName, time.Now().UTC()))
	for _, publicPort := range lb.PublicPorts {
		if publicPort < 0 || publicPort > 65535 {
			continue
		}
		port := uint32(publicPort)
		record := directPublicEndpointPort(asset.NativeID, dnsName, port, "tcp", model.PortObservationSourceServiceEndpoint, "public", getServiceName(port), directPublicPortEvidence("lightsail_lb", asset.NativeID, "lightsail_load_balancer_public_ports", "account_resource"))
		record.Confidence = directPublicConfidence("low", "Lightsail load balancer public DNS endpoint")
		result.Ports = append(result.Ports, record)
	}
}

func lightsailPortInfoPorts(portInfo lstypes.InstancePortInfo) []uint32 {
	if portInfo.FromPort < 0 {
		return nil
	}
	to := portInfo.ToPort
	if to == 0 {
		to = portInfo.FromPort
	}
	if to < portInfo.FromPort {
		return nil
	}
	limitToRangeStart := to-portInfo.FromPort > 10
	seen := make(map[uint32]struct{})
	var ports []uint32
	for port := portInfo.FromPort; port <= to; port++ {
		if port > 65535 {
			break
		}
		portUint := uint32(port)
		if _, ok := seen[portUint]; !ok {
			seen[portUint] = struct{}{}
			ports = append(ports, portUint)
		}
		if limitToRangeStart {
			break
		}
	}
	return ports
}

func lightsailInstancePortMetadata(instance lstypes.Instance) []map[string]interface{} {
	if instance.Networking == nil {
		return nil
	}
	ports := make([]map[string]interface{}, 0, len(instance.Networking.Ports))
	for _, port := range instance.Networking.Ports {
		ports = append(ports, map[string]interface{}{
			"access_type":      port.AccessType,
			"access_direction": port.AccessDirection,
			"access_from":      port.AccessFrom,
			"from_port":        port.FromPort,
			"to_port":          port.ToPort,
			"protocol":         port.Protocol,
			"cidrs":            port.Cidrs,
			"ipv6_cidrs":       port.Ipv6Cidrs,
			"common_name":      port.CommonName,
		})
	}
	return ports
}

func lightsailRegion(fallback string, location *lstypes.ResourceLocation) string {
	if location == nil || location.RegionName == "" {
		return fallback
	}
	return string(location.RegionName)
}

func lightsailAvailabilityZone(location *lstypes.ResourceLocation) string {
	if location == nil || location.AvailabilityZone == nil {
		return ""
	}
	return aws.ToString(location.AvailabilityZone)
}

func convertLightsailTags(tags []lstypes.Tag) map[string]string {
	out := make(map[string]string)
	for _, tag := range tags {
		if tag.Key != nil && tag.Value != nil {
			out[aws.ToString(tag.Key)] = aws.ToString(tag.Value)
		}
	}
	return out
}