// Package innerauth authenticates and authorizes Drive9 control-plane callers
// through the account-provider inner APIs.
package innerauth

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"

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

const (
	userAuthnPath   = "/inner/account-provider/api/v1/user/authn"
	userAuthzPath   = "/inner/account-provider/api/v1/user/authz"
	apiKeyAuthnPath = "/inner/account-provider/api/v1/api-key/authn"
	apiKeyAuthzPath = "/inner/account-provider/api/v1/api-key/authz"

	maxResponseBytes = 1 << 20
)

var (
	ErrBadRequest          = errors.New("account-provider rejected the request")
	ErrUnauthenticated     = errors.New("account-provider authentication failed")
	ErrForbidden           = errors.New("account-provider authorization denied")
	ErrUpstreamUnavailable = errors.New("account-provider unavailable")
	ErrInvalidResponse     = errors.New("account-provider returned an invalid response")
)

type ClientConfig struct {
	BaseURL    string
	Token      string
	HTTPClient *http.Client
}

type Client struct {
	baseURL string
	token   string
	http    *http.Client
}

const basicAuthUsername = "drive9"

func NewClient(cfg ClientConfig) (*Client, error) {
	baseURL := strings.TrimRight(strings.TrimSpace(cfg.BaseURL), "/")
	parsed, err := url.Parse(baseURL)
	if err != nil || parsed.Scheme == "" || parsed.Host == "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" {
		return nil, errors.New("innerauth: BaseURL must be an absolute URL without credentials, query, or fragment")
	}
	if !strings.EqualFold(parsed.Scheme, "https") && !isLoopbackHTTPURL(parsed) {
		return nil, errors.New("innerauth: BaseURL must use https unless it targets a loopback address")
	}
	if strings.TrimSpace(cfg.Token) == "" {
		return nil, errors.New("innerauth: Basic Auth token is required")
	}
	httpClient := cfg.HTTPClient
	if httpClient == nil {
		httpClient = &http.Client{Timeout: 15 * time.Second}
	}
	clientCopy := *httpClient
	clientCopy.CheckRedirect = func(*http.Request, []*http.Request) error {
		return http.ErrUseLastResponse
	}
	return &Client{
		baseURL: baseURL,
		token:   cfg.Token,
		http:    &clientCopy,
	}, nil
}

// BaseURL returns the configured inner API base URL.
func (c *Client) BaseURL() string {
	return c.baseURL
}

func isLoopbackHTTPURL(parsed *url.URL) bool {
	if !strings.EqualFold(parsed.Scheme, "http") {
		return false
	}
	host := parsed.Hostname()
	if strings.EqualFold(host, "localhost") {
		return true
	}
	ip := net.ParseIP(host)
	return ip != nil && ip.IsLoopback()
}

type UserAuthnResult struct {
	UserID uint64
}

type APIKeyAuthnResult struct {
	PublicKey      string
	OrganizationID uint64
	ResourceInfos  []APIKeyResourceInfo
}

type APIKeyResourceInfo struct {
	ScopeType  string `json:"scope_type"`
	ResourceID uint64 `json:"resource_id"`
	Role       string `json:"role"`
	RoleID     uint64 `json:"role_id"`
}

type BillingPlanProvider interface {
	BillingPlan(ctx context.Context, organizationID uint64) (*tenant.TiDBCloudOrganizationPlan, error)
}

func (c *Client) UserAuthn(ctx context.Context, token string) (result UserAuthnResult, resultErr error) {
	ctx = c.withTrace(ctx)
	start := time.Now()
	var upstreamTraceID string
	defer func() { observeInnerIAMOperation(ctx, "user_authn", start, upstreamTraceID, &resultErr) }()
	var response struct {
		CheckResult bool `json:"check_result"`
		UserInfo    *struct {
			ID uint64 `json:"id"`
		} `json:"user_info"`
	}
	var err error
	upstreamTraceID, err = c.post(ctx, userAuthnPath, struct {
		Auth0Token string `json:"auth0_token"`
	}{Auth0Token: token}, &response)
	if err != nil {
		return UserAuthnResult{}, err
	}
	if !response.CheckResult {
		return UserAuthnResult{}, ErrUnauthenticated
	}
	if response.UserInfo == nil || response.UserInfo.ID == 0 {
		return UserAuthnResult{}, ErrInvalidResponse
	}
	logInnerAPIResult(ctx, "iam", "user_authn", "ok",
		zap.Uint64("user_id", response.UserInfo.ID),
		zap.String("upstream_trace_id", upstreamTraceID))
	return UserAuthnResult{UserID: response.UserInfo.ID}, nil
}

