package tidbcloud

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	"github.com/Kong/go-pdk"

	"github.com/pingcap-inc/clinic/pkg/db"
	"github.com/pingcap-inc/clinic/pkg/db/model"
)

const (
	sessionIDHeader     = "x-session-id"
	clusterIDHeader     = "x-cluster-id"
	traceIDHeader       = "x-trace-id"
	b3TraceIDHeader     = "x-b3-traceid"
	b3SpanIDHeader      = "x-b3-spanid"
	b3SampledHeader     = "x-b3-sampled"
	requestIDHeader     = "x-request-id"
	kongRequestIDHeader = "x-kong-request-id"
	traceParentHeader   = "traceparent"
	TraceIDContextKey   = "dbaas_auth_trace_id"

	TraceIDResponseHeader = "X-Trace-ID"
)

var getAnyCluster = model.GetAnyCluster

var fallbackTraceIDCounter uint64

type OrgClusterValidation struct {
	Checked         bool
	SessionID       string
	ClusterID       string
	ClusterIDSource string
	ClusterTenantID uint64
}

type RequestTrace struct {
	ID     string
	Source string
}

type OrgClusterValidationLogEvent struct {
	Stage            string
	Status           string
	Reason           string
	Validation       OrgClusterValidation
	DBLookupDuration time.Duration
}

type OrgClusterValidationLogger func(OrgClusterValidationLogEvent)

func extractClusterID(kong *pdk.PDK) (string, string, error) {
	path, err := kong.Request.GetPath()
	if err != nil {
		return "", "", fmt.Errorf("get request path: %w", err)
	}

	if clusterID := clusterIDFromPath(path); clusterID != "" {
		return clusterID, "path", nil
	}

	clusterID, err := kong.Request.GetHeader(clusterIDHeader)
	if err != nil {
		return "", "", fmt.Errorf("get %s header: %w", clusterIDHeader, err)
	}
	clusterID = strings.TrimSpace(clusterID)
	if clusterID == "" {
		return "", "", fmt.Errorf("missing cluster id in path or %s header", clusterIDHeader)
	}

	return clusterID, "header", nil
}

// ValidateOrgCluster validates x-session-id when present. requireSessionID is
// only enabled for route groups whose public contract requires this header.
func ValidateOrgCluster(kong *pdk.PDK, requireSessionID bool, loggers ...OrgClusterValidationLogger) (OrgClusterValidation, error) {
	result := OrgClusterValidation{}
	logger := firstOrgClusterValidationLogger(loggers)

	sessionID, shouldValidate, err := sessionIDFromHeader(kong)
	if err != nil {
		emitOrgClusterValidationLog(logger, "session_id_read", "denied", err.Error(), result, 0)
		return result, err
	}
	if !shouldValidate {
		if requireSessionID {
			err := fmt.Errorf("x-session-id is required")
			emitOrgClusterValidationLog(logger, "session_id_check", "denied", err.Error(), result, 0)
			return result, err
		}
		// This check is opt-in for routes without a required rule: requests
		// without x-session-id continue with JWT authentication and Account RBAC.
		emitOrgClusterValidationLog(logger, "session_id_check", "skipped", "x-session-id is not present", result, 0)
		return result, nil
	}
	result.Checked = true
	result.SessionID = sessionID
	path, err := kong.Request.GetPath()
	if err != nil {
		pathErr := fmt.Errorf("get request path: %w", err)
		emitOrgClusterValidationLog(logger, "org_path_resolve", "denied", pathErr.Error(), result, 0)
		return result, pathErr
	}
	pathOrgID, hasPathOrgID, err := orgIDFromPath(path)
	if err != nil {
		emitOrgClusterValidationLog(logger, "org_path_resolve", "denied", err.Error(), result, 0)
		return result, err
	}
	if hasPathOrgID {
		if expectedSessionID := hashOrgID(strconv.FormatUint(pathOrgID, 10)); sessionID != expectedSessionID {
			mismatchErr := fmt.Errorf("org mismatch: x-session-id=%q path_org_id=%d", sessionID, pathOrgID)
			emitOrgClusterValidationLog(logger, "org_path_check", "denied", mismatchErr.Error(), result, 0)
			return result, mismatchErr
		}
		if clusterIDFromPath(path) == "" {
			emitOrgClusterValidationLog(logger, "org_path_check", "passed", "x-session-id matches path org; cluster tenant lookup is not required", result, 0)
			return result, nil
		}
	}

	// TEMPORARILY DISABLED: this block resolves a cluster ID, queries Clinic's
	// local DB through model.GetAnyCluster, and compares cluster.TenantID with
	// x-session-id. Keep it here for a future re-enable; the request currently
	// relies on the x-session-id/path-org check above plus Account RBAC.
	/*
		clusterID, source, err := extractClusterID(kong)
		if err != nil {
			emitOrgClusterValidationLog(logger, "cluster_id_resolve", "denied", err.Error(), result, 0)
			return result, err
		}
		result.ClusterID = clusterID
		result.ClusterIDSource = source

		emitOrgClusterValidationLog(logger, "cluster_db_lookup", "checking", "query cluster tenant by cluster id", result, 0)
		lookupStartedAt := time.Now()
		cluster, err := getAnyCluster(clusterID)
		lookupDuration := time.Since(lookupStartedAt)
		if err != nil {
			lookupErr := fmt.Errorf("get cluster by id %q from %s: %w", clusterID, source, err)
			emitOrgClusterValidationLog(logger, "cluster_db_lookup", "denied", lookupErr.Error(), result, lookupDuration)
			return result, lookupErr
		}
		if cluster == nil {
			lookupErr := fmt.Errorf("cluster %q from %s not found", clusterID, source)
			emitOrgClusterValidationLog(logger, "cluster_db_lookup", "denied", lookupErr.Error(), result, lookupDuration)
			return result, lookupErr
		}
		result.ClusterTenantID = cluster.TenantID

		if err := validateClusterTenant(cluster, sessionID, clusterID, source); err != nil {
			emitOrgClusterValidationLog(logger, "cluster_db_lookup", "denied", err.Error(), result, lookupDuration)
			return result, err
		}

		emitOrgClusterValidationLog(logger, "cluster_db_lookup", "passed", "cluster tenant matches x-session-id", result, lookupDuration)
	*/
	emitOrgClusterValidationLog(logger, "cluster_db_lookup", "skipped", "cluster tenant DB validation is temporarily disabled", result, 0)
	return result, nil
}

