package remoteimage

import (
	"context"
	"errors"
	"fmt"
	"regexp"
	"strings"
	"time"

	awssdk "github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/ec2"
	ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types"
	"github.com/aws/aws-sdk-go-v2/service/ssm"
)

const (
	FamilyAmazonLinux2023 = "amazon-linux-2023"

	ParameterAmazonLinux2023X8664 = "/aws/service/ami-amazon-linux-latest/al2023-ami-kernel-default-x86_64"
	ParameterAmazonLinux2023Arm64 = "/aws/service/ami-amazon-linux-latest/al2023-ami-kernel-default-arm64"

	FailureClassInstanceType = "aws-remote-image-instance-type-invalid"
	FailureClassSSMParameter = "aws-remote-image-ssm-parameter-invalid"
	FailureClassAMI          = "aws-remote-image-ami-invalid"
)

var concreteAMI = regexp.MustCompile(`^ami-([0-9a-f]{8}|[0-9a-f]{17})$`)

// Diagnosis is the stable, machine-readable handoff for a failed AWS image
// metadata dependency.
type Diagnosis struct {
	Dependency   string
	FailureClass string
	Field        string
	Remediation  string
}

type dependencyError struct {
	diagnosis Diagnosis
	err       error
}

func (e *dependencyError) Error() string {
	return fmt.Sprintf(
		"remote image dependency %s failed: %v; remediation: %s",
		e.diagnosis.Dependency,
		e.err,
		e.diagnosis.Remediation,
	)
}

func (e *dependencyError) Unwrap() error { return e.err }

// Diagnose returns actionable metadata when err came from an AWS image
// metadata dependency.
func Diagnose(err error) (Diagnosis, bool) {
	var target *dependencyError
	if !errors.As(err, &target) {
		return Diagnosis{}, false
	}
	return target.diagnosis, true
}

// EC2API is the read-only EC2 metadata surface used to derive a compatible AMI.
type EC2API interface {
	DescribeInstanceTypes(ctx context.Context, params *ec2.DescribeInstanceTypesInput,
		optFns ...func(*ec2.Options)) (*ec2.DescribeInstanceTypesOutput, error)
	DescribeImages(ctx context.Context, params *ec2.DescribeImagesInput,
		optFns ...func(*ec2.Options)) (*ec2.DescribeImagesOutput, error)
}

// SSMAPI is the read-only Parameter Store surface used to resolve AWS's
// regional Amazon Linux 2023 public parameter.
type SSMAPI interface {
	GetParameter(ctx context.Context, params *ssm.GetParameterInput,
		optFns ...func(*ssm.Options)) (*ssm.GetParameterOutput, error)
}

// Plan is the immutable remote-host image decision persisted before cloud mutation.
type Plan struct {
	Family         string    `json:"family"           yaml:"family"`
	InstanceType   string    `json:"instance_type"    yaml:"instance_type"`
	Architecture   string    `json:"architecture"     yaml:"architecture"`
	Parameter      string    `json:"parameter"        yaml:"parameter"`
	AMI            string    `json:"ami_id"           yaml:"ami_id"`
	Region         string    `json:"region"           yaml:"region"`
	RootDeviceName string    `json:"root_device_name" yaml:"root_device_name"`
	ResolvedAt     time.Time `json:"resolved_at"      yaml:"resolved_at"`
}

// Resolver derives and validates a concrete regional remote-host image.
type Resolver struct {
	EC2    EC2API
	SSM    SSMAPI
	Region string
	Now    func() time.Time
}

// Resolve derives the instance architecture, reads AWS's regional AL2023
// public parameter, and validates the returned concrete AMI before returning a
// plan suitable for persistence.
func (r *Resolver) Resolve(ctx context.Context, instanceType string) (*Plan, error) {
	if r == nil {
		return nil, fmt.Errorf("remote image resolver is required")
	}
	instanceType = strings.TrimSpace(instanceType)
	if instanceType == "" {
		return nil, fmt.Errorf("host.ec2.instance_type is required to derive the remote image")
	}
	region := strings.TrimSpace(r.Region)
	if region == "" {
		return nil, fmt.Errorf("target.region is required to derive the remote image")
	}
	if r.EC2 == nil {
		return nil, fmt.Errorf("EC2 metadata client is required to derive the remote image")
	}
	if r.SSM == nil {
		return nil, fmt.Errorf("SSM client is required to derive the remote image")
	}

	architecture, err := r.resolveArchitecture(ctx, instanceType)
	if err != nil {
		return nil, instanceTypeDependencyError(instanceType, region, err)
	}
	parameter := parameterForArchitecture(architecture)
	ami, err := r.resolveParameter(ctx, parameter)
	if err != nil {
		return nil, ssmDependencyError(parameter, region, err)
	}
	rootDeviceName, err := r.validateImage(ctx, ami, architecture)
	if err != nil {
		return nil, imageDependencyError(ami, architecture, region, err)
	}
	resolvedAt := time.Now().UTC()
	if r.Now != nil {
		resolvedAt = r.Now().UTC()
	}
	return &Plan{
		Family:         FamilyAmazonLinux2023,
		InstanceType:   instanceType,
		Architecture:   architecture,
		Parameter:      parameter,
		AMI:            ami,
		Region:         region,
		RootDeviceName: rootDeviceName,
		ResolvedAt:     resolvedAt,
	}, nil
}