func (c *Client) UserAuthz(ctx context.Context, userID uint64, requestURI, resourceURI string) (resultErr error) {
	ctx = c.withTrace(ctx)
	start := time.Now()
	var upstreamTraceID string
	defer func() {
		observeInnerIAMOperation(ctx, "user_authz", start, upstreamTraceID, &resultErr,
			zap.Uint64("user_id", userID),
			zap.String("request_uri", requestURI),
			zap.String("resource_uri", resourceURI))
	}()
	var response struct {
		HasPermission bool `json:"has_permission"`
	}
	var err error
	upstreamTraceID, err = c.post(ctx, userAuthzPath, struct {
		UserID      uint64 `json:"user_id"`
		RequestURI  string `json:"req_uri"`
		ResourceURI string `json:"resource_uri"`
	}{UserID: userID, RequestURI: requestURI, ResourceURI: resourceURI}, &response)
	if err != nil {
		return err
	}
	if !response.HasPermission {
		return ErrForbidden
	}
	logInnerAPIResult(ctx, "iam", "user_authz", "ok",
		zap.Uint64("user_id", userID),
		zap.String("request_uri", requestURI),
		zap.String("resource_uri", resourceURI),
		zap.String("upstream_trace_id", upstreamTraceID))
	return nil
}

func (c *Client) APIKeyAuthn(ctx context.Context, encodedAPIKey string) (result APIKeyAuthnResult, resultErr error) {
	ctx = c.withTrace(ctx)
	start := time.Now()
	var upstreamTraceID string
	defer func() { observeInnerIAMOperation(ctx, "api_key_authn", start, upstreamTraceID, &resultErr) }()
	var response struct {
		PublicKey     string               `json:"public_key"`
		AuthRet       bool                 `json:"auth_ret"`
		ResourceInfos []APIKeyResourceInfo `json:"resource_infos"`
	}
	var err error
	upstreamTraceID, err = c.post(ctx, apiKeyAuthnPath, struct {
		APIKey string `json:"api_key"`
	}{APIKey: encodedAPIKey}, &response)
	if err != nil {
		return APIKeyAuthnResult{}, err
	}
	if !response.AuthRet {
		return APIKeyAuthnResult{}, ErrUnauthenticated
	}
	organizationID, organizationErr := organizationIDFromResourceInfos(response.ResourceInfos)
	if strings.TrimSpace(response.PublicKey) == "" || organizationErr != nil {
		if organizationErr != nil {
			return APIKeyAuthnResult{}, fmt.Errorf("%w: %v", ErrInvalidResponse, organizationErr)
		}
		return APIKeyAuthnResult{}, ErrInvalidResponse
	}
	fields := []zap.Field{
		zap.String("mode", "inner"),
		zap.String("api", "iam"),
		zap.String("operation", "api_key_authn"),
		zap.String("result", "ok"),
		zap.Uint64("organization_id", organizationID),
		zap.Int("scope_count", len(response.ResourceInfos)),
		zap.String("upstream_trace_id", upstreamTraceID),
	}
	if projectScope, ok := SelectHighestProjectScope(response.ResourceInfos); ok {
		fields = append(fields,
			zap.Uint64("project_id", projectScope.ResourceID),
			zap.String("project_role", projectScope.Role))
	}
	logInnerAPIResult(ctx, "iam", "api_key_authn", "ok", fields...)
	return APIKeyAuthnResult{
		PublicKey:      response.PublicKey,
		OrganizationID: organizationID,
		ResourceInfos:  response.ResourceInfos,
	}, nil
}

func organizationIDFromResourceInfos(resourceInfos []APIKeyResourceInfo) (uint64, error) {
	var organizationID uint64
	for _, resourceInfo := range resourceInfos {
		if !strings.EqualFold(strings.TrimSpace(resourceInfo.ScopeType), "ORG") {
			continue
		}
		if resourceInfo.ResourceID == 0 {
			return 0, errors.New("organization scope has zero resource id")
		}
		if organizationID != 0 && organizationID != resourceInfo.ResourceID {
			return 0, errors.New("multiple organization scopes")
		}
		organizationID = resourceInfo.ResourceID
	}
	if organizationID == 0 {
		return 0, errors.New("organization scope is missing")
	}
	return organizationID, nil
}

