package exchange

import (
	"fmt"
	"github.com/alibaba/kt-connect/pkg/kt/command/general"
	opt "github.com/alibaba/kt-connect/pkg/kt/command/options"
	"github.com/alibaba/kt-connect/pkg/kt/service/cluster"
	"github.com/alibaba/kt-connect/pkg/kt/service/sshchannel"
	"github.com/alibaba/kt-connect/pkg/kt/transmission"
	"github.com/alibaba/kt-connect/pkg/kt/util"
	"github.com/rs/zerolog/log"
	coreV1 "k8s.io/api/core/v1"
	"strconv"
	"strings"
	"time"
)

func ByEphemeralContainer(resourceName string) error {
	log.Warn().Msgf("Experimental feature. It just works on kubernetes above v1.23, and it can NOT work with istio.")

	pods, err := getPodsOfResource(resourceName, opt.Get().Global.Namespace)

	for _, pod := range pods {
		if pod.Status.Phase != coreV1.PodRunning {
			log.Warn().Msgf("Pod %s is not running (%s), will not be exchanged", pod.Name, pod.Status.Phase)
			continue
		}
		privateKey, err2 := createEphemeralContainer(util.KtExchangeContainer, pod.Name)
		if err2 != nil {
			return err2
		}

		// record data
		opt.Store.Shadow = util.Append(opt.Store.Shadow, pod.Name)

		localSSHPort, err2 := transmission.ForwardPodToLocal(opt.Get().Exchange.Expose, pod.Name, privateKey)
		if err2 != nil {
			return err2
		}
		err = exchangeWithEphemeralContainer(opt.Get().Exchange.Expose, localSSHPort, privateKey)
		if err != nil {
			return err
		}
	}
	return nil
}


func getPodsOfResource(resourceName, namespace string) ([]coreV1.Pod, error) {
	resourceType, name, err := general.ParseResourceName(resourceName)
	if err != nil {
		return nil, err
	}

	switch resourceType {
	case "pod":
		pod, err := cluster.Ins().GetPod(name, namespace)
		if err != nil {
			return nil, err
		} else {
			return []coreV1.Pod{*pod}, nil
		}
	case "svc":
		fallthrough
	case "service":
		return getPodsOfService(name, namespace)
	}
	return nil, fmt.Errorf("invalid resource type: %s", resourceType)
}

func getPodsOfService(serviceName, namespace string) ([]coreV1.Pod, error) {
	svc, err := cluster.Ins().GetService(serviceName, namespace)
	if err != nil {
		return nil, err
	}
	pods, err := cluster.Ins().GetPodsByLabel(svc.Spec.Selector, namespace)
	if err != nil {
		return nil, err
	}
	return pods.Items, nil
}

func createEphemeralContainer(containerName, podName string) (string, error) {
	log.Info().Msgf("Adding ephemeral container for pod %s", podName)

	envs := make(map[string]string)
	privateKey, err := cluster.Ins().AddEphemeralContainer(containerName, podName, envs)
	if err != nil {
		return "", err
	}

	for i := 0; i < 10; i++ {
		log.Info().Msgf("Waiting for ephemeral container %s to be ready", containerName)
		ready, err2 := isEphemeralContainerReady(containerName, podName, opt.Get().Global.Namespace)
		if err2 != nil {
			return "", err2
		} else if ready {
			break
		}
		time.Sleep(5 * time.Second)
	}
	return privateKey, nil
}

func isEphemeralContainerReady(podName, containerName, namespace string) (bool, error) {
	pod, err := cluster.Ins().GetPod(podName, namespace)
	if err != nil {
		return false, err
	}
	cStats := pod.Status.EphemeralContainerStatuses
	for i := range cStats {
		if cStats[i].Name == containerName {
			if cStats[i].State.Running != nil {
				return true, nil
			} else if cStats[i].State.Terminated != nil {
				return false, fmt.Errorf("ephemeral container %s is terminated, code: %d",
					containerName, cStats[i].State.Terminated.ExitCode)
			}
		}
	}
	return false, nil
}

func exchangeWithEphemeralContainer(exposePorts string, localSSHPort int, privateKey string) error {
	// Get all listened ports on remote host
	listenedPorts, err := getListenedPorts(localSSHPort, privateKey)
	if err != nil {
		return err
	}

	redirectPorts, err := remoteRedirectPort(exposePorts, listenedPorts)
	if err != nil {
		return err
	}
	var redirectPortStr string
	for k, v := range redirectPorts {
		redirectPortStr += fmt.Sprintf("%d:%d,", k, v)
	}
	redirectPortStr = redirectPortStr[:len(redirectPortStr)-1]
	err = setupIptables(redirectPortStr, localSSHPort, privateKey)
	if err != nil {
		return err
	}

	if err = transmission.ForwardRemotePortsViaSshTunnel(exposePorts, localSSHPort, privateKey); err != nil {
		return err
	}

	return nil
}


func setupIptables(redirectPorts string, localSSHPort int, privateKey string) error {
	res, err := sshchannel.Ins().RunScript(
		privateKey,
		fmt.Sprintf("127.0.0.1:%d", localSSHPort),
		fmt.Sprintf("/setup_iptables.sh %s", redirectPorts))

	if err != nil {
		log.Error().Err(err).Msgf("Setup iptables failed, error")
	}

	log.Debug().Msgf("Run setup iptables result: %s", res)
	return err
}

func getListenedPorts(localSSHPort int, privateKey string) (map[int]struct{}, error) {
	result, err := sshchannel.Ins().RunScript(
		privateKey,
		fmt.Sprintf("127.0.0.1:%d", localSSHPort),
		`netstat -tuln | grep -E '^(tcp|udp|tcp6)' | grep LISTEN | awk '{print $4}' | awk -F: '{printf("%s\n", $NF)}'`)

	if err != nil {
		return nil, err
	}

	log.Debug().Msgf("Run get listened ports result: %s", result)
	var listenedPorts = make(map[int]struct{})
	// The result should be a string like
	// 38059
	// 22
	parts := strings.Split(result, "\n")
	for i := range parts {
		if len(parts[i]) > 0 {
			port, err2 := strconv.Atoi(parts[i])
			if err2 != nil {
				log.Warn().Err(err2).Msgf("Failed to fetch listened ports (got '%s')", parts[i])
				continue
			}
			listenedPorts[port] = struct{}{}
		}
	}

	return listenedPorts, nil
}

func remoteRedirectPort(exposePorts string, listenedPorts map[int]struct{}) (map[int]int, error) {
	portPairs := strings.Split(exposePorts, ",")
	redirectPort := make(map[int]int)
	for _, exposePort := range portPairs {
		_, remotePort, err := util.ParsePortMapping(exposePort)
		if err != nil {
			return nil, err
		}
		port := randPort(listenedPorts)
		if port == -1 {
			return nil, fmt.Errorf("failed to find redirect port for port: %d", remotePort)
		}
		redirectPort[remotePort] = port
	}

	return redirectPort, nil
}

func randPort(listenedPorts map[int]struct{}) int {
	for i := 0; i < 100; i++ {
		port := util.RandomPort()
		if _, exists := listenedPorts[port]; !exists {
			return port
		}
	}
	return -1
}

