package aws

import (
	"context"
	"fmt"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/ecs"

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

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

	return c.forEachECSCluster(ctx, svc, func(clusterArn string) error {
		if err := c.collectECSCluster(ctx, svc, clusterArn, region, result); err != nil {
			appendAWSPartialFailure(
				result,
				"service",
				fmt.Sprintf("%s:ecs:%s", region, clusterArn),
				fmt.Sprintf("failed to collect ECS cluster %s: %v", clusterArn, err),
			)
			c.report("warning", fmt.Sprintf("Failed to collect ECS cluster: %v", err), map[string]interface{}{
				"cluster_arn": clusterArn,
				"region":      region,
				"error":       err.Error(),
			})
		}
		return nil
	})
}

func (c *AWSCollector) forEachECSCluster(ctx context.Context, svc *ecs.Client, handle func(string) error) error {
	var nextToken *string

	for {
		resp, err := svc.ListClusters(ctx, &ecs.ListClustersInput{
			MaxResults: aws.Int32(100),
			NextToken:  nextToken,
		})
		if err != nil {
			return err
		}
		for _, clusterArn := range resp.ClusterArns {
			if err := handle(clusterArn); err != nil {
				return err
			}
		}
		if resp.NextToken == nil || *resp.NextToken == "" {
			break
		}
		nextToken = resp.NextToken
	}

	return nil
}

func (c *AWSCollector) collectECSCluster(ctx context.Context, svc *ecs.Client, clusterArn, region string, result *collector.AssetResult) error {
	// Get cluster details
	descResp, err := svc.DescribeClusters(ctx, &ecs.DescribeClustersInput{
		Clusters: []string{clusterArn},
	})
	if err != nil {
		return err
	}

	if len(descResp.Clusters) == 0 {
		return nil
	}

	cluster := descResp.Clusters[0]
	clusterName := ""
	if cluster.ClusterName != nil {
		clusterName = *cluster.ClusterName
	}

	// Create cluster asset
	asset := c.newBaseAsset(clusterArn, clusterName, string(model.AssetTypeCompute), "active", region)
	c.setService(&asset, "ecs")
	meta := map[string]interface{}{
		"cluster_arn":                          clusterArn,
		"status":                               cluster.Status,
		"running_tasks_count":                  cluster.RunningTasksCount,
		"pending_tasks_count":                  cluster.PendingTasksCount,
		"active_services_count":                cluster.ActiveServicesCount,
		"registered_container_instances_count": cluster.RegisteredContainerInstancesCount,
	}
	asset.Metadata = coremodel.JSONB(meta)
	result.Assets = append(result.Assets, asset)

	// Collect services in this cluster
	if err := c.collectECSServices(ctx, svc, clusterArn, clusterName, region, result); err != nil {
		appendAWSPartialFailure(
			result,
			"service",
			fmt.Sprintf("%s:ecs:%s:services", region, clusterArn),
			fmt.Sprintf("failed to collect ECS services for cluster %s: %v", clusterArn, err),
		)
		c.report("warning", fmt.Sprintf("Failed to collect ECS services for cluster %s: %v", clusterArn, err), map[string]interface{}{
			"cluster_arn": clusterArn,
			"region":      region,
			"error":       err.Error(),
		})
	}

	return nil
}

