package local

import (
	"fmt"
	"net/http"
	"net/url"
	"os"
	"sort"
	"sync"
	"sync/atomic"

	"github.com/pingcap/log"
	"go.uber.org/zap"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/kubernetes/scheme"
	restclient "k8s.io/client-go/rest"
	"k8s.io/client-go/tools/portforward"
	"k8s.io/client-go/transport/spdy"
)

type localForwardedPort int32

var nextForwardPort *localForwardedPort

func (c *localForwardedPort) increment() int32 {
	return atomic.AddInt32((*int32)(c), 1)
}

func init() {
	p := localForwardedPort(20000)
	nextForwardPort = &p
}

type PortForwardOptions struct {
	Namespace  string
	PodName    string
	Address    string
	LocalPort  uint16
	RemotePort uint16

	StopChannel  chan struct{}
	ReadyChannel chan struct{}
}

type PortForwarder struct {
	RESTClient *restclient.RESTClient
	Config     *restclient.Config
	client     *kubernetes.Clientset
}

func NewPortForwarder(cfg *restclient.Config) (*PortForwarder, error) {
	kubeClient, err := kubernetes.NewForConfig(cfg)
	if err != nil {
		return nil, err
	}

	cfg.GroupVersion = &schema.GroupVersion{Group: "", Version: "v1"}

	if cfg.APIPath == "" {
		cfg.APIPath = "/api"
	}
	if cfg.NegotiatedSerializer == nil {
		// This codec factory ensures the resources are not converted. Therefore, resources
		// will not be round-tripped through internal versions. Defaulting does not happen
		// on the client.
		cfg.NegotiatedSerializer = scheme.Codecs
	}
	RESTClient, err := restclient.RESTClientFor(cfg)
	if err != nil {
		return nil, err
	}
	return &PortForwarder{
		RESTClient: RESTClient,
		client:     kubeClient,
		Config:     cfg,
	}, nil
}

var forwardedPort = make(map[string]uint16)
var lock sync.RWMutex

func getServiceKey(namespace, name string, remotePort uint16) string {
	return fmt.Sprintf("service.%s.%s.%d", namespace, name, remotePort)
}
func getPodKey(namespace, name string, remotePort uint16) string {
	return fmt.Sprintf("pod.%s.%s.%d", namespace, name, remotePort)
}

func (f *PortForwarder) ForwardServicePort(namespace, name string, remotePort uint16) (uint16, error) {
	lock.Lock()
	defer lock.Unlock()
	if old, ok := forwardedPort[getServiceKey(namespace, name, remotePort)]; ok {
		return old, nil
	}

	service, err := f.client.CoreV1().Services(namespace).Get(name, metav1.GetOptions{})
	if err != nil {
		return 0, err
	}

	pod, err := f.attachablePodForService(service)
	if err != nil {
		return 0, err
	}
	log.Info("found service pod to attach", zap.String("name", pod.Name))

	localPort := uint16(nextForwardPort.increment())
	forwarder, err := f.getForwardForPod(pod, localPort, remotePort)
	if err != nil {
		return 0, err
	}
	go forwarder.ForwardPorts()
	forwardedPort[getServiceKey(namespace, name, remotePort)] = localPort
	return localPort, nil
}

func (f *PortForwarder) ForwardPodPort(namespace, podName string, remotePort uint16) (uint16, error) {
	lock.Lock()
	defer lock.Unlock()
	if old, ok := forwardedPort[getPodKey(namespace, podName, remotePort)]; ok {
		return old, nil
	}

	pod, err := f.client.CoreV1().Pods(namespace).Get(podName, metav1.GetOptions{})
	if err != nil {
		return 0, err
	}
	log.Info("found service pod to attach", zap.String("name", pod.Name))

	if pod.Status.Phase != corev1.PodRunning {
		return 0, fmt.Errorf("unable to forward port because pod is not running. Current status=%v", pod.Status.Phase)
	}
	localPort := uint16(nextForwardPort.increment())

	fw, err := f.getForwardForPod(pod, localPort, remotePort)
	if err != nil {
		return 0, err
	}
	go fw.ForwardPorts()
	forwardedPort[getPodKey(namespace, podName, remotePort)] = localPort
	return localPort, nil
}

func (f *PortForwarder) getForwarder(method string, url *url.URL, opts PortForwardOptions) (*portforward.PortForwarder, error) {
	transport, upgrader, err := spdy.RoundTripperFor(f.Config)
	if err != nil {
		return nil, err
	}
	dialer := spdy.NewDialer(upgrader, &http.Client{Transport: transport}, method, url)
	fw, err := portforward.NewOnAddresses(dialer, []string{opts.Address},
		[]string{fmt.Sprintf("%d:%d", opts.LocalPort, opts.RemotePort)}, opts.StopChannel, opts.ReadyChannel, os.Stdout, os.Stderr)
	if err != nil {
		return nil, err
	}
	return fw, nil
}

func (f *PortForwarder) getForwardForPod(pod *corev1.Pod, localPort, remotePort uint16) (*portforward.PortForwarder, error) {
	if pod.Status.Phase != corev1.PodRunning {
		return nil, fmt.Errorf("unable to forward port because pod is not running. Current status=%v", pod.Status.Phase)
	}

	options := PortForwardOptions{
		Namespace:  pod.Namespace,
		PodName:    pod.Name,
		Address:    "localhost",
		LocalPort:  localPort,
		RemotePort: remotePort,
	}

	req := f.RESTClient.Post().
		Resource("pods").
		Namespace(pod.Namespace).
		Name(pod.Name).
		SubResource("portforward")

	return f.getForwarder("POST", req.URL(), options)
}

//attachablePodForObject returns the pod to which to attach given an object.
func (f *PortForwarder) attachablePodForService(object *corev1.Service) (*corev1.Pod, error) {
	if object.Spec.Selector == nil || len(object.Spec.Selector) == 0 {
		return nil, fmt.Errorf("invalid service '%s': Service is defined without a selector", object.Name)
	}
	selector := labels.SelectorFromSet(object.Spec.Selector)

	sortBy := func(pods []*corev1.Pod) sort.Interface { return sort.Reverse(ActivePods(pods)) }
	pod, _, err := GetFirstPod(f.client.CoreV1(), object.Namespace, selector.String(), sortBy)
	return pod, err
}