package server

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"strings"

	"github.com/mem9-ai/drive9/pkg/innerauth"
	"github.com/mem9-ai/drive9/pkg/logger"
	"github.com/mem9-ai/drive9/pkg/tenant"
	"go.uber.org/zap"
)

const innerControlPlaneOrganizationPlaceholder = "{org_id}"

type innerAuthContextKey int

const innerAuthSubjectContextKey innerAuthContextKey = iota

var errInnerAuthMalformedRequest = errors.New("malformed inner authentication request")

func InnerAuthSubjectFromContext(ctx context.Context) *innerauth.ControlPlaneSubject {
	subject, _ := ctx.Value(innerAuthSubjectContextKey).(*innerauth.ControlPlaneSubject)
	return subject
}

func InnerAuthCredentialsFromContext(ctx context.Context) *tenant.CredentialProvisionRequest {
	subject := InnerAuthSubjectFromContext(ctx)
	if subject == nil {
		return nil
	}
	return subject.ExecutionCredentials
}

func controlPlaneCredentials(r *http.Request, publicKey, privateKey string) (tenant.CredentialProvisionRequest, error) {
	if credentials := InnerAuthCredentialsFromContext(r.Context()); credentials != nil {
		return *credentials, nil
	}
	return adminCredentials(publicKey, privateKey, r)
}

func controlPlaneOrganizationID(ctx context.Context) (string, bool) {
	subject := InnerAuthSubjectFromContext(ctx)
	if subject == nil || subject.OrganizationID == 0 {
		return "", false
	}
	return strconv.FormatUint(subject.OrganizationID, 10), true
}

func withInnerAuthSubject(ctx context.Context, subject *innerauth.ControlPlaneSubject) context.Context {
	return context.WithValue(ctx, innerAuthSubjectContextKey, subject)
}

func innerAuthURIs(subject innerauth.ControlPlaneSubject, orgResourceURI string) (requestURI, resourceURI string) {
	resourceURI = orgResourceURI
	if subject.Kind == innerauth.CredentialAPIKey {
		if projectScope, ok := innerauth.SelectHighestProjectScope(subject.ResourceInfos); ok {
			orgPrefix := "/drive9/orgs/" + innerControlPlaneOrganizationPlaceholder + "/"
			if strings.HasPrefix(resourceURI, orgPrefix) {
				resourceURI = strings.Replace(resourceURI, orgPrefix, orgPrefix+"projects/{project_id}/", 1)
			}
			requestURI = strings.ReplaceAll(resourceURI, innerControlPlaneOrganizationPlaceholder, strconv.FormatUint(subject.OrganizationID, 10))
			requestURI = strings.ReplaceAll(requestURI, "{project_id}", strconv.FormatUint(projectScope.ResourceID, 10))
			return requestURI, resourceURI
		}
	}
	requestURI = strings.ReplaceAll(resourceURI, innerControlPlaneOrganizationPlaceholder, strconv.FormatUint(subject.OrganizationID, 10))
	return requestURI, resourceURI
}