func (c *AWSCollector) collectECSServices(ctx context.Context, svc *ecs.Client, clusterArn, clusterName, region string, result *collector.AssetResult) error {
	return c.forEachECSServiceBatch(ctx, svc, clusterArn, 10, func(batch []string) error {
		descResp, err := svc.DescribeServices(ctx, &ecs.DescribeServicesInput{
			Cluster:  &clusterArn,
			Services: batch,
		})
		if err != nil {
			return err
		}

		for _, service := range descResp.Services {
			if service.ServiceArn == nil {
				continue
			}

			serviceArn := *service.ServiceArn
			serviceName := ""
			if service.ServiceName != nil {
				serviceName = *service.ServiceName
			}

			// Create service asset
			asset := c.newBaseAsset(serviceArn, serviceName, "ecs_service", "active", region)
			c.setService(&asset, "ecs")

			// Get task definition to extract port mappings
			var portMappings []map[string]interface{}
			if service.TaskDefinition != nil {
				var err error
				portMappings, err = c.getECSPortMappings(ctx, svc, *service.TaskDefinition)
				if err != nil {
					appendAWSPartialFailure(
						result,
						"service",
						fmt.Sprintf("%s:ecs:%s:task_definition", region, serviceArn),
						fmt.Sprintf("failed to collect ECS task definition for service %s: %v", serviceArn, err),
					)
					c.report("warning", fmt.Sprintf("Failed to collect ECS task definition for service %s: %v", serviceArn, err), map[string]interface{}{
						"cluster_arn":     clusterArn,
						"service_arn":     serviceArn,
						"task_definition": *service.TaskDefinition,
						"region":          region,
						"error":           err.Error(),
					})
				}
			}

			// Get load balancers
			lbConfigs := make([]map[string]interface{}, 0)
			for _, lb := range service.LoadBalancers {
				lbConfig := map[string]interface{}{}
				if lb.TargetGroupArn != nil {
					lbConfig["target_group_arn"] = *lb.TargetGroupArn
				}
				if lb.ContainerName != nil {
					lbConfig["container_name"] = *lb.ContainerName
				}
				if lb.ContainerPort != nil {
					lbConfig["container_port"] = *lb.ContainerPort
				}
				lbConfigs = append(lbConfigs, lbConfig)
			}

			meta := map[string]interface{}{
				"service_arn":     serviceArn,
				"cluster_name":    clusterName,
				"task_definition": service.TaskDefinition,
				"desired_count":   service.DesiredCount,
				"running_count":   service.RunningCount,
				"pending_count":   service.PendingCount,
				"launch_type":     service.LaunchType,
				"port_mappings":   portMappings,
				"load_balancers":  lbConfigs,
			}
			asset.Metadata = coremodel.JSONB(meta)
			result.Assets = append(result.Assets, asset)

			// Collect running tasks for this service to get actual IPs and ports
			if err := c.collectECSTasks(ctx, svc, clusterArn, serviceArn, region, portMappings, result); err != nil {
				appendAWSPartialFailure(
					result,
					"service",
					fmt.Sprintf("%s:ecs:%s:tasks", region, serviceArn),
					fmt.Sprintf("failed to collect ECS tasks for service %s: %v", serviceArn, err),
				)
				c.report("warning", fmt.Sprintf("Failed to collect ECS tasks for service %s: %v", serviceArn, err), map[string]interface{}{
					"cluster_arn": clusterArn,
					"service_arn": serviceArn,
					"region":      region,
					"error":       err.Error(),
				})
			}
		}
		return nil
	})
}

func (c *AWSCollector) forEachECSServiceBatch(ctx context.Context, svc *ecs.Client, clusterArn string, batchSize int, handle func([]string) error) error {
	if batchSize <= 0 {
		batchSize = 10
	}
	batch := make([]string, 0, batchSize)
	flush := func() error {
		if len(batch) == 0 {
			return nil
		}
		current := append([]string(nil), batch...)
		batch = batch[:0]
		return handle(current)
	}
	var nextToken *string

	for {
		resp, err := svc.ListServices(ctx, &ecs.ListServicesInput{
			MaxResults: aws.Int32(100),
			Cluster:    &clusterArn,
			NextToken:  nextToken,
		})
		if err != nil {
			return err
		}
		for _, serviceArn := range resp.ServiceArns {
			batch = append(batch, serviceArn)
			if len(batch) >= batchSize {
				if err := flush(); err != nil {
					return err
				}
			}
		}
		if resp.NextToken == nil || *resp.NextToken == "" {
			break
		}
		nextToken = resp.NextToken
	}

	return flush()
}

func (c *AWSCollector) getECSPortMappings(ctx context.Context, svc *ecs.Client, taskDefArn string) ([]map[string]interface{}, error) {
	resp, err := svc.DescribeTaskDefinition(ctx, &ecs.DescribeTaskDefinitionInput{
		TaskDefinition: &taskDefArn,
	})
	if err != nil {
		return nil, err
	}

	if resp.TaskDefinition == nil {
		return nil, nil
	}

	portMappings := make([]map[string]interface{}, 0)
	for _, containerDef := range resp.TaskDefinition.ContainerDefinitions {
		for _, pm := range containerDef.PortMappings {
			mapping := map[string]interface{}{}
			if containerDef.Name != nil {
				mapping["container_name"] = *containerDef.Name
			}
			if pm.ContainerPort != nil {
				mapping["container_port"] = *pm.ContainerPort
			}
			if pm.HostPort != nil {
				mapping["host_port"] = *pm.HostPort
			}
			if pm.Protocol != "" {
				mapping["protocol"] = string(pm.Protocol)
			}
			portMappings = append(portMappings, mapping)
		}
	}

	return portMappings, nil
}

