package proxy

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net/http"
	"net/http/httputil"
	"net/url"
	"sort"
	"strings"

	"github.com/gin-gonic/gin"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"
	corev1 "k8s.io/api/core/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/yaml"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	regionalpb "github.com/tidbcloud/infra-provider/api/tidbcloud/regional-server/v1"
	"github.com/tidbcloud/infra-provider/pkg/cloud/gcp"
	gkestatus "github.com/tidbcloud/infra-provider/pkg/gkeheterogeneous/status"
)

const (
	regionalServerAPIPathPrefix = "/api/dedicated/v1"
	clusterAgentGatewayPrefix   = "/cluster-agent"
	clusterAgentAPIPathPrefix   = "/api/v1"
)

type ClusterAgentProxy struct {
	kubeClient                   client.Client
	transport                    http.RoundTripper
	policyControlResourceManager gcp.ResourceManagerClient
}

func NewClusterAgentProxy(kubeClient client.Client) *ClusterAgentProxy {
	return &ClusterAgentProxy{
		kubeClient: kubeClient,
		transport:  customizedTransport(),
	}
}

func (p *ClusterAgentProxy) WithGKEPolicyControlResourceManager(resourceManager gcp.ResourceManagerClient) *ClusterAgentProxy {
	p.policyControlResourceManager = resourceManager
	return p
}

func (p *ClusterAgentProxy) ProxyGKEHeterogeneousDeploy(rp *ReverseProxyWrapper, list bool) gin.HandlerFunc {
	return func(c *gin.Context) {
		if rp == nil || rp.ReverseProxy == nil {
			c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": "cluster service proxy is not configured"})
			return
		}

		origin := rp.ReverseProxy
		proxy := *origin
		originModifyResponse := origin.ModifyResponse
		clusterID := strings.TrimSpace(c.Param("clusterId"))
		proxy.ModifyResponse = func(resp *http.Response) error {
			if originModifyResponse != nil {
				if err := originModifyResponse(resp); err != nil {
					return err
				}
			}
			return p.markGKEHeterogeneousDeployReady(c.Request.Context(), clusterID, resp, list)
		}
		proxy.ServeHTTP(c.Writer, c.Request)
	}
}

func (p *ClusterAgentProxy) markGKEHeterogeneousDeployReady(ctx context.Context, clusterID string, resp *http.Response, list bool) error {
	if p == nil || resp == nil || resp.Body == nil || resp.StatusCode != http.StatusOK {
		return nil
	}
	originalBody := resp.Body
	body, err := io.ReadAll(originalBody)
	if err != nil {
		return nil
	}
	_ = originalBody.Close()
	restoreHTTPResponseBody(resp, body)
	if len(bytes.TrimSpace(body)) == 0 {
		return nil
	}

	if list {
		deploys := &regionalpb.ListGKEHeterogeneousDeploysResponse{}
		if err := (protojson.UnmarshalOptions{DiscardUnknown: true}).Unmarshal(body, deploys); err != nil {
			return nil
		}
		statusSnapshot := p.gkeHeterogeneousStatusSnapshot(ctx, clusterID)
		changed := false
		for _, deploy := range deploys.GetDeploys() {
			if p.markGKEHeterogeneousDeployMessageReady(ctx, clusterID, deploy) {
				changed = true
			}
			if markGKEHeterogeneousDeployPodCounts(statusSnapshot, deploy) {
				changed = true
			}
		}
		if !changed {
			return nil
		}
		nextBody, err := (protojson.MarshalOptions{UseProtoNames: false}).Marshal(deploys)
		if err != nil {
			return nil
		}
		restoreHTTPResponseBody(resp, nextBody)
		return nil
	}

	deploy := &regionalpb.GKEHeterogeneousDeploy{}
	if err := (protojson.UnmarshalOptions{DiscardUnknown: true}).Unmarshal(body, deploy); err != nil {
		return nil
	}
	statusChanged := p.markGKEHeterogeneousDeployMessageReady(ctx, clusterID, deploy)
	podCountsChanged := markGKEHeterogeneousDeployPodCounts(p.gkeHeterogeneousStatusSnapshot(ctx, clusterID), deploy)
	if !statusChanged && !podCountsChanged {
		return nil
	}
	nextBody, err := (protojson.MarshalOptions{UseProtoNames: false}).Marshal(deploy)
	if err != nil {
		return nil
	}
	restoreHTTPResponseBody(resp, nextBody)
	return nil
}

