package portforward

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/pingcap/errors"
	"github.com/pingcap/log"
	"go.uber.org/zap"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	restclient "k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
)

type activeForward struct {
	localPort int32
	stopChan  chan struct{}
	readyChan chan struct{}
}

type Manager struct {
	clientSet  *kubernetes.Clientset
	kubeConfig *restclient.Config
	namespace  string // 可以设置一个默认的 namespace
	active     map[string]*activeForward
	mu         sync.Mutex
}

func NewPortForwardManager(kubeconfigPath string) *Manager {
	kubeCfg, err := clientcmd.BuildConfigFromFlags("", kubeconfigPath)
	if err != nil {
		log.Fatal("initialize dataplane kubeconfig failed", zap.Error(err))
	}
	clientSet, err := kubernetes.NewForConfig(kubeCfg)
	if err != nil {
		log.Fatal("initialize dataplane kube client failed", zap.Error(err))
	}
	return &Manager{
		clientSet:  clientSet,
		kubeConfig: kubeCfg,
		namespace:  "default",
		active:     make(map[string]*activeForward),
	}
}

func (m *Manager) GetForwardedAddress(ctx context.Context,
	serviceName string,
	serviceNamespace string, servicePort int32) (string, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	key := fmt.Sprintf("%s/%s:%d", serviceNamespace, serviceName, servicePort)
	if fwd, ok := m.active[key]; ok {
		return fmt.Sprintf("localhost:%d", fwd.localPort), nil
	}

	if serviceNamespace == "" {
		serviceNamespace = m.namespace
	}

	service, err := m.clientSet.CoreV1().Services(serviceNamespace).Get(ctx, serviceName, metav1.GetOptions{})
	if err != nil {
		return "", errors.Trace(err)
	}

	pod, err := getFirstPod(ctx, m.clientSet, service)
	if err != nil {
		return "", errors.Trace(err)
	}

	remotePort, err := lookupContainerPortNumberByServicePort(service, pod, servicePort)
	if err != nil {
		return "", errors.Trace(err)
	}

	localPort := getRandomPort()
	readyChan := make(chan struct{})
	stopChan := make(chan struct{}, 1)

	fwd := &activeForward{
		localPort: int32(localPort),
		stopChan:  stopChan,
		readyChan: readyChan,
	}

	go func() {
		err := forwardPorts(m.clientSet, m.kubeConfig, serviceNamespace, pod.Name, int32(localPort), remotePort, stopChan, readyChan)
		if err != nil {
			// 在生产环境中应该使用更健壮的错误处理，例如 logging
			fmt.Printf("Port forwarding for %s failed: %v\n", key, err)
		}
	}()

	select {
	case <-readyChan:
		m.active[key] = fwd
		fmt.Printf("Port forward ready for %s -> localhost:%d\n", key, localPort)
		return fmt.Sprintf("localhost:%d", localPort), nil
	case <-time.After(10 * time.Second):
		return "", errors.New("timed out waiting for port forward to be ready")
	}
}

func (m *Manager) CloseAll() {
	m.mu.Lock()
	defer m.mu.Unlock()
	for key, fwd := range m.active {
		close(fwd.stopChan)
		delete(m.active, key)
		fmt.Printf("Closed port forward for %s\n", key)
	}
}

// KubeTransport 是一个实现了 http.RoundTripper 的结构
type KubeTransport struct {
	manager          *Manager
	defaultTransport http.RoundTripper
}

func NewKubeTransport(manager *Manager) http.RoundTripper {
	return &KubeTransport{
		manager: manager,
		defaultTransport: &http.Transport{
			Proxy: http.ProxyFromEnvironment,
			DialContext: (&net.Dialer{
				Timeout:   30 * time.Second,
				KeepAlive: 30 * time.Second,
			}).DialContext,
			ForceAttemptHTTP2:     true,
			MaxIdleConns:          100,
			IdleConnTimeout:       90 * time.Second,
			TLSHandshakeTimeout:   10 * time.Second,
			ExpectContinueTimeout: 1 * time.Second,
		},
	}
}

func (t *KubeTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	host, portStr, err := net.SplitHostPort(req.URL.Host)
	if err != nil {
		// default http/80, https/443
		host = req.URL.Host
		if req.URL.Scheme == "https" {
			portStr = "443"
		} else {
			portStr = "80"
		}
	}

	port, _ := strconv.Atoi(portStr)

	parts := strings.Split(host, ".")
	if len(parts) < 2 {
		return nil, errors.New("invalid k8s service fqdn")
	}
	serviceName := parts[0]
	serviceNamespace := parts[1]

	localAddr, err := t.manager.GetForwardedAddress(req.Context(), serviceName, serviceNamespace, int32(port))
	if err != nil {
		return nil, err
	}
	newReq := req.Clone(req.Context())
	newReq.URL.Host = localAddr
	newReq.Host = req.Host
	fmt.Printf("Forwarding request for %s to %s\n", req.URL.String(), newReq.URL.String())
	return t.defaultTransport.RoundTrip(newReq)
}