// package cluster provides services and proxies related to the Kubernetes cluster.
package cluster

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/patrickmn/go-cache"
	"github.com/pingcap/errors"

	"github.com/tidbcloud/clusterproxy/resolve"
	"github.com/tidbcloud/clusterproxy/route"
	"github.com/tidbcloud/clusterproxy/route/router"
	"github.com/tidbcloud/clusterproxy/util"
	"github.com/tidbcloud/clusterproxy/util/cluster"
)

const (
	pdServiceFormat = "%s-pd"
	pdPort          = "2379"
	httpScheme      = "http"
	httpsScheme     = "https"
)

func init() {
	router.RegisterRoute(pdRoute)
}

var pdRoute = &route.Route{
	Name:    "pd",
	Path:    "pd",
	Version: route.V1,
	Schema: &route.Schema{
		PathRegex:       `.*`,
		RequiredHeaders: []string{},
		JSONBodyModel:   nil,
	},
	GetHandler: pdGetHandler,
}

type AuthType int

const (
	AuthTypeSQLUser AuthType = iota
	AuthTypeSharingCode

	dashboardLoginPath     = "/dashboard/api/user/login"
	dashboardLoginInfoPath = "/dashboard/api/user/login_info"
	tlsClientSecretFormat  = "%s-cluster-client-secret"
)

var (
	tokenCache = cache.New(20*time.Hour, 10*time.Minute)
)

type dashboardAuthPayload struct {
	Type     AuthType `json:"type"`
	Username string   `json:"username"`
	Password string   `json:"password"`
}

func pdGetHandler(c *route.Context, cl resolve.Cluster) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		scheme := httpsScheme
		namespace := cl.Namespace
		pdServiceName := fmt.Sprintf(pdServiceFormat, cl.Name)
		tlsClientSecretName := fmt.Sprintf(tlsClientSecretFormat, cl.Name)
		config, err := util.GetTLSConfig(c.Cli, namespace, cl.Name, tlsClientSecretName)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		if config == nil {
			scheme = httpScheme
		}
		target := cluster.ServiceURL(scheme, namespace, pdServiceName, pdPort)

		proxy := route.GetKubernetesServiceProxy(scheme, namespace, pdServiceName, pdPort, &http.Transport{
			TLSClientConfig: config,
			// Infra provider will call `/proxy/cluster/<cluster id>/pd/pd/api/v1/config` every 20 seconds.
			MaxIdleConns:    1,
			IdleConnTimeout: 30 * time.Second,
		})

		if strings.HasPrefix(r.URL.Path, "dashboard") {
			passwordFromParams := r.URL.Query().Get("password")
			password, err := base64.URLEncoding.DecodeString(passwordFromParams)
			if err != nil {
				http.Error(w, "failed to decode the password", http.StatusBadRequest)
				return
			}
			token, err := getToken(namespace, cl.Name, target, config, string(password))
			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			r.Header.Set("Authorization", "Bearer "+token)
		}
		proxy.ServeHTTP(w, r)
	})
}

func getToken(namespace, name string, target *url.URL, config *tls.Config, password string) (string, error) {
	var token string
	// Cache tokens per clusters as a single namespace can host multiple PD services.
	key := fmt.Sprintf("%s/%s", namespace, name)
	if cachedToken, found := tokenCache.Get(key); found {
		return cachedToken.(string), nil
	}

	// get public key from dashboard login info
	loginInfoURL := target.String() + dashboardLoginInfoPath
	loginInfoRes, err := fetchPublicKey(loginInfoURL, config)
	if err != nil {
		return "", errors.Annotatef(err, "pdaAPI fetch login info")
	}

	loginURL := target.String() + dashboardLoginPath
	username := "cloud_admin"
	if password == "" {
		username = "root"
	}

	pwd := password
	if len(loginInfoRes.SQLAuthPublicKey) > 0 {
		// parse public key
		rsaPub, err := parsePublicKeyBase64(loginInfoRes.SQLAuthPublicKey)
		if err != nil {
			return "", errors.Annotatef(err, "pdaAPI parse public key")
		}
		// encrypt password by public key
		pwd, err = encryptByPublicKey(password, rsaPub)
		if err != nil {
			return "", errors.Annotatef(err, "pdaAPI encrypt password")
		}
	}

	payload := dashboardAuthPayload{
		Type:     AuthTypeSQLUser,
		Username: username,
		Password: pwd,
	}
	res, err := fetchToken(loginURL, payload, config)
	if err != nil {
		return "", errors.Annotatef(err, "pdaAPI fetch token for %s", username)
	}
	token = res.Token
	tokenCache.Set(key, token, cache.DefaultExpiration)
	return token, nil
}

type loginInfoResponse struct {
	SQLAuthPublicKey string `json:"sql_auth_public_key"`
}

func fetchPublicKey(loginInfoURL string, tlsConfig *tls.Config) (*loginInfoResponse, error) {
	const defaultTimeout = 5 * time.Second
	client := &http.Client{
		Timeout:   defaultTimeout,
		Transport: &http.Transport{TLSClientConfig: tlsConfig},
	}

	res, err := client.Get(loginInfoURL) // nolint:bodyclose
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	body, err := io.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}
	loginInfoResponse := &loginInfoResponse{}
	if res.StatusCode == http.StatusOK {
		err = json.Unmarshal(body, loginInfoResponse)
		if err != nil {
			return nil, err
		}
	}
	return loginInfoResponse, nil
}

func parsePublicKeyBase64(base64Key string) (*rsa.PublicKey, error) {
	decodedString, err := base64.StdEncoding.DecodeString(base64Key)
	if err != nil {
		return nil, err
	}
	pub, err := x509.ParsePKIXPublicKey(decodedString)
	if err != nil {
		return nil, err
	}
	rsaPub, ok := pub.(*rsa.PublicKey)
	if !ok {
		return nil, fmt.Errorf("failed to parse RSA public key")
	}
	return rsaPub, nil
}

func encryptByPublicKey(plainText string, publicKey *rsa.PublicKey) (string, error) {
	encryptedText, err := rsa.EncryptPKCS1v15(rand.Reader, publicKey, []byte(plainText))
	if err != nil {
		return "", err
	}
	encodedText := base64.StdEncoding.EncodeToString(encryptedText)
	return encodedText, nil
}

type tokenResponse struct {
	Token  string    `json:"token"`
	Expire time.Time `json:"expire"`
}

func fetchToken(loginURL string, payload dashboardAuthPayload, tlsConfig *tls.Config) (*tokenResponse, error) {
	const defaultTimeout = 5 * time.Second
	client := &http.Client{
		Timeout:   defaultTimeout,
		Transport: &http.Transport{TLSClientConfig: tlsConfig},
	}

	data, err := json.Marshal(payload)
	if err != nil {
		return nil, err
	}

	res, err := client.Post(loginURL, "application/json", bytes.NewBuffer(data)) // nolint:bodyclose
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	body, err := io.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}
	if res.StatusCode == http.StatusOK {
		tokenResponse := &tokenResponse{}
		err = json.Unmarshal(body, tokenResponse)
		if err != nil {
			return nil, err
		}
		return tokenResponse, nil
	}
	return nil, fmt.Errorf("error response %d URL %s, response body: %s", res.StatusCode, loginURL, string(body[:]))
}