func matchInnerControlPlaneRoute(method, path string) (string, bool) {
	const (
		adminTenantsModify    = "/drive9/orgs/{org_id}/admin/tenants/modify"
		adminTenantsView      = "/drive9/orgs/{org_id}/admin/tenants/view"
		tenantPoolModify      = "/drive9/orgs/{org_id}/admin/tenant-pool/modify"
		tenantPoolView        = "/drive9/orgs/{org_id}/admin/tenant-pool/view"
		quotaModify           = "/drive9/orgs/{org_id}/quota/modify"
		quotaView             = "/drive9/orgs/{org_id}/quota/view"
		extractConfigView     = "/drive9/orgs/{org_id}/extract-config/view"
		extractConfigModify   = "/drive9/orgs/{org_id}/extract-config/modify"
		objectBackendView     = "/drive9/orgs/{org_id}/object-backend/view"
		objectBackendModify   = "/drive9/orgs/{org_id}/object-backend/modify"
		objectNamespaceView   = "/drive9/orgs/{org_id}/object-namespace/view"
		objectNamespaceModify = "/drive9/orgs/{org_id}/object-namespace/modify"
		tokensModify          = "/drive9/orgs/{org_id}/tokens/modify"
		tokensView            = "/drive9/orgs/{org_id}/tokens/view"
	)

	switch {
	case method == http.MethodPost && path == "/v1/provision":
		return adminTenantsModify, true
	case path == "/v1/admin/tenants":
		switch method {
		case http.MethodPost:
			return adminTenantsModify, true
		case http.MethodGet:
			return adminTenantsView, true
		}
	case strings.HasPrefix(path, "/v1/admin/tenants/"):
		parts := splitAdminTenantPath(path)
		switch {
		case len(parts) == 1 && method == http.MethodDelete:
			return adminTenantsModify, true
		case len(parts) == 1 && method == http.MethodGet:
			return adminTenantsView, true
		case len(parts) == 2 && parts[1] == "embedding-config" && method == http.MethodGet:
			return adminTenantsView, true
		case len(parts) == 2 && parts[1] == "embedding-config" && method == http.MethodPut:
			return adminTenantsModify, true
		case len(parts) == 2 && parts[1] == "quota" && method == http.MethodPost:
			return quotaModify, true
		case len(parts) == 3 && parts[1] == "extract-config" && method == http.MethodGet:
			return extractConfigView, true
		case len(parts) == 3 && parts[1] == "extract-config" && method == http.MethodPut:
			return extractConfigModify, true
		case len(parts) == 2 && parts[1] == "object-namespace" && method == http.MethodGet:
			return objectNamespaceView, true
		case len(parts) == 2 && parts[1] == "object-namespace" && (method == http.MethodPut || method == http.MethodDelete):
			return objectNamespaceModify, true
		}
	case path == "/v1/admin/object-backends":
		switch method {
		case http.MethodGet:
			return objectBackendView, true
		case http.MethodPost:
			return objectBackendModify, true
		}
	case strings.HasPrefix(path, "/v1/admin/object-backends/"):
		switch method {
		case http.MethodGet:
			return objectBackendView, true
		case http.MethodPatch, http.MethodPut, http.MethodDelete:
			return objectBackendModify, true
		}
	case path == "/v1/admin/tenant-pool":
		switch method {
		case http.MethodPost, http.MethodPatch, http.MethodDelete:
			return tenantPoolModify, true
		case http.MethodGet:
			return tenantPoolView, true
		}
	case path == "/v1/quota":
		switch method {
		case http.MethodGet:
			return quotaView, true
		case http.MethodPost:
			return quotaModify, true
		}
	case path == "/v1/tokens" && method == http.MethodGet:
		return tokensView, true
	case path == "/v1/tokens/generate" && method == http.MethodPost:
		return tokensModify, true
	case path == "/v1/tokens/revoke" && method == http.MethodPost:
		return tokensModify, true
	case strings.HasPrefix(path, "/v1/tokens/"):
		switch {
		case method == http.MethodDelete && tokenIDFromPath(path) != "":
			return tokensModify, true
		case method == http.MethodPost && tokenIDFromStatusPath(path, "/activate") != "":
			return tokensModify, true
		case method == http.MethodPost && tokenIDFromStatusPath(path, "/deactivate") != "":
			return tokensModify, true
		}
	}
	return "", false
}

func splitAdminTenantPath(path string) []string {
	const prefix = "/v1/admin/tenants/"
	if !strings.HasPrefix(path, prefix) {
		return nil
	}
	rest := strings.Trim(strings.TrimPrefix(path, prefix), "/")
	if rest == "" {
		return nil
	}
	parts := strings.Split(rest, "/")
	for _, part := range parts {
		if strings.TrimSpace(part) == "" {
			return nil
		}
	}
	return parts
}

func innerControlPlaneAuthMiddleware(dispatcher *innerauth.Dispatcher, next http.Handler) http.Handler {
	if !dispatcher.Enabled() {
		return next
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		resourceURI, matched := matchInnerControlPlaneRoute(r.Method, r.URL.Path)
		if !matched {
			next.ServeHTTP(w, r)
			return
		}

		authorization := strings.TrimSpace(r.Header.Get(innerauth.TiDBCloudAuthorizationHeader))
		hasDrive9Bearer := bearerToken(r) != ""
		if hasDrive9Bearer && supportsDrive9Bearer(r.Method, r.URL.Path) {
			if authorization != "" || strings.TrimSpace(r.Header.Get(quotaPublicKeyHeader)) != "" || strings.TrimSpace(r.Header.Get(quotaPrivateKeyHeader)) != "" {
				writeInnerAuthError(w, innerauth.ErrAmbiguousCredentials)
				return
			}
			next.ServeHTTP(w, r)
			return
		}
		credentials, err := readInnerControlPlaneCredentials(r)
		if err != nil {
			writeInnerAuthError(w, err)
			return
		}
		if hasDrive9Bearer {
			if authorization != "" || credentials.PublicKey != "" || credentials.PrivateKey != "" {
				writeInnerAuthError(w, innerauth.ErrAmbiguousCredentials)
				return
			}
			writeInnerAuthError(w, innerauth.ErrUnauthenticated)
			return
		}
		organizationID := uint64(0)
		kind, err := innerauth.ClassifyCredentials(authorization, credentials)
		if err != nil {
			writeInnerAuthError(w, err)
			return
		}
		if kind == innerauth.CredentialJWT {
			organizationID, err = innerauth.ParseOrganizationID(r.URL.Query().Get("orgId"))
			if err != nil {
				writeInnerAuthError(w, err)
				return
			}
		}

		subject, err := dispatcher.Authenticate(r.Context(), authorization, organizationID, credentials)
		if err != nil {
			writeInnerAuthError(w, err)
			return
		}
		requestURI, resourceURI := innerAuthURIs(subject, resourceURI)
		if err := dispatcher.Authorize(r.Context(), subject, requestURI, resourceURI); err != nil {
			if errors.Is(err, innerauth.ErrForbidden) {
				logger.Warn(r.Context(), "inner_control_plane_authorization_denied",
					zap.String("credential_kind", string(subject.Kind)),
					zap.String("resource_uri", resourceURI),
					zap.Uint64("organization_id", subject.OrganizationID))
			}
			writeInnerAuthError(w, err)
			return
		}
		next.ServeHTTP(w, r.WithContext(withInnerAuthSubject(r.Context(), &subject)))
	})
}

