package main

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/Kong/go-pdk"
	"github.com/Kong/go-pdk/server"
	"github.com/imdario/mergo"
	"github.com/pingcap/log"
	sdkauth "github.com/tidbcloud/tiinf/sdk/auth"
	"go.uber.org/atomic"
	"go.uber.org/zap"

	"github.com/pingcap-inc/clinic/pkg/account_sdk"
	"github.com/pingcap-inc/clinic/pkg/api_gateway/auth/tidbcloud"
	"github.com/pingcap-inc/clinic/pkg/api_gateway/rbac"
	"github.com/pingcap-inc/clinic/pkg/config"
	"github.com/pingcap-inc/clinic/pkg/db"
)

const UserHeader = "X-CLOUD-USER"

type Config struct {
	SecretConfigPath          *string                   `json:"secret-config-path"`
	DBConfig                  *DBConfig                 `json:"db"`
	TiDBCloudAccount          TiDBCloudAccountConfig    `json:"tidb-cloud-account"`
	RBAC                      RBACConfig                `json:"rbac"`
	SessionIDValidationConfig SessionIDValidationConfig `json:"session-id-validation"`
}

type DBConfig struct {
	User     string `json:"user"`
	Host     string `json:"host"`
	Database string `json:"database"`
	Password string `json:"password"`
	Port     int    `json:"port"`
}

type TiDBCloudAccountConfig struct {
	Endpoint string `json:"endpoint"`
}

type RBACConfig struct {
	ServerEndpoint string            `json:"server-endpoint"`
	Whitelist      []string          `json:"whitelist"`
	ResourceMap    map[string]string `json:"resources"`
	// Mode           string            `json:"mode"`
}

type SessionIDValidationConfig struct {
	DeployTypeCacheTTL  string                    `json:"deploy-type-cache-ttl"`
	DeployTypeCacheSize int                       `json:"deploy-type-cache-size"`
	AllowedDeployTypes  []string                  `json:"allowed-deploy-types"`
	Rules               []SessionIDValidationRule `json:"rules"`
}

type SessionIDValidationRule struct {
	Method           string `json:"method"`
	Path             string `json:"path"`
	Required         bool   `json:"required"`
	ClusterIDPattern string `json:"cluster-id-pattern"`
}

type Response struct {
	Code    int      `json:"code"`
	Message string   `json:"message"`
	Details []string `json:""`
	//{
	//	"code": 16,
	//	"message": "unauthorized",
	//	"details": []
	//}
}

func New() any {
	return &Config{}
}

var (
	Version  = "0.2"
	Priority = 1
)

var (
	// TODO: add comments
	// sess *session2.SessionManager
	// allowIPNet   []*net.IPNet
	inited atomic.Bool
	once   sync.Once
	// externalAuth *account_sdk.Client
	externalAuth *account_sdk.Client

	// TODO: add comments
	// auth *authorization.Authorization

	rbacChecker rbac.PermissionChecker

	deployTypeValidator *clusterDeployTypeValidator
)

func main() {
	// config
	err := server.StartServer(New, Version, Priority)
	if err != nil {
		log.Error("start server failed", zap.Error(err))
	}
}

func (conf *Config) MergeSecret(kong *pdk.PDK) error {
	if conf.SecretConfigPath == nil || *conf.SecretConfigPath == "" {
		return nil
	}

	data, err := os.ReadFile(*conf.SecretConfigPath)
	if err != nil {
		_ = kong.Log.Err("[dbaas_auth] failed to read secret config: ", err.Error(), ", path: ", *conf.SecretConfigPath)
		return err
	}

	secretConfig := Config{}
	if err := json.Unmarshal(data, &secretConfig); err != nil {
		_ = kong.Log.Err("[dbaas_auth] failed to unmarshal secret config: ", err.Error(), ", path: ", *conf.SecretConfigPath)
		return err
	}

	if err := mergo.Merge(conf, &secretConfig, mergo.WithOverride); err != nil {
		_ = kong.Log.Err("[dbaas_auth] failed to merge secret config: ", err.Error())
		return err
	}

	return nil
}