func (p *ClusterAgentProxy) gkeHeterogeneousStatusSnapshot(ctx context.Context, clusterID string) *gkestatus.Snapshot {
	if p == nil || p.kubeClient == nil {
		return nil
	}
	cluster, _, errMsg := p.resolveGKEClusterForAgentRequest(ctx, clusterID)
	if errMsg != "" {
		return nil
	}
	cm := &corev1.ConfigMap{}
	if err := p.kubeClient.Get(ctx, client.ObjectKey{
		Namespace: cluster.Namespace,
		Name:      gkestatus.ConfigMapNamePrefix + cluster.Name,
	}, cm); err != nil {
		return nil
	}
	data := strings.TrimSpace(cm.Data[gkestatus.ConfigMapDataKey])
	if data == "" {
		return nil
	}
	snapshot := &gkestatus.Snapshot{}
	if err := yaml.Unmarshal([]byte(data), snapshot); err != nil {
		return nil
	}
	if !gkeHeterogeneousStatusSnapshotUsable(snapshot) {
		return nil
	}
	return snapshot
}

func gkeHeterogeneousStatusSnapshotUsable(snapshot *gkestatus.Snapshot) bool {
	return snapshot != nil &&
		snapshot.Enabled &&
		snapshot.ValidIntent &&
		snapshot.ProbeStatus == gkestatus.ProbeStatusSuccess
}

func markGKEHeterogeneousDeployPodCounts(snapshot *gkestatus.Snapshot, deploy *regionalpb.GKEHeterogeneousDeploy) bool {
	if snapshot == nil || deploy == nil || !deploy.GetEnabled() || strings.TrimSpace(deploy.GetComputeClassName()) == "" {
		return false
	}
	for _, component := range snapshot.Components {
		if !gkeHeterogeneousStatusComponentMatchesDeploy(component, deploy) {
			continue
		}
		podCounts := regionalGKEHeterogeneousDeployPodCounts(component)
		if len(podCounts) == 0 {
			return false
		}
		if gkeHeterogeneousDeployPodCountsEqual(deploy.GetPodCounts(), podCounts) {
			return false
		}
		deploy.PodCounts = podCounts
		return true
	}
	return false
}

func gkeHeterogeneousDeployPodCountsEqual(a, b []*regionalpb.GKEHeterogeneousDeployPodCount) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if !proto.Equal(a[i], b[i]) {
			return false
		}
	}
	return true
}

func gkeHeterogeneousStatusComponentMatchesDeploy(component gkestatus.ComponentStatus, deploy *regionalpb.GKEHeterogeneousDeploy) bool {
	if strings.TrimSpace(component.ComputeClassName) != strings.TrimSpace(deploy.GetComputeClassName()) {
		return false
	}
	if strings.TrimSpace(component.TiDBGroupID) != strings.TrimSpace(deploy.GetTidbGroupId()) {
		return false
	}
	switch deploy.GetComponentType() {
	case regionalpb.ComponentType_TIKV:
		return strings.EqualFold(strings.TrimSpace(component.ComponentType), string(cloudv1.ComponentTiKV))
	case regionalpb.ComponentType_TIDB:
		return strings.EqualFold(strings.TrimSpace(component.ComponentType), string(cloudv1.ComponentTiDB))
	default:
		return false
	}
}