func (c *Client) APIKeyAuthz(ctx context.Context, publicKey, requestURI, resourceURI string) (resultErr error) {
	ctx = c.withTrace(ctx)
	start := time.Now()
	var upstreamTraceID string
	defer func() {
		observeInnerIAMOperation(ctx, "api_key_authz", start, upstreamTraceID, &resultErr,
			zap.String("request_uri", requestURI),
			zap.String("resource_uri", resourceURI))
	}()
	var response struct {
		AuthRet bool `json:"auth_ret"`
	}
	var err error
	upstreamTraceID, err = c.post(ctx, apiKeyAuthzPath, struct {
		PublicKey   string `json:"public_key"`
		RequestURI  string `json:"req_uri"`
		ResourceURI string `json:"resource_uri"`
	}{PublicKey: publicKey, RequestURI: requestURI, ResourceURI: resourceURI}, &response)
	if err != nil {
		return err
	}
	if !response.AuthRet {
		return ErrForbidden
	}
	logInnerAPIResult(ctx, "iam", "api_key_authz", "ok",
		zap.String("request_uri", requestURI),
		zap.String("resource_uri", resourceURI),
		zap.String("upstream_trace_id", upstreamTraceID))
	return nil
}

func (c *Client) BillingPlan(ctx context.Context, organizationID uint64) (result *tenant.TiDBCloudOrganizationPlan, resultErr error) {
	ctx = c.withTrace(ctx)
	start := time.Now()
	var upstreamTraceID string
	defer func() {
		metrics.RecordTiDBCloudOpenAPIRequest("billing", "get_plan", innerAPIResult(resultErr))
		metrics.RecordTiDBCloudOpenAPIRequestDuration("billing", "get_plan", time.Since(start))
		if resultErr != nil {
			fields := []zap.Field{
				zap.Uint64("organization_id", organizationID),
				zap.String("upstream_trace_id", upstreamTraceID),
				zap.Error(resultErr),
			}
			logInnerAPIResult(ctx, "billing", "get_plan", innerAPIResult(resultErr), fields...)
		}
	}()
	if organizationID == 0 {
		return nil, ErrBadRequest
	}
	requestPath := fmt.Sprintf("/inner/billing/tenants/%d/plans", organizationID)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.baseURL+requestPath, nil)
	if err != nil {
		return nil, fmt.Errorf("innerauth: build Billing request: %w", err)
	}
	c.setTraceHeaders(req, ctx)
	req.SetBasicAuth(basicAuthUsername, c.token)
	resp, err := c.http.Do(req)
	if err != nil {
		if ctxErr := ctx.Err(); ctxErr != nil {
			return nil, ctxErr
		}
		c.logInnerHTTPTransportFailure(ctx, http.MethodGet, requestPath, err)
		return nil, fmt.Errorf("%w: Billing request failed", ErrUpstreamUnavailable)
	}
	defer func() { _ = resp.Body.Close() }()
	upstreamTraceID = traceid.FromResponse(resp)
	raw, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes+1))
	if err != nil {
		return nil, fmt.Errorf("%w: read Billing response", ErrUpstreamUnavailable)
	}
	if statusErr := mapStatus(resp.StatusCode); statusErr != nil {
		if len(raw) > maxResponseBytes {
			c.logInnerHTTPFailure(ctx, http.MethodGet, requestPath, resp, nil, c.token)
		} else {
			c.logInnerHTTPFailure(ctx, http.MethodGet, requestPath, resp, raw, c.token)
		}
		return nil, statusErr
	}
	if len(raw) > maxResponseBytes {
		c.logInnerHTTPFailure(ctx, http.MethodGet, requestPath, resp, nil, c.token)
		return nil, fmt.Errorf("%w: Billing response too large", ErrInvalidResponse)
	}
	c.logInnerHTTPFailure(ctx, http.MethodGet, requestPath, resp, raw, c.token)
	plan, err := tenant.ParseTiDBCloudOrganizationPlan(raw, strconv.FormatUint(organizationID, 10))
	if err != nil {
		return nil, err
	}
	logInnerAPIResult(ctx, "billing", "get_plan", "ok",
		zap.Uint64("organization_id", organizationID),
		zap.String("effective_plan", plan.EffectivePlan),
		zap.String("poc_credits", plan.POCCredits),
		zap.Bool("is_free", plan.IsFree),
		zap.String("upstream_trace_id", upstreamTraceID))
	return plan, nil
}

