package main

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"errors"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/Kong/go-pdk"
	baseapi "github.com/tidbcloud/aisupport/packages/sdk-go/contracts/api/base"
	authv1 "github.com/tidbcloud/aisupport/packages/sdk-go/contracts/api/core-server/auth/v1"
	userv1 "github.com/tidbcloud/aisupport/packages/sdk-go/contracts/api/core-server/user/v1"
	authconsts "github.com/tidbcloud/aisupport/packages/sdk-go/infra/auth/consts"
	"github.com/tidbcloud/aisupport/packages/sdk-go/infra/errs"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
)

const (
	defaultRequestTimeout = 30 * time.Second
)

type Config struct {
	AuthServiceAddr     string `json:"auth_service_addr"`
	AuthServiceInsecure bool   `json:"auth_service_insecure"`
	RequestTimeoutMS    int    `json:"request_timeout_ms"`

	mu         sync.Mutex
	authClient authv1.AuthServiceClient
	conn       *grpc.ClientConn
}

type credential struct {
	accessToken string
	apiToken    string
	apiKey      string
}

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

func (c *Config) Access(kong *pdk.PDK) {
	method, err := kong.Request.GetMethod()
	if err != nil {
		abort(kong, err)
		return
	}
	if method == http.MethodOptions {
		return
	}

	ctx, cancel := context.WithTimeout(context.Background(), c.requestTimeout())
	defer cancel()

	cred, err := c.credential(kong)
	if err != nil {
		abort(kong, errs.APIErrorNew(http.StatusUnauthorized, err))
		return
	}

	client, err := c.client()
	if err != nil {
		abort(kong, err)
		return
	}

	clientIP, err := kong.Client.GetForwardedIp()
	if err != nil {
		abort(kong, err)
		return
	}
	ctx = metadata.AppendToOutgoingContext(ctx, strings.ToLower(authconsts.HeaderClientIP), clientIP)

	authnResp, err := client.Authenticate(ctx, authenticateRequest(cred))
	if err != nil {
		code := http.StatusUnauthorized
		if status.Code(err) == codes.PermissionDenied {
			code = http.StatusForbidden
		}
		abort(kong, errs.APIErrorNew(code, err))
		return
	}

	profile := authnResp.PrincipalProfile

	uri, err := kong.Request.GetPath()
	if err != nil {
		abort(kong, err)
		return
	}
	authzResp, err := client.Authorize(ctx, authorizeRequest(authnResp, uri, method))
	if err != nil {
		abort(kong, errs.APIErrorNew(http.StatusForbidden, err))
		return
	}
	if !authzResp.GetAuthorized() {
		abort(kong, errs.APIErrorf(http.StatusForbidden, "permission denied"))
		return
	}

	if err := setIdentityHeaders(kong, profile, clientIP); err != nil {
		abort(kong, err)
	}
}

func (c *Config) credential(kong *pdk.PDK) (credential, error) {
	authHeader, err := kong.Request.GetHeader(authconsts.HeaderAuthorization)
	if err != nil {
		return credential{}, err
	}
	if token := accessTokenFromAuthHeader(authHeader); token != "" {
		return credential{accessToken: token}, nil
	}

	apiToken, err := kong.Request.GetHeader(authconsts.HeaderAPIToken)
	if err != nil {
		return credential{}, err
	}
	if apiToken != "" {
		return credential{apiToken: apiToken}, nil
	}

	apiKey, err := kong.Request.GetHeader(authconsts.HeaderAPIKey)
	if err != nil {
		return credential{}, err
	}
	if apiKey != "" {
		return credential{apiKey: apiKey}, nil
	}

	return credential{}, errors.New("miss credential")
}

func accessTokenFromAuthHeader(header string) string {
	if header == "" {
		return ""
	}
	fields := strings.Split(header, " ")
	if len(fields) != 2 || strings.ToLower(fields[0]) != "bearer" {
		return ""
	}
	return fields[1]
}

func authenticateRequest(cred credential) *authv1.AuthenticateRequest {
	req := &authv1.AuthenticateRequest{}
	switch {
	case cred.accessToken != "":
		req.AccessToken = &cred.accessToken
	case cred.apiToken != "":
		req.ApiToken = &cred.apiToken
	case cred.apiKey != "":
		req.ApiKey = &cred.apiKey
	}
	return req
}

func authorizeRequest(authnResp *authv1.AuthenticateResponse, uri string, method string) *authv1.AuthorizeRequest {
	profile := authnResp.PrincipalProfile
	req := &authv1.AuthorizeRequest{
		Uri:    uri,
		Method: method,
	}
	switch profile.OperatorType {
	case authconsts.OperatorTypeAPIKey:
		req.ApiKeyName = &profile.OperatorName
	case authconsts.OperatorTypeAPIToken:
		req.UserEmail = &profile.UserEmail
		req.ApiTokenPrivilege = &authnResp.ApiToken.Privilege
	default:
		req.UserEmail = &profile.UserEmail
	}
	return req
}

func setIdentityHeaders(kong *pdk.PDK, profile *userv1.PrincipalProfile, clientIP string) error {
	data, err := json.Marshal(profile)
	if err != nil {
		return err
	}
	headers := map[string]string{
		authconsts.HeaderOperatorType:     profile.OperatorType,
		authconsts.HeaderOperator:         profile.OperatorName,
		authconsts.HeaderUserEmail:        profile.UserEmail,
		authconsts.HeaderPrincipalProfile: string(data),
	}
	if clientIP != "" {
		headers[authconsts.HeaderClientIP] = clientIP
	}

	for name, value := range headers {
		if err := kong.ServiceRequest.SetHeader(name, value); err != nil {
			return err
		}
		if err := kong.ServiceRequest.SetHeader(authconsts.GRPCMetadataKeyFromHTTPHeader(name), value); err != nil {
			return err
		}
	}
	return nil
}

func (c *Config) client() (authv1.AuthServiceClient, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.authClient != nil {
		return c.authClient, nil
	}
	if c.AuthServiceAddr == "" {
		return nil, errors.New("auth_service_addr is required")
	}

	opts := []grpc.DialOption{}
	if c.AuthServiceInsecure {
		opts = append(opts, grpc.WithTransportCredentials(insecure.NewCredentials()))
	} else {
		opts = append(opts, grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{})))
	}

	conn, err := grpc.Dial(c.AuthServiceAddr, opts...)
	if err != nil {
		return nil, err
	}
	c.conn = conn
	c.authClient = authv1.NewAuthServiceClient(conn)
	return c.authClient, nil
}

func (c *Config) requestTimeout() time.Duration {
	if c.RequestTimeoutMS > 0 {
		return time.Duration(c.RequestTimeoutMS) * time.Millisecond
	}
	return defaultRequestTimeout
}

func abort(kong *pdk.PDK, err error) {
	aerr := errs.APIErrorWrap(err)
	body, _ := json.Marshal(&baseapi.BaseResponse{
		ErrCode:    aerr.ErrCode(),
		ErrMessage: aerr.ErrMessage(),
	})

	headers := map[string][]string{
		"Content-Type": {"application/json"},
	}
	if origin, _ := kong.Request.GetHeader("Origin"); origin != "" {
		headers["Access-Control-Allow-Origin"] = []string{origin}
		headers["Access-Control-Allow-Credentials"] = []string{"true"}
		headers["Vary"] = []string{"Origin"}
	}
	kong.Response.Exit(aerr.ErrCode(), body, headers)
}