func regionalGKEHeterogeneousDeployPodCounts(component gkestatus.ComponentStatus) []*regionalpb.GKEHeterogeneousDeployPodCount {
	counts := make(map[string]int32, len(component.NodeUsages))
	for _, usage := range component.NodeUsages {
		counts[gkeHeterogeneousMachineType(usage.MachineType)] += usage.PodCount
	}
	if len(counts) == 0 {
		for _, group := range component.PodGroups {
			counts[gkeHeterogeneousMachineType(group.ActualInstanceType)] += group.PodCount
		}
	}

	machineTypes := make([]string, 0, len(counts))
	for machineType := range counts {
		machineTypes = append(machineTypes, machineType)
	}
	sort.Strings(machineTypes)

	out := make([]*regionalpb.GKEHeterogeneousDeployPodCount, 0, len(machineTypes))
	for _, machineType := range machineTypes {
		out = append(out, &regionalpb.GKEHeterogeneousDeployPodCount{
			MachineType: machineType,
			PodCount:    counts[machineType],
		})
	}
	return out
}

func gkeHeterogeneousMachineType(machineType string) string {
	machineType = strings.TrimSpace(machineType)
	if machineType == "" {
		return "unknown"
	}
	return machineType
}

func (p *ClusterAgentProxy) markGKEHeterogeneousDeployMessageReady(ctx context.Context, clusterID string, deploy *regionalpb.GKEHeterogeneousDeploy) bool {
	if !p.gkeHeterogeneousDeployReadyFromClusterAgent(ctx, clusterID, deploy) {
		return false
	}
	deploy.Status = regionalpb.GKEHeterogeneousDeployStatus_GKE_HETEROGENEOUS_DEPLOY_STATUS_HETEROGENEOUS_READY
	return true
}

func restoreHTTPResponseBody(resp *http.Response, body []byte) {
	resp.Body = io.NopCloser(bytes.NewReader(body))
	resp.ContentLength = int64(len(body))
	resp.Header.Set("Content-Length", fmt.Sprintf("%d", len(body)))
}

func (p *ClusterAgentProxy) gkeHeterogeneousDeployReadyFromClusterAgent(ctx context.Context, clusterID string, deploy *regionalpb.GKEHeterogeneousDeploy) bool {
	if deploy == nil || !deploy.GetEnabled() {
		return false
	}
	if deploy.GetStatus() != regionalpb.GKEHeterogeneousDeployStatus_GKE_HETEROGENEOUS_DEPLOY_STATUS_READY_FOR_GATEKEEPER {
		return false
	}
	computeClassName := strings.TrimSpace(deploy.GetComputeClassName())
	if computeClassName == "" {
		return false
	}
	cluster, _, errMsg := p.resolveGKEClusterForAgentGateway(ctx, clusterID)
	if errMsg != "" {
		return false
	}
	return p.gkeHeterogeneousComputeClassReady(ctx, cluster, computeClassName) &&
		p.gkeHeterogeneousGatekeeperAssignReady(ctx, cluster, deploy)
}

func (p *ClusterAgentProxy) gkeHeterogeneousComputeClassReady(ctx context.Context, cluster *cloudv1.Cluster, computeClassName string) bool {
	query := url.Values{}
	query.Set("name", computeClassName)
	out := &regionalpb.GetGKEHeterogeneousComputeClassResponse{}
	if !p.getClusterAgentProto(ctx, cluster, fmt.Sprintf("/clusters/%s/gkeHeterogeneous/computeClasses", cluster.Name), query, out) {
		return false
	}
	for _, computeClass := range out.GetComputeClasses() {
		if strings.TrimSpace(computeClass.GetName()) == computeClassName {
			return true
		}
	}
	return false
}