func firstOrgClusterValidationLogger(loggers []OrgClusterValidationLogger) OrgClusterValidationLogger {
	if len(loggers) == 0 {
		return nil
	}

	return loggers[0]
}

func emitOrgClusterValidationLog(logger OrgClusterValidationLogger, stage, status, reason string, validation OrgClusterValidation, lookupDuration time.Duration) {
	if logger == nil {
		return
	}

	logger(OrgClusterValidationLogEvent{
		Stage:            stage,
		Status:           status,
		Reason:           reason,
		Validation:       validation,
		DBLookupDuration: lookupDuration,
	})
}

func RequestTraceID(kong *pdk.PDK) RequestTrace {
	headers, err := kong.Request.GetHeaders(-1)
	if err != nil {
		return RequestTrace{ID: "unavailable", Source: "request-headers-error"}
	}

	return requestTraceFromHeaders(headers)
}

func EnsureRequestTraceID(kong *pdk.PDK) (RequestTrace, error) {
	trace := RequestTraceID(kong)
	if trace.ID == "" || trace.ID == "missing" || trace.ID == "unavailable" {
		trace = RequestTrace{ID: generateTraceID(), Source: "generated"}
	}

	var errs []string
	if err := kong.Ctx.SetShared(TraceIDContextKey, trace.ID); err != nil {
		errs = append(errs, fmt.Sprintf("set ctx trace id: %s", err.Error()))
	}
	if err := kong.ServiceRequest.SetHeader(TraceIDResponseHeader, trace.ID); err != nil {
		errs = append(errs, fmt.Sprintf("set upstream %s: %s", TraceIDResponseHeader, err.Error()))
	}
	spanID := generateSpanID()
	if err := kong.ServiceRequest.SetHeader(b3TraceIDHeader, trace.ID); err != nil {
		errs = append(errs, fmt.Sprintf("set upstream %s: %s", b3TraceIDHeader, err.Error()))
	}
	if err := kong.ServiceRequest.SetHeader(b3SpanIDHeader, spanID); err != nil {
		errs = append(errs, fmt.Sprintf("set upstream %s: %s", b3SpanIDHeader, err.Error()))
	}
	if err := kong.ServiceRequest.SetHeader(b3SampledHeader, "1"); err != nil {
		errs = append(errs, fmt.Sprintf("set upstream %s: %s", b3SampledHeader, err.Error()))
	}
	if traceparent := traceparentFromTraceIDAndSpanID(trace.ID, spanID); traceparent != "" {
		if err := kong.ServiceRequest.SetHeader(traceParentHeader, traceparent); err != nil {
			errs = append(errs, fmt.Sprintf("set upstream %s: %s", traceParentHeader, err.Error()))
		}
	}

	if len(errs) > 0 {
		return trace, errors.New(strings.Join(errs, "; "))
	}
	return trace, nil
}

func requestTraceFromHeaders(headers map[string][]string) RequestTrace {
	// traceparent is Kong's canonical distributed-tracing value. Prefer it
	// over compatibility headers when both are present.
	if traceID := traceIDFromTraceParent(firstHeaderValue(headers, traceParentHeader)); traceID != "" {
		return RequestTrace{ID: traceID, Source: traceParentHeader}
	}

	for _, header := range []string{traceIDHeader, b3TraceIDHeader} {
		if traceID := strings.ToLower(firstHeaderValue(headers, header)); isValidTraceID(traceID) {
			return RequestTrace{ID: traceID, Source: header}
		}
	}

	for _, header := range []string{requestIDHeader, kongRequestIDHeader} {
		if traceID := strings.ToLower(firstHeaderValue(headers, header)); isValidTraceID(traceID) {
			return RequestTrace{ID: traceID, Source: header}
		}
	}

	return RequestTrace{ID: "missing", Source: "none"}
}