func (conf Config) Access(kong *pdk.PDK) {
	once.Do(func() {
		init := true
		if err := conf.MergeSecret(kong); err != nil {
			_ = kong.Log.Err("[dbaas_auth] merge secret config failed, err: ", err.Error())
			init = false
		}

		if conf.DBConfig == nil {
			_ = kong.Log.Err("[dbaas_auth] db config is required")
			init = false
		} else {
			err := db.Initialize(&config.DBConfig{
				User:     conf.DBConfig.User,
				Host:     conf.DBConfig.Host,
				Password: conf.DBConfig.Password,
				Port:     uint16(conf.DBConfig.Port),
				Database: conf.DBConfig.Database,
			}, true)
			if err != nil {
				_ = kong.Log.Err("[dbaas_auth] init db failed, err: ", err.Error())
				init = false
			}
		}

		validator, err := newClusterDeployTypeValidator(conf.SessionIDValidationConfig)
		if err != nil {
			_ = kong.Log.Err("[dbaas_auth] init deploy type validator failed, err: ", err.Error())
			init = false
		} else {
			deployTypeValidator = validator
		}

		cli, err := account_sdk.New(conf.TiDBCloudAccount.Endpoint, account_sdk.TimeoutOption(5*time.Second))
		if err != nil {
			_ = kong.Log.Err("[dbaas_auth] init external auth client failed, err: ", err.Error())
			init = false
		} else {
			externalAuth = cli
		}

		sdkauth.InitAccountClientByEndpoint(conf.RBAC.ServerEndpoint, 5*time.Second)
		rbacchecker, err := rbac.NewDbaasRBACChecker(conf.RBAC.ResourceMap, conf.RBAC.Whitelist)
		if err != nil {
			_ = kong.Log.Err("[dbaas_auth] init rbac server failed, err: ", err.Error())
			init = false
		} else {
			rbacChecker = rbacchecker
		}

		inited.Store(init)
		if init {
			_ = kong.Log.Info("[dbaas_auth] init done")
		}
	})

	responseHeaders := make(map[string][]string)
	responseHeaders["Content-Type"] = append(responseHeaders["Content-Type"], "application/json")
	if !inited.Load() {
		_ = kong.Log.Warn("[dbaas_auth] wait plugin inited or init failed")
		kong.Response.Exit(http.StatusInternalServerError, ExitResponse(16, "internal server error"), responseHeaders)
		return
	}

	user, authErr := tidbcloud.MWTryAuth(kong, externalAuth)
	if authErr != nil {
		_ = kong.Log.Err("[dbaas_auth] auth failed: " + authErr.Message())
		kong.Response.Exit(http.StatusUnauthorized, ExitResponse(16, "unauthorized"), responseHeaders)
		return
	}

	trace, traceErr := tidbcloud.EnsureRequestTraceID(kong)
	if traceErr != nil {
		_ = kong.Log.Warn("[dbaas_auth] ensure request trace failed: ", traceErr.Error())
	}
	addTraceResponseHeaders(responseHeaders, trace)

	rule, shouldValidateSessionID, validationPath, err := conf.sessionIDValidationRule(kong)
	if err != nil {
		_ = kong.Log.Err("[dbaas_auth] resolve x-session-id validation scope failed: ", err.Error())
		kong.Response.Exit(http.StatusInternalServerError, ExitResponse(16, "internal server error"), responseHeaders)
		return
	}
	if shouldValidateSessionID {
		_, err = tidbcloud.ValidateOrgCluster(kong, rule.Required, func(event tidbcloud.OrgClusterValidationLogEvent) {
			logOrgClusterValidationEvent(kong, trace, event)
		})
		if err != nil {
			kong.Response.Exit(http.StatusForbidden, ExitResponse(16, "forbidden"), responseHeaders)
			return
		}

		if rule.ClusterIDPattern != "" && len(conf.SessionIDValidationConfig.AllowedDeployTypes) > 0 {
			validation, validateErr := deployTypeValidator.Validate(rule, conf.SessionIDValidationConfig.AllowedDeployTypes, validationPath)
			logClusterDeployTypeValidationEvent(kong, trace, validation)
			if validateErr != nil {
				kong.Response.Exit(http.StatusForbidden, ExitResponse(16, "forbidden"), responseHeaders)
				return
			}
		}
	} else {
		_ = kong.Log.Info(fmt.Sprintf(
			"[dbaas_auth] org cluster validation trace_id=%q trace_source=%s stage=validation_scope status=skipped checked=false reason=%q",
			trace.ID,
			trace.Source,
			fmt.Sprintf("path %q is not configured for x-session-id validation", validationPath),
		))
	}

	data, err := json.Marshal(user)
	if err != nil {
		_ = kong.Log.Err("[dbaas_auth] marshal failed: ", err.Error())
		// kong.Response.Exit(http.StatusUnauthorized, ExitResponse(16, "unauthorized"), responseHeaders)
		kong.Response.Exit(http.StatusInternalServerError, ExitResponse(16, "internal server error"), responseHeaders)
		return
	}

	ok, err := rbacChecker.Handle(kong, user)
	if err != nil || !ok {
		if err != nil {
			_ = kong.Log.Err("[dbaas_auth] rbac check failed: ", err.Error())
		} else {
			_ = kong.Log.Warn("[dbaas_auth] rbac permission denied")
		}
		kong.Response.Exit(http.StatusForbidden, ExitResponse(16, "forbidden"), responseHeaders)
		return
	}

	err = kong.ServiceRequest.SetHeader(UserHeader, base64.StdEncoding.EncodeToString(data))
	if err != nil {
		_ = kong.Log.Err("[dbaas_auth] set header failed: ", err.Error())
		kong.Response.Exit(http.StatusInternalServerError, ExitResponse(16, "internal server error"), responseHeaders)
		return
	}
}

