package handler

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/pingcap/log"
	clientgo "github.com/tidbcloud/tiflash-manager/pkg/client-go"
	"github.com/tidbcloud/tiflash-manager/pkg/config"
	"github.com/tidbcloud/tiflash-manager/pkg/core"
	"go.uber.org/zap"
	corev1 "k8s.io/api/core/v1"
)

type Server struct {
	ctx         context.Context
	pools       *core.Pools
	tenantsPool *core.TenantsPool
	conf        *config.Config
}

func NewServer(ctx context.Context, conf *config.Config, pools *core.Pools, tenantsPool *core.TenantsPool) *Server {
	return &Server{
		ctx,
		pools,
		tenantsPool,
		conf,
	}
}

func (s *Server) GetCluster(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")

	keyspaceID := r.URL.Query().Get(clientgo.KeyspaceIDKey)
	tenantToken := r.URL.Query().Get(clientgo.TenantTokenKey)
	if keyspaceID == "" && tenantToken == "" {
		resp := clientgo.Response{
			Error: clientgo.NewError(http.StatusBadRequest, "keyspace_id or tenant_token are required"),
		}
		writeResponse(w, http.StatusBadRequest, resp.String())
		return
	}

	var preferredPool string
	if tenantToken != "" {
		preferredPool = s.tenantsPool.GetKeyspacePool(tenantToken)
	} else if keyspaceID != "" {
		preferredPool = s.tenantsPool.GetKeyspacePool(keyspaceID)
	}

	poolName, ps, err := s.resolveClusterPods(r.Context(), preferredPool)
	if err != nil {
		resp := clientgo.Response{
			Error: clientgo.NewError(http.StatusServiceUnavailable, err.Error()),
		}
		writeResponse(w, http.StatusServiceUnavailable, resp.String())
		return
	}

	resp := clientgo.Response{
		Cluster: &clientgo.Cluster{
			Pool: poolName,
			Port: s.conf.Tiflash.Port,
			Pods: ps,
		},
	}

	writeResponse(w, http.StatusOK, resp.String())
}

func k8sPods2ClientPods(pods []*corev1.Pod, dnsFormat, pool string) []*clientgo.Pod {
	clientPods := make([]*clientgo.Pod, 0, len(pods))
	for _, pod := range pods {
		if pod.Status.PodIP == "" {
			log.Warn("Skipping pod without IP", zap.String("podName", pod.Name), zap.String("poolName", pool))
			continue
		}

		dns, err := getPodSvcDNS(pod, dnsFormat, pool)
		if err != nil {
			log.Warn("Skipping pod with invalid service DNS", zap.String("podName", pod.Name), zap.String("poolName", pool), zap.Error(err))
			continue
		}

		clientPods = append(clientPods, &clientgo.Pod{
			Name:    pod.Name,
			Address: dns,
		})
	}
	return clientPods
}

func getPodSvcDNS(pod *corev1.Pod, serviceDNSFormat string, poolName string) (string, error) {
	index, ok := pod.Labels["apps.kubernetes.io/pod-index"]
	if !ok {
		return "", fmt.Errorf("pod index not found")
	}

	return strings.ReplaceAll(strings.ReplaceAll(serviceDNSFormat, "{POOL}", poolName), "{INDEX}", index), nil
}

func (s *Server) Healthz(w http.ResponseWriter, r *http.Request) {
	writeResponse(w, http.StatusOK, "OK")
}

func (s *Server) DumpClusters(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	clusters := s.pools.GetPools(r.Context())
	if len(clusters) == 0 {
		resp := clientgo.DumpResponse{
			Error:     clientgo.NewError(http.StatusNotFound, "no clusters found"),
			Timestamp: time.Now().UnixNano(),
		}
		writeResponse(w, http.StatusNotFound, resp.String())
		return
	}
	var clientClusters []*clientgo.Cluster
	for poolName, pods := range clusters {
		ps := k8sPods2ClientPods(pods, s.conf.Tiflash.ServiceDNS, poolName)

		clientClusters = append(clientClusters, &clientgo.Cluster{
			Pool: poolName,
			Port: s.conf.Tiflash.Port,
			Pods: ps,
		})
	}
	resp := clientgo.DumpResponse{
		Clusters:  clientClusters,
		Timestamp: time.Now().UnixNano(),
	}
	writeResponse(w, http.StatusOK, resp.String())
}