func firstHeaderValue(headers map[string][]string, header string) string {
	for key, values := range headers {
		if !strings.EqualFold(key, header) || len(values) == 0 {
			continue
		}

		if value := strings.TrimSpace(values[0]); value != "" {
			return value
		}
	}

	return ""
}

func traceIDFromTraceParent(traceParent string) string {
	parts := strings.Split(strings.TrimSpace(traceParent), "-")
	if len(parts) != 4 || len(parts[1]) != 32 {
		return ""
	}

	traceID := strings.ToLower(parts[1])
	if !isValidTraceID(traceID) {
		return ""
	}
	return traceID
}

func traceparentFromTraceID(traceID string) string {
	return traceparentFromTraceIDAndSpanID(traceID, generateSpanID())
}

func traceparentFromTraceIDAndSpanID(traceID, spanID string) string {
	if !isValidTraceID(traceID) {
		return ""
	}

	return fmt.Sprintf("00-%s-%s-01", strings.ToLower(traceID), spanID)
}

func isValidTraceID(traceID string) bool {
	if len(traceID) != 32 || strings.Trim(traceID, "0") == "" {
		return false
	}

	_, err := hex.DecodeString(traceID)
	return err == nil
}

func generateTraceID() string {
	var traceID [16]byte
	if _, err := rand.Read(traceID[:]); err == nil && !allZero(traceID[:]) {
		return hex.EncodeToString(traceID[:])
	}

	binary.BigEndian.PutUint64(traceID[:8], uint64(time.Now().UnixNano()))
	binary.BigEndian.PutUint64(traceID[8:], atomic.AddUint64(&fallbackTraceIDCounter, 1))
	return hex.EncodeToString(traceID[:])
}

func generateSpanID() string {
	var spanID [8]byte
	if _, err := rand.Read(spanID[:]); err == nil && !allZero(spanID[:]) {
		return hex.EncodeToString(spanID[:])
	}

	binary.BigEndian.PutUint64(spanID[:], uint64(time.Now().UnixNano())^atomic.AddUint64(&fallbackTraceIDCounter, 1))
	return hex.EncodeToString(spanID[:])
}

func allZero(data []byte) bool {
	for _, b := range data {
		if b != 0 {
			return false
		}
	}

	return true
}

func sessionIDFromHeader(kong *pdk.PDK) (string, bool, error) {
	sessionID, err := kong.Request.GetHeader(sessionIDHeader)
	if err != nil {
		return "", false, fmt.Errorf("get %s header: %w", sessionIDHeader, err)
	}
	sessionID = strings.TrimSpace(sessionID)
	if sessionID == "" {
		return "", false, nil
	}

	if !isValidOrgIDHash(sessionID) {
		return "", true, fmt.Errorf("invalid %s header %q: must be a 64-character lowercase SHA-256 hex string", sessionIDHeader, sessionID)
	}

	return sessionID, true, nil
}

func hashOrgID(orgID string) string {
	digest := sha256.Sum256([]byte(strings.TrimSpace(orgID)))
	return hex.EncodeToString(digest[:])
}

func isValidOrgIDHash(value string) bool {
	if len(value) != sha256.Size*2 || value != strings.ToLower(value) {
		return false
	}

	_, err := hex.DecodeString(value)
	return err == nil
}

func clusterIDFromPath(path string) string {
	parts := strings.Split(path, "/")
	for i := 0; i < len(parts)-1; i++ {
		if parts[i] != "clusters" {
			continue
		}

		clusterID := strings.TrimSpace(parts[i+1])
		if clusterID != "" {
			return clusterID
		}
	}

	return ""
}

func orgIDFromPath(path string) (uint64, bool, error) {
	parts := strings.Split(path, "/")
	for i := 0; i < len(parts)-1; i++ {
		if parts[i] != "orgs" {
			continue
		}

		orgIDText := strings.TrimSpace(parts[i+1])
		if orgIDText == "" {
			return 0, true, fmt.Errorf("missing org id in request path")
		}
		orgID, err := strconv.ParseUint(orgIDText, 10, 64)
		if err != nil || orgID == 0 {
			return 0, true, fmt.Errorf("invalid org id in request path %q", orgIDText)
		}
		return orgID, true, nil
	}

	return 0, false, nil
}

func validateClusterTenant(cluster *db.Cluster, sessionID, clusterID, source string) error {
	if hashOrgID(strconv.FormatUint(cluster.TenantID, 10)) != sessionID {
		return fmt.Errorf("cluster org mismatch: cluster_id=%q source=%s x-session-id=%q cluster_tenant_id=%d", clusterID, source, sessionID, cluster.TenantID)
	}

	return nil
}