func innerAPIResult(err error) string {
	switch {
	case err == nil:
		return "ok"
	case errors.Is(err, ErrUnauthenticated), errors.Is(err, ErrForbidden):
		return "denied"
	case errors.Is(err, ErrBadRequest):
		return "client_error"
	case errors.Is(err, ErrInvalidResponse):
		return "protocol_error"
	case errors.Is(err, tenant.ErrTiDBCloudBillingResponseInvalid):
		return "protocol_error"
	case errors.Is(err, context.Canceled):
		return "canceled"
	case errors.Is(err, context.DeadlineExceeded):
		return "timeout"
	default:
		return "upstream_error"
	}
}

func observeInnerIAMOperation(ctx context.Context, operation string, start time.Time, upstreamTraceID string, resultErr *error, fields ...zap.Field) {
	metrics.RecordTiDBCloudOpenAPIRequest("iam", operation, innerAPIResult(*resultErr))
	metrics.RecordTiDBCloudOpenAPIRequestDuration("iam", operation, time.Since(start))
	if *resultErr == nil {
		return
	}
	fields = append([]zap.Field{
		zap.String("upstream_trace_id", upstreamTraceID),
		zap.Error(*resultErr),
	}, fields...)
	logInnerAPIResult(ctx, "iam", operation, innerAPIResult(*resultErr), fields...)
}

func logInnerAPIResult(ctx context.Context, api, operation, result string, fields ...zap.Field) {
	fields = append([]zap.Field{
		zap.String("mode", "inner"),
		zap.String("api", api),
		zap.String("operation", operation),
		zap.String("result", result),
	}, fields...)
	switch result {
	case "ok":
		logger.Info(ctx, "tidbcloud_api_result", fields...)
	case "denied":
		logger.Warn(ctx, "tidbcloud_api_result", fields...)
	default:
		logger.Error(ctx, "tidbcloud_api_result", fields...)
	}
}

func (c *Client) post(ctx context.Context, requestPath string, requestBody, responseBody any) (string, error) {
	ctx = c.withTrace(ctx)
	rawBody, err := json.Marshal(requestBody)
	if err != nil {
		return "", fmt.Errorf("innerauth: encode account-provider request: %w", err)
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+requestPath, bytes.NewReader(rawBody))
	if err != nil {
		return "", fmt.Errorf("innerauth: build account-provider request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	c.setTraceHeaders(req, ctx)
	req.SetBasicAuth(basicAuthUsername, c.token)

	resp, err := c.http.Do(req)
	if err != nil {
		if ctxErr := ctx.Err(); ctxErr != nil {
			return "", ctxErr
		}
		c.logInnerHTTPTransportFailure(ctx, http.MethodPost, requestPath, err)
		return "", fmt.Errorf("%w: request failed", ErrUpstreamUnavailable)
	}
	defer func() { _ = resp.Body.Close() }()
	upstreamTraceID := traceid.FromResponse(resp)
	rawResponse, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes+1))
	if err != nil {
		return upstreamTraceID, fmt.Errorf("%w: read response", ErrUpstreamUnavailable)
	}
	secrets := append([]string{c.token}, requestSecretValues(rawBody)...)
	if statusErr := mapStatus(resp.StatusCode); statusErr != nil {
		if len(rawResponse) > maxResponseBytes {
			c.logInnerHTTPFailure(ctx, http.MethodPost, requestPath, resp, nil, secrets...)
		} else {
			c.logInnerHTTPFailure(ctx, http.MethodPost, requestPath, resp, rawResponse, secrets...)
		}
		return upstreamTraceID, statusErr
	}
	if len(rawResponse) > maxResponseBytes {
		c.logInnerHTTPFailure(ctx, http.MethodPost, requestPath, resp, nil, secrets...)
		return upstreamTraceID, fmt.Errorf("%w: response too large", ErrInvalidResponse)
	}
	c.logInnerHTTPFailure(ctx, http.MethodPost, requestPath, resp, rawResponse, secrets...)
	if err := json.Unmarshal(rawResponse, responseBody); err != nil {
		return upstreamTraceID, fmt.Errorf("%w: malformed JSON", ErrInvalidResponse)
	}
	return upstreamTraceID, nil
}