// ValidatePinned verifies a persisted plan against the current instance-type
// and EC2 image metadata without consulting SSM. This preserves the original
// AMI decision across retries even when AWS advances the public parameter.
func (r *Resolver) ValidatePinned(ctx context.Context, plan *Plan, instanceType string) error {
	if r == nil {
		return fmt.Errorf("remote image resolver is required")
	}
	if plan == nil {
		return fmt.Errorf("persisted remote image plan is required")
	}
	instanceType = strings.TrimSpace(instanceType)
	if instanceType == "" {
		return fmt.Errorf("host.ec2.instance_type is required to validate the pinned remote image")
	}
	region := strings.TrimSpace(r.Region)
	if region == "" {
		return fmt.Errorf("target.region is required to validate the pinned remote image")
	}
	if strings.TrimSpace(plan.Family) != FamilyAmazonLinux2023 {
		return fmt.Errorf("pinned remote image family %q is unsupported; expected %q", plan.Family, FamilyAmazonLinux2023)
	}
	if pinnedType := strings.TrimSpace(plan.InstanceType); pinnedType != instanceType {
		return fmt.Errorf("pinned remote image instance type %q does not match configured instance type %q", pinnedType, instanceType)
	}
	if pinnedRegion := strings.TrimSpace(plan.Region); pinnedRegion != region {
		return fmt.Errorf("pinned remote image region %q does not match target region %q", pinnedRegion, region)
	}
	architecture := strings.TrimSpace(plan.Architecture)
	if architecture != string(ec2types.ArchitectureTypeX8664) && architecture != string(ec2types.ArchitectureTypeArm64) {
		return fmt.Errorf("pinned remote image architecture %q is unsupported; expected x86_64 or arm64", architecture)
	}
	expectedParameter := parameterForArchitecture(architecture)
	if parameter := strings.TrimSpace(plan.Parameter); parameter != expectedParameter {
		return fmt.Errorf("pinned remote image parameter %q does not match %s architecture parameter %q", parameter, architecture, expectedParameter)
	}
	ami := strings.TrimSpace(plan.AMI)
	if !concreteAMI.MatchString(ami) {
		return fmt.Errorf("pinned remote image AMI %q is not a concrete AMI id", ami)
	}
	rootDeviceName := strings.TrimSpace(plan.RootDeviceName)
	if rootDeviceName == "" {
		return fmt.Errorf("pinned remote image %q root device name is empty", ami)
	}
	if plan.ResolvedAt.IsZero() {
		return fmt.Errorf("pinned remote image %q resolved_at is required", ami)
	}
	if r.EC2 == nil {
		return fmt.Errorf("EC2 metadata client is required to validate the pinned remote image")
	}

	currentArchitecture, err := r.resolveArchitecture(ctx, instanceType)
	if err != nil {
		return instanceTypeDependencyError(instanceType, region, err)
	}
	if currentArchitecture != architecture {
		return fmt.Errorf(
			"pinned remote image architecture %q does not match instance type %q architecture %q",
			architecture,
			instanceType,
			currentArchitecture,
		)
	}
	currentRootDeviceName, err := r.validateImage(ctx, ami, architecture)
	if err != nil {
		return imageDependencyError(ami, architecture, region, err)
	}
	if currentRootDeviceName != rootDeviceName {
		return fmt.Errorf(
			"pinned remote image %q root device %q does not match current EC2 metadata %q",
			ami,
			rootDeviceName,
			currentRootDeviceName,
		)
	}
	return nil
}

func (r *Resolver) resolveArchitecture(ctx context.Context, instanceType string) (string, error) {
	out, err := r.EC2.DescribeInstanceTypes(ctx, &ec2.DescribeInstanceTypesInput{
		InstanceTypes: []ec2types.InstanceType{ec2types.InstanceType(instanceType)},
	})
	if err != nil {
		return "", fmt.Errorf("DescribeInstanceTypes %q for remote image architecture: %w", instanceType, err)
	}
	if out == nil || len(out.InstanceTypes) != 1 {
		count := 0
		if out != nil {
			count = len(out.InstanceTypes)
		}
		return "", fmt.Errorf("DescribeInstanceTypes %q returned %d instance types; expected exactly one", instanceType, count)
	}
	returnedType := strings.TrimSpace(string(out.InstanceTypes[0].InstanceType))
	if returnedType != instanceType {
		return "", fmt.Errorf("DescribeInstanceTypes %q returned instance type %q", instanceType, returnedType)
	}
	processor := out.InstanceTypes[0].ProcessorInfo
	if processor == nil || len(processor.SupportedArchitectures) != 1 {
		var architectures []ec2types.ArchitectureType
		if processor != nil {
			architectures = processor.SupportedArchitectures
		}
		return "", fmt.Errorf(
			"instance type %q must report exactly one supported architecture (x86_64 or arm64); got %v",
			instanceType,
			architectures,
		)
	}
	architecture := string(processor.SupportedArchitectures[0])
	switch architecture {
	case string(ec2types.ArchitectureTypeX8664), string(ec2types.ArchitectureTypeArm64):
		return architecture, nil
	default:
		return "", fmt.Errorf(
			"instance type %q reports unsupported architecture %q; cp9 requires exactly one of x86_64 or arm64",
			instanceType,
			architecture,
		)
	}
}

