// 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 status

import (
	"errors"
	"strconv"
	"strings"
)

const (
	// ContainerStatusNone is the zero state of a container; this container has not completed pull
	ContainerStatusNone ContainerStatus = iota
	// ContainerManifestPulled represents a container which has had its image manifest pulled
	ContainerManifestPulled
	// ContainerPulled represents a container which has had the image pulled
	ContainerPulled
	// ContainerCreated represents a container that has been created
	ContainerCreated
	// ContainerRunning represents a container that has started
	ContainerRunning
	// ContainerResourcesProvisioned represents a container that has completed provisioning all of its
	// resources. Non-internal containers (containers present in the task definition) transition to
	// this state without doing any additional work. However, containers that are added to a task
	// by the ECS Agent would possibly need to perform additional actions before they can be
	// considered "ready" and contribute to the progress of a task. For example, the "pause" container
	// would be provisioned by invoking CNI plugins
	ContainerResourcesProvisioned
	// ContainerStopped represents a container that has stopped
	ContainerStopped
	// ContainerZombie is an "impossible" state that is used as the maximum
	ContainerZombie
)

const (
	// ContainerHealthUnknown is the initial status of container health
	ContainerHealthUnknown ContainerHealthStatus = iota
	// ContainerHealthy represents the status of container health check when returned healthy
	ContainerHealthy
	// ContainerUnhealthy represents the status of container health check when returned unhealthy
	ContainerUnhealthy
)

// ContainerStatus is an enumeration of valid states in the container lifecycle
type ContainerStatus int32

// ContainerHealthStatus is an enumeration of container health check status
type ContainerHealthStatus int32

var containerStatusMap = map[string]ContainerStatus{
	"NONE":                  ContainerStatusNone,
	"MANIFEST_PULLED":       ContainerManifestPulled,
	"PULLED":                ContainerPulled,
	"CREATED":               ContainerCreated,
	"RUNNING":               ContainerRunning,
	"RESOURCES_PROVISIONED": ContainerResourcesProvisioned,
	"STOPPED":               ContainerStopped,
}

// BackendStatus returns the container health status recognized by backend
func (healthStatus ContainerHealthStatus) BackendStatus() string {
	switch healthStatus {
	case ContainerHealthy:
		return "HEALTHY"
	case ContainerUnhealthy:
		return "UNHEALTHY"
	default:
		return "UNKNOWN"
	}
}

// String returns the readable description of the container health status
func (healthStatus ContainerHealthStatus) String() string {
	return healthStatus.BackendStatus()
}

// String returns a human readable string representation of this object
func (cs ContainerStatus) String() string {
	for k, v := range containerStatusMap {
		if v == cs {
			return k
		}
	}
	return "NONE"
}

// ShouldReportToBackend returns true if the container status is recognized as a
// valid state by ECS. Note that not all container statuses are recognized by ECS
// or map to ECS states
func (cs *ContainerStatus) ShouldReportToBackend(steadyStateStatus ContainerStatus) bool {
	return *cs == steadyStateStatus || *cs == ContainerStopped
}

// BackendStatusString maps the internal container status in Agent to a backend recognized
// status string.
func (cs ContainerStatus) BackendStatusString() string {
	switch cs {
	case ContainerRunning:
		return "RUNNING"
	case ContainerStopped:
		return "STOPPED"
	default:
		return "PENDING"
	}
}

// Terminal returns true if the container status is STOPPED
func (cs ContainerStatus) Terminal() bool {
	return cs == ContainerStopped
}

