package main

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

	"github.com/Kong/go-pdk"
	authconsts "github.com/tidbcloud/aisupport/packages/sdk-go/infra/auth/consts"
)

const (
	defaultJWKSURL        = "https://forge.cdn.prod.atlassian-dev.net/.well-known/jwks.json"
	defaultRequestTimeout = 5 * time.Second
)

type Config struct {
	JWKSURL                        string   `json:"jwks_url"`
	AllowedAppIDs                  []string `json:"allowed_app_ids"`
	AllowedInstallationIDs         []string `json:"allowed_installation_ids"`
	AllowUnregisteredInstallations bool     `json:"allow_unregistered_installations"`
	RequirePrincipal               bool     `json:"require_principal"`
	RequestTimeoutMS               int      `json:"request_timeout_ms"`

	mu        sync.Mutex
	validator *forgeValidator
}

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

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

	rawToken, err := bearerToken(kong)
	if err != nil {
		abort(kong, http.StatusUnauthorized, "unauthorized")
		return
	}

	validator, err := c.getValidator()
	if err != nil {
		kong.Log.Err("invalid forge-auth configuration: ", err.Error())
		abort(kong, http.StatusInternalServerError, "internal authentication error")
		return
	}

	timeout := defaultRequestTimeout
	if c.RequestTimeoutMS > 0 {
		timeout = time.Duration(c.RequestTimeoutMS) * time.Millisecond
	}
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	identity, err := validator.validate(ctx, rawToken)
	if err != nil {
		kong.Log.Warn("forge authentication failed: ", err.Error())
		abort(kong, http.StatusUnauthorized, "unauthorized")
		return
	}
	if !identity.InstallationAllowlisted {
		kong.Log.Notice(
			"accepted Forge installation in bootstrap mode: app_id=", identity.AppID,
			" installation_id=", identity.InstallationID,
		)
	}

	if err := setForgeIdentityHeaders(kong, identity); err != nil {
		kong.Log.Err("set Forge identity headers: ", err.Error())
		abort(kong, http.StatusInternalServerError, "internal authentication error")
	}
}

func (c *Config) getValidator() (*forgeValidator, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.validator != nil {
		return c.validator, nil
	}
	validator, err := newForgeValidator(c)
	if err != nil {
		return nil, err
	}
	c.validator = validator
	return validator, nil
}

func bearerToken(kong *pdk.PDK) (string, error) {
	authorization, err := kong.Request.GetHeader(authconsts.HeaderAuthorization)
	if err != nil {
		return "", err
	}
	parts := strings.Fields(authorization)
	if len(parts) != 2 || !strings.EqualFold(parts[0], "Bearer") || parts[1] == "" {
		return "", errors.New("missing Forge Invocation Token")
	}
	return parts[1], nil
}

func setForgeIdentityHeaders(kong *pdk.PDK, identity forgeIdentity) error {
	// OAuth tokens are unnecessary for jsm-syncer. Remove all credentials before
	// forwarding the request so downstream services and logs cannot expose them.
	for _, header := range []string{
		authconsts.HeaderAuthorization,
		"X-Forge-OAuth-System",
		"X-Forge-OAuth-User",
	} {
		if err := kong.ServiceRequest.ClearHeader(header); err != nil {
			return err
		}
	}

	headers := map[string]string{
		authconsts.HeaderForgeAppID:          identity.AppID,
		authconsts.HeaderForgeInstallationID: identity.InstallationID,
		authconsts.HeaderForgeInvocationID:   identity.InvocationID,
		authconsts.HeaderAtlassianAccountID:  identity.Principal,
	}
	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 abort(kong *pdk.PDK, status int, message string) {
	body, _ := json.Marshal(map[string]any{
		"errcode": status,
		"errmsg":  message,
	})
	kong.Response.Exit(status, body, map[string][]string{
		"Content-Type": {"application/json"},
	})
}