func (r *Resolver) resolveParameter(ctx context.Context, parameter string) (string, error) {
	out, err := r.SSM.GetParameter(ctx, &ssm.GetParameterInput{Name: awssdk.String(parameter)})
	if err != nil {
		return "", fmt.Errorf("SSM GetParameter %q for regional Amazon Linux 2023 AMI: %w", parameter, err)
	}
	ami := ""
	if out != nil && out.Parameter != nil {
		ami = strings.TrimSpace(awssdk.ToString(out.Parameter.Value))
	}
	if !concreteAMI.MatchString(ami) {
		return "", fmt.Errorf("SSM GetParameter %q returned invalid concrete AMI %q", parameter, ami)
	}
	return ami, nil
}

func (r *Resolver) validateImage(ctx context.Context, ami, architecture string) (string, error) {
	out, err := r.EC2.DescribeImages(ctx, &ec2.DescribeImagesInput{ImageIds: []string{ami}})
	if err != nil {
		return "", fmt.Errorf("DescribeImages validate remote AMI %q: %w", ami, err)
	}
	if out == nil || len(out.Images) != 1 {
		count := 0
		if out != nil {
			count = len(out.Images)
		}
		return "", fmt.Errorf("DescribeImages remote AMI %q returned %d images; expected exactly one visible image", ami, count)
	}
	image := out.Images[0]
	if imageID := strings.TrimSpace(awssdk.ToString(image.ImageId)); imageID != ami {
		return "", fmt.Errorf("DescribeImages remote AMI %q returned image id %q", ami, imageID)
	}
	if image.State != ec2types.ImageStateAvailable {
		return "", fmt.Errorf("remote AMI %q must be available; got %q", ami, image.State)
	}
	if imageArchitecture := string(image.Architecture); imageArchitecture != architecture {
		return "", fmt.Errorf(
			"remote AMI %q architecture %q is incompatible with instance architecture %q",
			ami,
			imageArchitecture,
			architecture,
		)
	}
	rootDeviceName := strings.TrimSpace(awssdk.ToString(image.RootDeviceName))
	if rootDeviceName == "" {
		return "", fmt.Errorf("remote AMI %q root device name is empty", ami)
	}
	return rootDeviceName, nil
}

func parameterForArchitecture(architecture string) string {
	if architecture == string(ec2types.ArchitectureTypeArm64) {
		return ParameterAmazonLinux2023Arm64
	}
	return ParameterAmazonLinux2023X8664
}

func instanceTypeDependencyError(instanceType, region string, err error) error {
	return &dependencyError{
		diagnosis: Diagnosis{
			Dependency:   "EC2 DescribeInstanceTypes",
			FailureClass: FailureClassInstanceType,
			Field:        "host.ec2.instance_type",
			Remediation: fmt.Sprintf(
				"verify host.ec2.instance_type %q is available in target.region %q, reports exactly one x86_64 or arm64 architecture, and the role allows ec2:DescribeInstanceTypes",
				instanceType,
				region,
			),
		},
		err: err,
	}
}

func ssmDependencyError(parameter, region string, err error) error {
	return &dependencyError{
		diagnosis: Diagnosis{
			Dependency:   "SSM GetParameter",
			FailureClass: FailureClassSSMParameter,
			Field:        "target.region",
			Remediation: fmt.Sprintf(
				"verify public parameter %q is available in target.region %q and the role allows ssm:GetParameter; do not set host.ec2.ami",
				parameter,
				region,
			),
		},
		err: err,
	}
}

func imageDependencyError(ami, architecture, region string, err error) error {
	return &dependencyError{
		diagnosis: Diagnosis{
			Dependency:   "EC2 DescribeImages",
			FailureClass: FailureClassAMI,
			Field:        "target.region",
			Remediation: fmt.Sprintf(
				"verify resolved AMI %q is visible and available in target.region %q, matches architecture %q, and the role allows ec2:DescribeImages; do not set host.ec2.ami",
				ami,
				region,
				architecture,
			),
		},
		err: err,
	}
}