func (p *ClusterAgentProxy) gkeHeterogeneousGatekeeperAssignReady(ctx context.Context, cluster *cloudv1.Cluster, deploy *regionalpb.GKEHeterogeneousDeploy) bool {
	query := url.Values{}
	if deploy.GetComponentType() == regionalpb.ComponentType_TIDB && strings.TrimSpace(deploy.GetTidbGroupId()) != "" {
		query.Set("tidbGroupId", strings.TrimSpace(deploy.GetTidbGroupId()))
	}
	out := &regionalpb.GKEHeterogeneousGatekeeperAssign{}
	path := fmt.Sprintf("/clusters/%s/gkeHeterogeneous/gatekeeperAssigns/%s", cluster.Name, deploy.GetComponentType().String())
	if !p.getClusterAgentProto(ctx, cluster, path, query, out) {
		return false
	}
	return out.GetComponentType() == deploy.GetComponentType() &&
		strings.TrimSpace(out.GetTidbGroupId()) == strings.TrimSpace(deploy.GetTidbGroupId()) &&
		strings.TrimSpace(out.GetComputeClassName()) == strings.TrimSpace(deploy.GetComputeClassName())
}

func (p *ClusterAgentProxy) getClusterAgentProto(ctx context.Context, cluster *cloudv1.Cluster, suffix string, query url.Values, out proto.Message) bool {
	if p == nil || p.transport == nil || cluster == nil || out == nil {
		return false
	}
	target, err := url.Parse(strings.TrimSpace(cluster.Status.Gateway.Endpoint))
	if err != nil || target.Scheme == "" || target.Host == "" {
		return false
	}
	agentPrefix := fmt.Sprintf("%s/%s%s", clusterAgentGatewayPrefix, cluster.PrefixedId(), clusterAgentAPIPathPrefix)
	target.Path = joinURLPath(target.EscapedPath(), joinURLPath(agentPrefix, suffix))
	target.RawPath = ""
	target.RawQuery = query.Encode()

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, target.String(), nil)
	if err != nil {
		return false
	}
	resp, err := (&http.Client{Transport: p.transport}).Do(req)
	if err != nil {
		return false
	}
	defer func() { _ = resp.Body.Close() }()
	if resp.StatusCode != http.StatusOK {
		return false
	}
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return false
	}
	return (protojson.UnmarshalOptions{DiscardUnknown: true}).Unmarshal(body, out) == nil
}

func (p *ClusterAgentProxy) Proxy(c *gin.Context) {
	if p == nil || p.kubeClient == nil {
		c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": "cluster agent proxy is not configured"})
		return
	}

	clusterID := strings.TrimSpace(c.Param("clusterId"))
	if clusterID == "" {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "cluster_id is required"})
		return
	}
	cluster, statusCode, errMsg := p.resolveGKEClusterForAgentRequest(c.Request.Context(), clusterID)
	if errMsg != "" {
		c.AbortWithStatusJSON(statusCode, gin.H{"error": errMsg})
		return
	}
	if statusCode, errMsg := validateGKEClusterAgentGateway(cluster, clusterID); errMsg != "" {
		c.AbortWithStatusJSON(statusCode, gin.H{"error": errMsg})
		return
	}
	if statusCode, errMsg := rewriteGKEHeterogeneousGatekeeperAssignKind(c, cluster); errMsg != "" {
		c.AbortWithStatusJSON(statusCode, gin.H{"error": errMsg})
		return
	}
	if p.rejectGKEHeterogeneousGatekeeperAssignNotReady(c, cluster) {
		return
	}
	if p.rejectDeleteGKEHeterogeneousComputeClassInUse(c, cluster) {
		return
	}

	dropTiDBGroupQueryForTiKVGatekeeperAssign(c)
	p.proxyToCluster(c, cluster, clusterID)
}

func (p *ClusterAgentProxy) proxyToCluster(c *gin.Context, cluster *cloudv1.Cluster, publicClusterID string) {
	endpoint := strings.TrimSpace(cluster.Status.Gateway.Endpoint)
	target, err := url.Parse(endpoint)
	if err != nil {
		c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("parse cluster gateway endpoint: %s", err)})
		return
	}
	if target.Scheme == "" || target.Host == "" {
		c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": "cluster gateway endpoint must include scheme and host"})
		return
	}

	rp := &httputil.ReverseProxy{
		Transport: p.transport,
		Director: func(req *http.Request) {
			rewriteClusterAgentRequest(req, target, publicClusterID, cluster.Name, cluster.PrefixedId())
		},
	}
	rp.ServeHTTP(c.Writer, c.Request)
}