func (c *AWSCollector) collectECSTasks(ctx context.Context, svc *ecs.Client, clusterArn, serviceArn, region string, portMappings []map[string]interface{}, result *collector.AssetResult) error {
	return c.forEachECSTaskBatch(ctx, svc, clusterArn, serviceArn, 100, func(batch []string) error {
		descResp, err := svc.DescribeTasks(ctx, &ecs.DescribeTasksInput{
			Cluster: &clusterArn,
			Tasks:   batch,
		})
		if err != nil {
			return err
		}

		for _, task := range descResp.Tasks {
			if task.TaskArn == nil {
				continue
			}

			taskArn := *task.TaskArn

			// Extract task IP based on network mode
			var taskIP string
			networkMode := "bridge" // default

			// For awsvpc mode, extract ENI IP
			if len(task.Attachments) > 0 {
				for _, attachment := range task.Attachments {
					if attachment.Type != nil && *attachment.Type == "ElasticNetworkInterface" {
						networkMode = "awsvpc"
						// Extract private IP from attachment details
						for _, detail := range attachment.Details {
							if detail.Name != nil && *detail.Name == "privateIPv4Address" && detail.Value != nil {
								taskIP = *detail.Value
								break
							}
						}
					}
				}
			}

			// Create task asset
			asset := c.newBaseAsset(taskArn, taskArn, string(model.AssetTypeCompute), "running", region)
			c.setService(&asset, "ecs")
			meta := map[string]interface{}{
				"task_arn":            taskArn,
				"task_definition_arn": task.TaskDefinitionArn,
				"desired_status":      task.DesiredStatus,
				"last_status":         task.LastStatus,
				"network_mode":        networkMode,
				"task_ip":             taskIP,
			}
			asset.Metadata = coremodel.JSONB(meta)
			result.Assets = append(result.Assets, asset)

			// Add IP as IP asset if available
			if taskIP != "" {
				result.IPAssets = append(result.IPAssets, model.IPAsset{
					Asset:     model.Asset{NativeID: asset.NativeID},
					IPAddress: taskIP,
					IsPublic:  false, // Task IPs are in VPC, private by default
				})

				// Add port entries for each port mapping
				for _, pm := range portMappings {
					containerPortVal, hasContainerPort := pm["container_port"]
					protocolVal, hasProtocol := pm["protocol"]

					if !hasContainerPort {
						continue
					}

					var port uint32
					switch v := containerPortVal.(type) {
					case int32:
						port = uint32(v)
					case int:
						port = uint32(v)
					case uint32:
						port = v
					default:
						continue
					}

					protocol := "tcp"
					if hasProtocol {
						if p, ok := protocolVal.(string); ok {
							protocol = p
						}
					}

					// For ECS tasks, exposure depends on security groups
					// We'll mark as private by default (should cross-reference with security groups)
					result.Ports = append(result.Ports, model.IPPortIndex{
						Asset:             model.Asset{NativeID: asset.NativeID},
						IPAddress:         taskIP,
						Port:              port,
						Protocol:          protocol,
						ObservationSource: model.PortObservationSourceTaskMapping,
						ServiceName:       getServiceName(port),
						Exposure:          "private", // TODO: cross-reference with security groups
						Status:            "open",
					})
				}
			}
		}
		return nil
	})
}

func (c *AWSCollector) forEachECSTaskBatch(ctx context.Context, svc *ecs.Client, clusterArn, serviceArn string, batchSize int, handle func([]string) error) error {
	if batchSize <= 0 {
		batchSize = 100
	}
	batch := make([]string, 0, batchSize)
	flush := func() error {
		if len(batch) == 0 {
			return nil
		}
		current := append([]string(nil), batch...)
		batch = batch[:0]
		return handle(current)
	}
	var nextToken *string

	for {
		resp, err := svc.ListTasks(ctx, &ecs.ListTasksInput{
			MaxResults:  aws.Int32(100),
			Cluster:     &clusterArn,
			ServiceName: &serviceArn,
			NextToken:   nextToken,
		})
		if err != nil {
			return err
		}
		for _, taskArn := range resp.TaskArns {
			batch = append(batch, taskArn)
			if len(batch) >= batchSize {
				if err := flush(); err != nil {
					return err
				}
			}
		}
		if resp.NextToken == nil || *resp.NextToken == "" {
			break
		}
		nextToken = resp.NextToken
	}

	return flush()
}