func (s *Server) GetTopology(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	keyspaceID := r.FormValue("keyspace_id")
	rgName := r.FormValue("resource_group")
	if len(keyspaceID) == 0 {
		resp := clientgo.NewTopologyResult(errors.New("keyspace id cannot be empty"), nil)
		writeResponse(w, http.StatusBadRequest, resp.String())
		return
	}

	ip, _ := getIP(r)
	traceID := r.FormValue("trace_id")

	log.Info("GetTopology called",
		zap.String("keyspaceID", keyspaceID),
		zap.String("rgName", rgName),
		zap.String("traceID", r.FormValue("trace_id")),
		zap.String("clientIP", ip),
	)

	ctx := context.WithValue(r.Context(), "trace_id", traceID)
	ctx = context.WithValue(ctx, "client_ip", ip)

	preferredPool := s.tenantsPool.GetKeyspacePool(keyspaceID)
	_, topo, err := s.resolveTopology(ctx, preferredPool)
	if err != nil {
		resp := clientgo.NewTopologyResult(err, nil)
		writeResponse(w, http.StatusServiceUnavailable, resp.String())
		return
	}
	resp := clientgo.NewTopologyResult(nil, topo)
	writeResponse(w, http.StatusOK, resp.String())
}

// resolveClusterPods returns the first pool that has at least one valid store.
// If the preferred pool is empty or unhealthy, the default pool is used as fallback.
func (s *Server) resolveClusterPods(ctx context.Context, preferredPool string) (string, []*clientgo.Pod, error) {
	return firstValidPool(preferredPool, s.conf.Tiflash.DefaultPool, func(poolName string) ([]*clientgo.Pod, bool, error) {
		pods, ok := s.pools.GetPool(ctx, poolName)
		if !ok {
			return nil, false, nil
		}

		clientPods := k8sPods2ClientPods(pods, s.conf.Tiflash.ServiceDNS, poolName)
		if len(clientPods) == 0 {
			log.Warn("Pool has no valid stores, trying next pool", zap.String("poolName", poolName))
			return nil, false, nil
		}
		return clientPods, true, nil
	})
}

func (s *Server) resolveTopology(ctx context.Context, preferredPool string) (string, []string, error) {
	return firstValidPool(preferredPool, s.conf.Tiflash.DefaultPool, func(poolName string) ([]string, bool, error) {
		pods, ok := s.pools.GetPool(ctx, poolName)
		if !ok {
			return nil, false, nil
		}

		clientPods := k8sPods2ClientPods(pods, s.conf.Tiflash.ServiceDNS, poolName)
		if len(clientPods) == 0 {
			log.Warn("Pool has no valid topology stores, trying next pool", zap.String("poolName", poolName))
			return nil, false, nil
		}

		topo := make([]string, 0, len(clientPods))
		for _, pod := range clientPods {
			topo = append(topo, fmt.Sprintf("%s:%v", pod.Address, s.conf.Tiflash.Port))
		}
		return topo, true, nil
	})
}

func candidatePools(preferredPool, defaultPool string) []string {
	if preferredPool == "" || preferredPool == defaultPool {
		return []string{defaultPool}
	}
	return []string{preferredPool, defaultPool}
}

func firstValidPool[T any](preferredPool, defaultPool string, build func(poolName string) (T, bool, error)) (string, T, error) {
	var zero T
	var lastErr error

	for _, poolName := range candidatePools(preferredPool, defaultPool) {
		value, ok, err := build(poolName)
		if err != nil {
			lastErr = err
			continue
		}
		if !ok {
			continue
		}
		return poolName, value, nil
	}

	if lastErr != nil {
		return "", zero, fmt.Errorf("no valid store found in preferred pool or default pool: %w", lastErr)
	}
	return "", zero, fmt.Errorf("no valid store found in preferred pool or default pool")
}

func writeResponse(w http.ResponseWriter, status int, message string) {
	w.WriteHeader(status)
	io.WriteString(w, message)
}

func getIP(r *http.Request) (string, error) {
	ips := r.Header.Get("X-Forwarded-For")
	splitIps := strings.Split(ips, ",")

	if len(splitIps) > 0 {
		// get last IP in list since ELB prepends other user defined IPs, meaning the last one is the actual client IP.
		netIP := net.ParseIP(splitIps[len(splitIps)-1])
		if netIP != nil {
			return netIP.String(), nil
		}
	}

	ip, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return "", err
	}

	netIP := net.ParseIP(ip)
	if netIP != nil {
		ip := netIP.String()
		if ip == "::1" {
			return "127.0.0.1", nil
		}
		return ip, nil
	}

	return "", errors.New("IP not found")
}