func supportsDrive9Bearer(method, path string) bool {
	// Quota only supports Drive9 bearer reads; token routes support all mapped methods.
	return (method == http.MethodGet && path == "/v1/quota") || strings.HasPrefix(path, "/v1/tokens")
}

func readInnerControlPlaneCredentials(r *http.Request) (tenant.CredentialProvisionRequest, error) {
	credentials := tenant.CredentialProvisionRequest{
		PublicKey:  strings.TrimSpace(r.Header.Get(quotaPublicKeyHeader)),
		PrivateKey: strings.TrimSpace(r.Header.Get(quotaPrivateKeyHeader)),
	}
	if credentials.PublicKey != "" && credentials.PrivateKey != "" {
		return credentials, nil
	}
	if r.Body == nil || (r.Method != http.MethodPost && r.Method != http.MethodPatch && r.Method != http.MethodDelete) {
		return credentials, nil
	}

	body, err := io.ReadAll(io.LimitReader(r.Body, maxCredentialProvisionBodyBytes+1))
	r.Body = io.NopCloser(bytes.NewReader(body))
	if err != nil {
		return tenant.CredentialProvisionRequest{}, fmt.Errorf("%w: read request body", errInnerAuthMalformedRequest)
	}
	if len(body) > int(maxCredentialProvisionBodyBytes) {
		return tenant.CredentialProvisionRequest{}, fmt.Errorf("%w: request body too large", errInnerAuthMalformedRequest)
	}
	if len(bytes.TrimSpace(body)) == 0 {
		return credentials, nil
	}
	var bodyCredentials struct {
		PublicKey  string `json:"public_key"`
		PrivateKey string `json:"private_key"`
	}
	if err := json.Unmarshal(body, &bodyCredentials); err != nil {
		if credentials.PublicKey == "" || credentials.PrivateKey == "" {
			return tenant.CredentialProvisionRequest{}, fmt.Errorf("%w: invalid JSON body", errInnerAuthMalformedRequest)
		}
		return credentials, nil
	}
	if value := strings.TrimSpace(bodyCredentials.PublicKey); credentials.PublicKey == "" && value != "" {
		credentials.PublicKey = value
	}
	if value := strings.TrimSpace(bodyCredentials.PrivateKey); credentials.PrivateKey == "" && value != "" {
		credentials.PrivateKey = value
	}
	return credentials, nil
}

func writeInnerAuthError(w http.ResponseWriter, err error) {
	status := http.StatusBadGateway
	message := "inner authentication upstream failure"
	switch {
	case errors.Is(err, innerauth.ErrMalformedAuthorization),
		errors.Is(err, innerauth.ErrInvalidJWTOrganizationID),
		errors.Is(err, innerauth.ErrAmbiguousCredentials),
		errors.Is(err, innerauth.ErrBadRequest),
		errors.Is(err, tenant.ErrPartialCredentials),
		errors.Is(err, errInnerAuthMalformedRequest):
		status = http.StatusBadRequest
		message = "invalid credentials"
		if errors.Is(err, innerauth.ErrAmbiguousCredentials) {
			message = "ambiguous credentials"
		}
	case errors.Is(err, innerauth.ErrUnauthenticated), errors.Is(err, tenant.ErrCredentialsRequired):
		status = http.StatusUnauthorized
		message = "invalid credentials"
	case errors.Is(err, innerauth.ErrForbidden):
		status = http.StatusForbidden
		message = "forbidden"
	case errors.Is(err, innerauth.ErrDefaultCredentialsUnavailable):
		status = http.StatusServiceUnavailable
		message = "server credentials unavailable"
	}
	errJSON(w, status, message)
}