package innerauth

import (
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/mem9-ai/drive9/pkg/metrics"
	"github.com/mem9-ai/drive9/pkg/tenant"
)

var (
	ErrMalformedAuthorization        = errors.New("malformed TiDB Cloud authorization header")
	ErrInvalidJWTOrganizationID      = errors.New("JWT requires a numeric non-zero orgId query parameter")
	ErrAmbiguousCredentials          = errors.New("JWT and API key credentials cannot be used together")
	ErrDefaultCredentialsUnavailable = errors.New("server default TiDB Cloud credentials are unavailable")
	ErrInnerAuthDisabled             = errors.New("inner authentication is disabled")
)

const TiDBCloudAuthorizationHeader = "X-TiDB-Cloud-Authorization"

type CredentialKind string

const (
	CredentialJWT    CredentialKind = "jwt"
	CredentialAPIKey CredentialKind = "api_key"
)

type ControlPlaneSubject struct {
	OrganizationID uint64
	UserID         uint64
	PublicKey      string
	Kind           CredentialKind
	ResourceInfos  []APIKeyResourceInfo
	// ExecutionCredentials is the server-side credential used for downstream
	// TiDB Cloud OpenAPI calls; it is separate from the caller's authz key.
	ExecutionCredentials *tenant.CredentialProvisionRequest
}

type AccountProvider interface {
	UserAuthn(ctx context.Context, token string) (UserAuthnResult, error)
	UserAuthz(ctx context.Context, userID uint64, requestURI, resourceURI string) error
	APIKeyAuthn(ctx context.Context, encodedAPIKey string) (APIKeyAuthnResult, error)
	APIKeyAuthz(ctx context.Context, publicKey, requestURI, resourceURI string) error
}

type DefaultCredentialProvider = tenant.DefaultCredentialProvider

type Dispatcher struct {
	account     AccountProvider
	defaults    DefaultCredentialProvider
	apiKeyCache *apiKeyCache
}

func NewDispatcher(account AccountProvider, defaults DefaultCredentialProvider, cacheTTL time.Duration) *Dispatcher {
	return &Dispatcher{account: account, defaults: defaults, apiKeyCache: newAPIKeyCache(cacheTTL)}
}

func (d *Dispatcher) Enabled() bool {
	return d != nil && d.account != nil
}

func (d *Dispatcher) Authenticate(ctx context.Context, authorization string, organizationID uint64, credentials tenant.CredentialProvisionRequest) (ControlPlaneSubject, error) {
	if !d.Enabled() {
		return ControlPlaneSubject{}, ErrInnerAuthDisabled
	}
	kind, err := ClassifyCredentials(authorization, credentials)
	if err != nil {
		return ControlPlaneSubject{}, err
	}
	switch kind {
	case CredentialJWT:
		token, err := BearerJWT(authorization)
		if err != nil {
			return ControlPlaneSubject{}, err
		}
		if organizationID == 0 {
			return ControlPlaneSubject{}, ErrInvalidJWTOrganizationID
		}
		identity, err := d.account.UserAuthn(ctx, token)
		if err != nil {
			return ControlPlaneSubject{}, err
		}
		executionCredentials, err := d.defaultExecutionCredentials()
		if err != nil {
			return ControlPlaneSubject{}, err
		}
		return ControlPlaneSubject{
			OrganizationID:       organizationID,
			UserID:               identity.UserID,
			Kind:                 CredentialJWT,
			ExecutionCredentials: &executionCredentials,
		}, nil
	case CredentialAPIKey:
		encoded, err := EncodeAPIKey(credentials)
		if err != nil {
			return ControlPlaneSubject{}, err
		}
		identity, cached := d.apiKeyCache.getAuthn(encoded)
		metrics.RecordTiDBCloudInnerAuthCacheRequest("api_key_authn", cacheResult(cached))
		if !cached {
			identity, err = d.account.APIKeyAuthn(ctx, encoded)
			if err != nil {
				return ControlPlaneSubject{}, err
			}
		}
		if strings.TrimSpace(identity.PublicKey) != strings.TrimSpace(credentials.PublicKey) {
			return ControlPlaneSubject{}, ErrInvalidResponse
		}
		executionCredentials, err := d.defaultExecutionCredentials()
		if err != nil {
			return ControlPlaneSubject{}, err
		}
		if !cached {
			d.apiKeyCache.rememberAuthn(encoded, identity)
		}
		return ControlPlaneSubject{
			OrganizationID:       identity.OrganizationID,
			PublicKey:            identity.PublicKey,
			Kind:                 CredentialAPIKey,
			ResourceInfos:        identity.ResourceInfos,
			ExecutionCredentials: &executionCredentials,
		}, nil
	default:
		return ControlPlaneSubject{}, tenant.ErrCredentialsRequired
	}
}