func (c *Client) withTrace(ctx context.Context) context.Context {
	if strings.TrimSpace(traceid.FromContext(ctx)) != "" {
		return ctx
	}
	return traceid.With(ctx, traceid.Generate())
}

func (c *Client) setTraceHeaders(req *http.Request, ctx context.Context) {
	traceID := strings.TrimSpace(traceid.FromContext(ctx))
	if traceID == "" {
		return
	}
	req.Header.Set("X-Trace-ID", traceID)
	req.Header.Set("X-Debug-Trace-Id", traceID)
}

func (c *Client) logInnerHTTPTransportFailure(ctx context.Context, method, requestPath string, err error) {
	logger.Error(ctx, "tidbcloud_inner_api_request_failed",
		zap.String("method", method),
		zap.String("path", requestPath),
		zap.Int("status", 0),
		zap.String("trace_id", traceid.FromContext(ctx)),
		zap.String("result", "transport_error"),
		zap.Error(err))
}

func (c *Client) logInnerHTTPFailure(ctx context.Context, method, requestPath string, resp *http.Response, raw []byte, secrets ...string) {
	if resp == nil {
		return
	}
	var envelope struct {
		AuthRet       *bool  `json:"auth_ret"`
		HasPermission *bool  `json:"has_permission"`
		CheckResult   *bool  `json:"check_result"`
		Message       string `json:"message"`
		CheckMessage  string `json:"check_message"`
	}
	_ = json.Unmarshal(raw, &envelope)
	denied := resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden ||
		(envelope.AuthRet != nil && !*envelope.AuthRet) ||
		(envelope.HasPermission != nil && !*envelope.HasPermission) ||
		(envelope.CheckResult != nil && !*envelope.CheckResult)
	if strings.TrimSpace(envelope.Message) == "" {
		envelope.Message = envelope.CheckMessage
	}
	if resp.StatusCode >= http.StatusOK && resp.StatusCode < http.StatusMultipleChoices && !denied {
		return
	}
	result := "http_error"
	if denied {
		result = "denied"
	}
	fields := []zap.Field{
		zap.String("method", method),
		zap.String("path", requestPath),
		zap.Int("status", resp.StatusCode),
		zap.String("trace_id", traceid.FromContext(ctx)),
		zap.String("upstream_trace_id", traceid.FromResponse(resp)),
		zap.String("result", result),
		zap.String("message", safeInnerMessage(envelope.Message, secrets...)),
	}
	if denied {
		logger.Warn(ctx, "tidbcloud_inner_api_request_failed", fields...)
		return
	}
	logger.Error(ctx, "tidbcloud_inner_api_request_failed", fields...)
}

func requestSecretValues(raw []byte) []string {
	var fields map[string]json.RawMessage
	if err := json.Unmarshal(raw, &fields); err != nil {
		return nil
	}
	secrets := make([]string, 0, 3)
	for _, field := range []string{"api_key", "public_key", "auth0_token"} {
		var value string
		if err := json.Unmarshal(fields[field], &value); err == nil && strings.TrimSpace(value) != "" {
			secrets = append(secrets, value)
		}
	}
	return secrets
}

func safeInnerMessage(message string, secrets ...string) string {
	message = strings.Join(strings.Fields(message), " ")
	for _, secret := range secrets {
		if secret != "" {
			message = strings.ReplaceAll(message, secret, "[REDACTED]")
		}
	}
	if len(message) > 512 {
		message = message[:512]
		for !utf8.ValidString(message) {
			message = message[:len(message)-1]
		}
	}
	return message
}

func mapStatus(status int) error {
	switch {
	case status >= http.StatusOK && status < http.StatusMultipleChoices:
		return nil
	case status == http.StatusBadRequest:
		return ErrBadRequest
	case status == http.StatusUnauthorized:
		return ErrUnauthenticated
	case status == http.StatusForbidden:
		return ErrForbidden
	default:
		return fmt.Errorf("%w: status %d", ErrUpstreamUnavailable, status)
	}
}