// resolveGKEClusterForAgentRequest resolves the Cluster CR before gateway
// validation so regional-only preflight can inspect Cluster metadata. Proxy paths
// that need actual dataplane state must still validate the cluster-agent gateway.
func (p *ClusterAgentProxy) resolveGKEClusterForAgentRequest(ctx context.Context, clusterID string) (*cloudv1.Cluster, int, string) {
	if p == nil || p.kubeClient == nil {
		return nil, http.StatusInternalServerError, "cluster agent proxy is not configured"
	}

	var clusters cloudv1.ClusterList
	if err := p.kubeClient.List(ctx, &clusters, client.MatchingLabels{cloudv1.ClusterTagKey: clusterID}); err != nil {
		return nil, http.StatusInternalServerError, fmt.Sprintf("list cluster %s: %s", clusterID, err)
	}
	if len(clusters.Items) == 0 {
		return nil, http.StatusNotFound, fmt.Sprintf("cluster %s not found", clusterID)
	}
	if len(clusters.Items) > 1 {
		return nil, http.StatusInternalServerError, fmt.Sprintf("found more than one cluster with id %s", clusterID)
	}

	cluster := clusters.Items[0]
	if cluster.GetProviderType() != cloudv1.ProviderTypeGKE {
		return nil, http.StatusPreconditionFailed, fmt.Sprintf("cluster %s is not GKE", clusterID)
	}
	return &cluster, 0, ""
}

// resolveGKEClusterForAgentGateway resolves the Cluster CR and verifies that the
// request can be forwarded through the cluster-agent gateway.
func (p *ClusterAgentProxy) resolveGKEClusterForAgentGateway(ctx context.Context, clusterID string) (*cloudv1.Cluster, int, string) {
	cluster, statusCode, errMsg := p.resolveGKEClusterForAgentRequest(ctx, clusterID)
	if errMsg != "" {
		return nil, statusCode, errMsg
	}
	if statusCode, errMsg := validateGKEClusterAgentGateway(cluster, clusterID); errMsg != "" {
		return nil, statusCode, errMsg
	}
	return cluster, 0, ""
}

func validateGKEClusterAgentGateway(cluster *cloudv1.Cluster, clusterID string) (int, string) {
	if cluster == nil {
		return http.StatusInternalServerError, "cluster is required"
	}
	if cluster.Status.Gateway == nil || strings.TrimSpace(cluster.Status.Gateway.Endpoint) == "" {
		return http.StatusPreconditionFailed, fmt.Sprintf("cluster %s gateway endpoint is not ready", clusterID)
	}
	return 0, ""
}