// shouldValidateSessionID limits x-session-id checks to externally exposed
// routes explicitly listed in the Kong plugin config. A matched rule may
// require the header, while other matched routes keep the optional behavior.
// Rules are anchored regular expressions over the request path; an empty
// method matches every HTTP method.
func (conf Config) sessionIDValidationRule(kong *pdk.PDK) (SessionIDValidationRule, bool, string, error) {
	path, err := kong.Request.GetPath()
	if err != nil {
		return SessionIDValidationRule{}, false, "", fmt.Errorf("get request path: %w", err)
	}
	method, err := kong.Request.GetMethod()
	if err != nil {
		return SessionIDValidationRule{}, false, path, fmt.Errorf("get request method: %w", err)
	}

	rule, shouldValidate, err := sessionIDValidationRuleForRequest(method, path, conf.SessionIDValidationConfig.Rules)
	if err != nil {
		return SessionIDValidationRule{}, false, path, err
	}
	return rule, shouldValidate, path, nil
}

func matchesSessionIDValidationRules(method, path string, rules []SessionIDValidationRule) (bool, error) {
	_, matched, err := sessionIDValidationRuleForRequest(method, path, rules)
	return matched, err
}

func sessionIDValidationRuleForRequest(method, path string, rules []SessionIDValidationRule) (SessionIDValidationRule, bool, error) {
	for _, rule := range rules {
		if rule.Method != "" && !strings.EqualFold(rule.Method, method) {
			continue
		}
		pattern, err := regexp.Compile(rule.Path)
		if err != nil {
			return SessionIDValidationRule{}, false, fmt.Errorf("compile session-id-validation path %q: %w", rule.Path, err)
		}
		if pattern.MatchString(path) {
			return rule, true, nil
		}
	}

	return SessionIDValidationRule{}, false, nil
}

func (conf Config) Response(kong *pdk.PDK) {
	traceID, err := kong.Ctx.GetSharedString(tidbcloud.TraceIDContextKey)
	if err != nil || traceID == "" {
		return
	}

	if err := kong.Response.SetHeader(tidbcloud.TraceIDResponseHeader, traceID); err != nil {
		_ = kong.Log.Warn("[dbaas_auth] set response trace header failed: ", err.Error())
	}
}

func ExitResponse(code int, message string) []byte {
	resp := Response{
		Code:    code,
		Message: message,
	}

	data, _ := json.Marshal(resp)
	return data
}

func logOrgClusterValidationEvent(kong *pdk.PDK, trace tidbcloud.RequestTrace, event tidbcloud.OrgClusterValidationLogEvent) {
	entry := orgClusterValidationLog(trace, event)
	if event.Status == "denied" {
		_ = kong.Log.Warn(entry)
		return
	}

	_ = kong.Log.Info(entry)
}

func orgClusterValidationLog(trace tidbcloud.RequestTrace, event tidbcloud.OrgClusterValidationLogEvent) string {
	validation := event.Validation
	return fmt.Sprintf(
		"[dbaas_auth] org cluster validation trace_id=%q trace_source=%s stage=%s status=%s checked=%t session_id=%q cluster_id=%q cluster_id_source=%s cluster_tenant_id=%d db_lookup_ms=%d reason=%q",
		trace.ID,
		trace.Source,
		event.Stage,
		event.Status,
		validation.Checked,
		validation.SessionID,
		validation.ClusterID,
		validation.ClusterIDSource,
		validation.ClusterTenantID,
		event.DBLookupDuration.Milliseconds(),
		event.Reason,
	)
}

func logClusterDeployTypeValidationEvent(kong *pdk.PDK, trace tidbcloud.RequestTrace, validation ClusterDeployTypeValidation) {
	entry := fmt.Sprintf(
		"[dbaas_auth] cluster deploy type validation trace_id=%q trace_source=%s status=%s cluster_id=%q deploy_type=%q allowed_deploy_types=%q lookup_source=%s cache_hit=%t db_lookup_ms=%d reason=%q",
		trace.ID,
		trace.Source,
		validation.Status,
		validation.ClusterID,
		validation.DeployType,
		strings.Join(validation.AllowedDeployTypes, ","),
		validation.LookupSource,
		validation.CacheHit,
		validation.DBLookupDuration.Milliseconds(),
		validation.Reason,
	)
	if validation.Status == "denied" {
		_ = kong.Log.Warn(entry)
		return
	}

	_ = kong.Log.Info(entry)
}

func addTraceResponseHeaders(headers map[string][]string, trace tidbcloud.RequestTrace) {
	if trace.ID == "" || trace.ID == "missing" || trace.ID == "unavailable" {
		return
	}

	headers[tidbcloud.TraceIDResponseHeader] = []string{trace.ID}
}