// UnmarshalJSON overrides the logic for parsing the JSON-encoded ContainerStatus data
func (cs *ContainerStatus) UnmarshalJSON(b []byte) error {
	if strings.ToLower(string(b)) == "null" {
		*cs = ContainerStatusNone
		return nil
	}
	if b[0] != '"' || b[len(b)-1] != '"' {
		*cs = ContainerStatusNone
		return errors.New("container status unmarshal: status must be a string or null; Got " + string(b))
	}
	strStatus := string(b[1 : len(b)-1])
	// 'UNKNOWN' and 'DEAD' for Compatibility with v1.0.0 state files
	if strStatus == "UNKNOWN" {
		*cs = ContainerStatusNone
		return nil
	}
	if strStatus == "DEAD" {
		*cs = ContainerStopped
		return nil
	}

	// Before MarshalText method was implemented for ContainerStatus, ContainerStatus was
	// marshaled as an integer in cases that depend on MarshalText such as keys in a JSON object.
	// To make unmarshaling backwards-compatible, check if the marshaled text is an integer
	// and map it to container status.
	if intStatus, err := strconv.Atoi(strStatus); err == nil {
		// This map is only for making text unmarshaling compatible with old state files.
		// Updates are NOT needed to this map as newer container states are introduced.
		intContainerStatusMap := map[int]ContainerStatus{
			0: ContainerStatusNone,
			1: ContainerPulled,
			2: ContainerCreated,
			3: ContainerRunning,
			4: ContainerResourcesProvisioned,
			5: ContainerStopped,
			6: ContainerZombie,
		}
		if stat, ok := intContainerStatusMap[intStatus]; ok {
			*cs = stat
			return nil
		}
	}

	stat, ok := containerStatusMap[strStatus]
	if !ok {
		*cs = ContainerStatusNone
		return errors.New("container status unmarshal: unrecognized status")
	}
	*cs = stat
	return nil
}

// MarshalJSON overrides the logic for JSON-encoding the ContainerStatus type
func (cs *ContainerStatus) MarshalJSON() ([]byte, error) {
	if cs == nil {
		return nil, nil
	}
	return []byte(`"` + cs.String() + `"`), nil
}

// Marshals a container status to its text form.
// In some cases such as a map with ContainerStatus as keys, MarshalText method is used to
// marshal container statuses by functions in the encoding package. Without this method
// container statuses will be marshaled as integers which is undesirable.
func (cs ContainerStatus) MarshalText() ([]byte, error) {
	return []byte(cs.String()), nil
}

// Unmarshals a container status from its text form.
func (cs *ContainerStatus) UnmarshalText(b []byte) error {
	strStatus := string(b)
	stat, ok := containerStatusMap[strStatus]
	if !ok {
		*cs = ContainerStatusNone
		return errors.New("container status text unmarshal: unrecognized status: " + strStatus)
	}
	*cs = stat
	return nil
}

// UnmarshalJSON overrides the logic for parsing the JSON-encoded container health data
func (healthStatus *ContainerHealthStatus) UnmarshalJSON(b []byte) error {
	*healthStatus = ContainerHealthUnknown

	if strings.ToLower(string(b)) == "null" {
		return nil
	}
	if b[0] != '"' || b[len(b)-1] != '"' {
		return errors.New("container health status unmarshal: status must be a string or null; Got " + string(b))
	}

	strStatus := string(b[1 : len(b)-1])
	switch strStatus {
	case "UNKNOWN":
	// The health status is already set to ContainerHealthUnknown initially
	case "HEALTHY":
		*healthStatus = ContainerHealthy
	case "UNHEALTHY":
		*healthStatus = ContainerUnhealthy
	default:
		return errors.New("container health status unmarshal: unrecognized status: " + string(b))
	}
	return nil
}

// MarshalJSON overrides the logic for JSON-encoding the ContainerHealthStatus type
func (healthStatus *ContainerHealthStatus) MarshalJSON() ([]byte, error) {
	if healthStatus == nil {
		return nil, nil
	}
	return []byte(`"` + healthStatus.String() + `"`), nil
}

// IsRunning returns true if the container status is either RUNNING or RESOURCES_PROVISIONED
func (cs ContainerStatus) IsRunning() bool {
	return cs == ContainerRunning || cs == ContainerResourcesProvisioned
}