func rewriteGKEHeterogeneousGatekeeperAssignKind(c *gin.Context, cluster *cloudv1.Cluster) (int, string) {
	if c.Request.Method != http.MethodPost || !strings.HasSuffix(c.FullPath(), "/gkeHeterogeneous/gatekeeperAssigns") {
		return 0, ""
	}
	body, err := io.ReadAll(c.Request.Body)
	if err != nil {
		return http.StatusBadRequest, fmt.Sprintf("read gatekeeper assign request body: %s", err)
	}
	c.Request.Body = io.NopCloser(bytes.NewReader(body))
	if len(bytes.TrimSpace(body)) == 0 {
		return http.StatusBadRequest, "gatekeeper assign request body is required"
	}

	req := &regionalpb.CreateGKEHeterogeneousGatekeeperAssignRequest{}
	if err := (protojson.UnmarshalOptions{DiscardUnknown: true}).Unmarshal(body, req); err != nil {
		return http.StatusBadRequest, fmt.Sprintf("decode gatekeeper assign request body: %s", err)
	}

	if req.GetComponentType() != regionalpb.ComponentType_TIDB {
		return 0, ""
	}

	groupID := strings.TrimSpace(req.GetTidbGroupId())
	expectedKind, statusCode, errMsg := expectedGKEHeterogeneousTiDBGroupKind(cluster, groupID)
	if errMsg != "" {
		return statusCode, errMsg
	}
	if req.GetTidbGroupKind() != regionalpb.TiDBGroupKind_TIDB_GROUP_KIND_UNSPECIFIED && req.GetTidbGroupKind() != expectedKind {
		return http.StatusBadRequest, fmt.Sprintf("tidbGroupKind does not match tidbGroupId %s, want %s", groupID, expectedKind)
	}
	req.TidbGroupKind = expectedKind
	body, err = (protojson.MarshalOptions{UseProtoNames: false}).Marshal(req)
	if err != nil {
		return http.StatusInternalServerError, fmt.Sprintf("encode gatekeeper assign request body: %s", err)
	}
	c.Request.Body = io.NopCloser(bytes.NewReader(body))
	c.Request.ContentLength = int64(len(body))
	c.Request.Header.Set("Content-Length", fmt.Sprintf("%d", len(body)))
	return 0, ""
}

func expectedGKEHeterogeneousTiDBGroupKind(cluster *cloudv1.Cluster, groupID string) (regionalpb.TiDBGroupKind, int, string) {
	groupID = strings.TrimSpace(groupID)
	if groupID == "" {
		return regionalpb.TiDBGroupKind_TIDB_GROUP_KIND_UNSPECIFIED, http.StatusBadRequest, "tidbGroupId is required for TiDB gatekeeper assign"
	}
	if cluster == nil {
		return regionalpb.TiDBGroupKind_TIDB_GROUP_KIND_UNSPECIFIED, http.StatusInternalServerError, "cluster is required"
	}
	if groupID == strings.TrimSpace(cluster.GetDefaultGroupID()) {
		return regionalpb.TiDBGroupKind_TIDB_GROUP_KIND_DEFAULT, 0, ""
	}
	if _, ok := cluster.GetExtraGroupSpec(groupID); ok {
		return regionalpb.TiDBGroupKind_TIDB_GROUP_KIND_EXTRA, 0, ""
	}
	return regionalpb.TiDBGroupKind_TIDB_GROUP_KIND_UNSPECIFIED, http.StatusBadRequest, fmt.Sprintf("TiDB group %s not found in cluster %s", groupID, cluster.Name)
}

func rewriteClusterAgentRequest(req *http.Request, target *url.URL, publicClusterID, agentClusterID, agentNamespace string) {
	req.URL.Scheme = target.Scheme
	req.URL.Host = target.Host
	req.Host = target.Host

	// The public cluster id is only stable at the regional boundary. The
	// cluster-agent gateway and agent-side namespace checks use Cluster.Name.
	suffix := strings.TrimPrefix(req.URL.Path, regionalServerAPIPathPrefix)
	suffix = strings.Replace(suffix, fmt.Sprintf("/clusters/%s", publicClusterID), fmt.Sprintf("/clusters/%s", agentClusterID), 1)
	agentPrefix := fmt.Sprintf("%s/%s%s", clusterAgentGatewayPrefix, agentNamespace, clusterAgentAPIPathPrefix)
	req.URL.Path = joinURLPath(target.EscapedPath(), joinURLPath(agentPrefix, suffix))
	req.URL.RawPath = ""
}

func joinURLPath(a, b string) string {
	if a == "" {
		return b
	}
	if b == "" {
		return a
	}
	aSlash := strings.HasSuffix(a, "/")
	bSlash := strings.HasPrefix(b, "/")
	switch {
	case aSlash && bSlash:
		return a + b[1:]
	case !aSlash && !bSlash:
		return a + "/" + b
	default:
		return a + b
	}
}