package service

import (
	"context"
	"os/exec"

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

	"github.com/tidbcloud/infra-provider/pkg/cloud/aws"
	"github.com/tidbcloud/infra-provider/pkg/cloud/azure"
	"github.com/tidbcloud/infra-provider/pkg/cloud/gcp"
)

// NodeRestartProvider is an interface to restart a pod's node
type NodeRestartProvider interface {
	Restart(context.Context, string) error
}

type awsNodeRestarter struct {
	awsCli aws.Client
}

func NewAWSNodeRestarter(awsCli aws.Client) NodeRestartProvider {
	return &awsNodeRestarter{awsCli: awsCli}
}

func (r *awsNodeRestarter) Restart(ctx context.Context, providerID string) error {
	region, instanceID, err := parseEKSNodeProviderID(providerID)
	if err != nil {
		return err
	}
	_, err = r.awsCli.EC2(nil, aws.WithRegion(region)).
		RebootInstances(ctx, &ec2.RebootInstancesInput{InstanceIds: []string{instanceID}})
	return err
}

type gcpNodeRestarter struct {
	gcpCli gcp.Client
}

func NewGCPNodeRestarter(gcpCli gcp.Client) NodeRestartProvider {
	return &gcpNodeRestarter{gcpCli: gcpCli}
}

func (r *gcpNodeRestarter) Restart(ctx context.Context, providerID string) error {
	project, zone, instanceName, err := parseGKENodeProviderID(providerID)
	if err != nil {
		return err
	}
	return r.gcpCli.RestartInstance(ctx, project, zone, instanceName)
}

type azureNodeRestarter struct {
	azCli azure.Client
}

func NewAzureNodeRestarter(azCli azure.Client) NodeRestartProvider {
	return &azureNodeRestarter{azCli: azCli}
}

func (r *azureNodeRestarter) Restart(ctx context.Context, providerID string) error {
	subscription, resourceGroup, vmss, instanceID, err := parseAKSNodeProviderID(providerID)
	if err != nil {
		return err
	}
	return r.azCli.RestartVM(ctx, subscription, resourceGroup, vmss, instanceID)
}

// To use kindNodeRestarter, the cluster-agent should be run on host machine
type kindNodeRestarter struct{}

func NewKindNodeRestarter() NodeRestartProvider {
	return &kindNodeRestarter{}
}

func (r *kindNodeRestarter) Restart(ctx context.Context, providerID string) error {
	node, err := parseKindNodeProviderID(providerID)
	if err != nil {
		return err
	}
	return exec.Command("docker", "restart", node).Run()
}