func (d *Dispatcher) defaultExecutionCredentials() (tenant.CredentialProvisionRequest, error) {
	if d.defaults == nil {
		return tenant.CredentialProvisionRequest{}, ErrDefaultCredentialsUnavailable
	}
	credentials, ok := d.defaults.DefaultCredentials()
	if !ok || strings.TrimSpace(credentials.PublicKey) == "" || strings.TrimSpace(credentials.PrivateKey) == "" {
		return tenant.CredentialProvisionRequest{}, ErrDefaultCredentialsUnavailable
	}
	return credentials, nil
}

func SelectHighestProjectScope(resourceInfos []APIKeyResourceInfo) (APIKeyResourceInfo, bool) {
	priority := map[string]int{
		"project:readonly": 1,
		"project:dev":      2,
		"project:owner":    3,
	}
	var selected APIKeyResourceInfo
	selectedPriority := 0
	for _, resourceInfo := range resourceInfos {
		if !strings.EqualFold(strings.TrimSpace(resourceInfo.ScopeType), "PROJECT") {
			continue
		}
		if resourceInfo.ResourceID == 0 {
			continue
		}
		currentPriority := priority[strings.ToLower(strings.TrimSpace(resourceInfo.Role))]
		if currentPriority > selectedPriority {
			selected = resourceInfo
			selectedPriority = currentPriority
		}
	}
	return selected, selectedPriority > 0
}

func (d *Dispatcher) Authorize(ctx context.Context, subject ControlPlaneSubject, requestURI, resourceURI string) error {
	if !d.Enabled() {
		return ErrInnerAuthDisabled
	}
	switch subject.Kind {
	case CredentialJWT:
		return d.account.UserAuthz(ctx, subject.UserID, requestURI, resourceURI)
	case CredentialAPIKey:
		cached := d.apiKeyCache.hasAuthz(subject.PublicKey, requestURI, resourceURI)
		metrics.RecordTiDBCloudInnerAuthCacheRequest("api_key_authz", cacheResult(cached))
		if cached {
			return nil
		}
		if err := d.account.APIKeyAuthz(ctx, subject.PublicKey, requestURI, resourceURI); err != nil {
			return err
		}
		d.apiKeyCache.rememberAuthz(subject.PublicKey, requestURI, resourceURI)
		return nil
	default:
		return ErrUnauthenticated
	}
}

func cacheResult(hit bool) string {
	if hit {
		return "hit"
	}
	return "miss"
}

func ClassifyCredentials(authorization string, credentials tenant.CredentialProvisionRequest) (CredentialKind, error) {
	hasJWT := strings.TrimSpace(authorization) != ""
	hasPublicKey := credentials.PublicKey != ""
	hasPrivateKey := credentials.PrivateKey != ""
	if hasJWT && (hasPublicKey || hasPrivateKey) {
		return "", ErrAmbiguousCredentials
	}
	if hasPublicKey != hasPrivateKey {
		return "", tenant.ErrPartialCredentials
	}
	if hasJWT {
		return CredentialJWT, nil
	}
	if hasPublicKey {
		return CredentialAPIKey, nil
	}
	return "", tenant.ErrCredentialsRequired
}

func BearerJWT(authorization string) (string, error) {
	parts := strings.Fields(authorization)
	if len(parts) != 2 || !strings.EqualFold(parts[0], "Bearer") || !validJWTShape(parts[1]) {
		return "", ErrMalformedAuthorization
	}
	return parts[1], nil
}

func JWTFromHeader(header http.Header) (string, error) {
	return BearerJWT(header.Get(TiDBCloudAuthorizationHeader))
}

func ParseOrganizationID(raw string) (uint64, error) {
	organizationID, err := strconv.ParseUint(strings.TrimSpace(raw), 10, 64)
	if err != nil || organizationID == 0 {
		return 0, ErrInvalidJWTOrganizationID
	}
	return organizationID, nil
}

func EncodeAPIKey(credentials tenant.CredentialProvisionRequest) (string, error) {
	if credentials.PublicKey == "" && credentials.PrivateKey == "" {
		return "", tenant.ErrCredentialsRequired
	}
	if credentials.PublicKey == "" || credentials.PrivateKey == "" {
		return "", tenant.ErrPartialCredentials
	}
	return base64.URLEncoding.EncodeToString([]byte(credentials.PublicKey + ":" + credentials.PrivateKey)), nil
}

func validJWTShape(token string) bool {
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		return false
	}
	for _, part := range parts {
		if part == "" {
			return false
		}
		if _, err := decodeJWTPart(part); err != nil {
			return false
		}
	}
	return true
}

func decodeJWTPart(part string) ([]byte, error) {
	raw, err := base64.RawURLEncoding.DecodeString(part)
	if err == nil {
		return raw, nil
	}
	raw, paddedErr := base64.URLEncoding.DecodeString(part)
	if paddedErr != nil {
		return nil, fmt.Errorf("decode JWT part: %w", paddedErr)
	}
	return raw, nil
}