// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file is distributed
// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
// express or implied. See the License for the specific language governing
// permissions and limitations under the License.

package tasknetworkconfig

import (
	ni "github.com/aws/amazon-ecs-agent/ecs-agent/netlib/model/networkinterface"

	"github.com/aws/aws-sdk-go-v2/service/ecs/types"
	"github.com/pkg/errors"
)

// TaskNetworkConfig is the top level network data structure associated with a task.
type TaskNetworkConfig struct {
	NetworkNamespaces []*NetworkNamespace
	NetworkMode       types.NetworkMode
}

func New(networkMode types.NetworkMode, netNSs ...*NetworkNamespace) (*TaskNetworkConfig, error) {
	if networkMode != types.NetworkModeAwsvpc &&
		networkMode != types.NetworkModeBridge &&
		networkMode != types.NetworkModeHost &&
		networkMode != types.NetworkModeNone {
		return nil, errors.New("invalid network mode: " + string(networkMode))
	}

	return &TaskNetworkConfig{
		NetworkNamespaces: netNSs,
		NetworkMode:       networkMode,
	}, nil
}

// GetPrimaryInterface returns the interface with index 0 inside the network namespace
// with index 0 associated with the task's network config.
func (tnc *TaskNetworkConfig) GetPrimaryInterface() *ni.NetworkInterface {
	if tnc != nil && tnc.GetPrimaryNetNS() != nil {
		return tnc.GetPrimaryNetNS().GetPrimaryInterface()
	}
	return nil
}

// GetPrimaryNetNS returns the netns with index 0 associated with the task's network config.
func (tnc *TaskNetworkConfig) GetPrimaryNetNS() *NetworkNamespace {
	for _, netns := range tnc.NetworkNamespaces {
		if netns.Index == 0 {
			return netns
		}
	}

	return nil
}

// GetEniNamesToAssociationProtocolMapping returns a map of ENI names to
// interface association protocols (like tunnel/veth).
func (tnc *TaskNetworkConfig) GetEniNamesToAssociationProtocolMapping() map[string]string {
	eniNameToAssociationProtocol := make(map[string]string)
	for _, netNS := range tnc.NetworkNamespaces {
		for _, iface := range netNS.NetworkInterfaces {
			if iface.Name != "" {
				eniNameToAssociationProtocol[iface.Name] = iface.InterfaceAssociationProtocol
			}
		}
	}
	return eniNameToAssociationProtocol
}

// GetInterfaceNamesToNetNSMapping returns a map where key is interface name and value is the netns
// in which the interface exists.
func (tnc *TaskNetworkConfig) GetInterfaceNamesToNetNSMapping() map[string]*NetworkNamespace {
	name2NetNS := make(map[string]*NetworkNamespace)
	for _, netNS := range tnc.NetworkNamespaces {
		for _, iface := range netNS.NetworkInterfaces {
			if iface.Name != "" {
				name2NetNS[iface.Name] = netNS
			}
		}
	}
	return name2NetNS
}
