package apptoken

import (
	"context"
	"fmt"
	"time"

	"github.com/hashicorp/vault/sdk/framework"
	"github.com/hashicorp/vault/sdk/logical"
)

// appRole describes the tokens issued from creds/<name>.
type appRole struct {
	// Audience is the JWT `aud` claim; it defaults to the role name and is the
	// value the consuming app checks. Empty means "use the role name".
	Audience string `json:"audience"`
	// TTL / MaxTTL bound the token lifetime.
	TTL    time.Duration `json:"ttl"`
	MaxTTL time.Duration `json:"max_ttl"`
	// AllowedSubjects, when non-empty, restricts which requesting identities
	// (Vault entity id or token display name) may mint a token from this role.
	AllowedSubjects []string `json:"allowed_subjects"`
	// Claims are extra string claims merged into every issued token. They may
	// not override registered JWT claims (iss, sub, aud, exp, ...).
	Claims map[string]string `json:"claims"`
}

func (r *appRole) audience(name string) string {
	if r.Audience != "" {
		return r.Audience
	}
	return name
}

func (r *appRole) toResponseData() map[string]interface{} {
	return map[string]interface{}{
		"audience":         r.Audience,
		"ttl":              int64(r.TTL.Seconds()),
		"max_ttl":          int64(r.MaxTTL.Seconds()),
		"allowed_subjects": r.AllowedSubjects,
		"claims":           r.Claims,
	}
}

func pathRole(b *appTokenBackend) *framework.Path {
	return &framework.Path{
		Pattern: "roles/" + framework.GenericNameRegex("name"),
		DisplayAttrs: &framework.DisplayAttributes{
			OperationPrefix: "apptoken",
			OperationSuffix: "role",
		},
		Fields: map[string]*framework.FieldSchema{
			"name": {
				Type:        framework.TypeLowerCaseString,
				Description: "Name of the role (typically the app name).",
				Required:    true,
			},
			"audience": {
				Type:        framework.TypeString,
				Description: "JWT audience (aud) claim. Defaults to the role name.",
			},
			"ttl": {
				Type:        framework.TypeDurationSecond,
				Description: "Default token TTL for this role.",
			},
			"max_ttl": {
				Type:        framework.TypeDurationSecond,
				Description: "Maximum token TTL for this role.",
			},
			"allowed_subjects": {
				Type:        framework.TypeCommaStringSlice,
				Description: "Optional allowlist of requesting identities (Vault entity id or token display name) permitted to mint tokens from this role. Empty means any authorized caller.",
			},
			"claims": {
				Type:        framework.TypeKVPairs,
				Description: "Extra key=value claims merged into each issued token. Registered JWT claims cannot be overridden.",
			},
		},
		Operations: map[logical.Operation]framework.OperationHandler{
			logical.ReadOperation:   &framework.PathOperation{Callback: b.pathRoleRead},
			logical.CreateOperation: &framework.PathOperation{Callback: b.pathRoleWrite},
			logical.UpdateOperation: &framework.PathOperation{Callback: b.pathRoleWrite},
			logical.DeleteOperation: &framework.PathOperation{Callback: b.pathRoleDelete},
		},
		ExistenceCheck:  b.pathRoleExistenceCheck,
		HelpSynopsis:    "Manage per-app token roles.",
		HelpDescription: "A role defines the audience, TTLs, subject allowlist and custom claims of the JWTs issued from creds/<name>.",
	}
}

func pathRolesList(b *appTokenBackend) *framework.Path {
	return &framework.Path{
		Pattern: "roles/?$",
		DisplayAttrs: &framework.DisplayAttributes{
			OperationPrefix: "apptoken",
			OperationSuffix: "roles",
		},
		Operations: map[logical.Operation]framework.OperationHandler{
			logical.ListOperation: &framework.PathOperation{Callback: b.pathRolesList},
		},
		HelpSynopsis: "List the configured roles.",
	}
}

func (b *appTokenBackend) pathRoleExistenceCheck(ctx context.Context, req *logical.Request, data *framework.FieldData) (bool, error) {
	role, err := b.getRole(ctx, req.Storage, data.Get("name").(string))
	if err != nil {
		return false, err
	}
	return role != nil, nil
}

func (b *appTokenBackend) pathRolesList(ctx context.Context, req *logical.Request, _ *framework.FieldData) (*logical.Response, error) {
	entries, err := req.Storage.List(ctx, roleStoragePrefix)
	if err != nil {
		return nil, err
	}
	return logical.ListResponse(entries), nil
}

func (b *appTokenBackend) pathRoleRead(ctx context.Context, req *logical.Request, data *framework.FieldData) (*logical.Response, error) {
	role, err := b.getRole(ctx, req.Storage, data.Get("name").(string))
	if err != nil {
		return nil, err
	}
	if role == nil {
		return nil, nil
	}
	return &logical.Response{Data: role.toResponseData()}, nil
}

func (b *appTokenBackend) pathRoleWrite(ctx context.Context, req *logical.Request, data *framework.FieldData) (*logical.Response, error) {
	name := data.Get("name").(string)
	if name == "" {
		return logical.ErrorResponse("role name is required"), nil
	}

	role, err := b.getRole(ctx, req.Storage, name)
	if err != nil {
		return nil, err
	}
	if role == nil {
		role = &appRole{}
	}

	if v, ok := data.GetOk("audience"); ok {
		role.Audience = v.(string)
	}
	if v, ok := data.GetOk("ttl"); ok {
		role.TTL = time.Duration(v.(int)) * time.Second
	}
	if v, ok := data.GetOk("max_ttl"); ok {
		role.MaxTTL = time.Duration(v.(int)) * time.Second
	}
	if v, ok := data.GetOk("allowed_subjects"); ok {
		role.AllowedSubjects = v.([]string)
	}
	if v, ok := data.GetOk("claims"); ok {
		role.Claims = v.(map[string]string)
	}

	if role.TTL < 0 || role.MaxTTL < 0 {
		return logical.ErrorResponse("ttl and max_ttl must not be negative"), nil
	}
	if role.MaxTTL != 0 && role.TTL > role.MaxTTL {
		return logical.ErrorResponse("ttl must not be greater than max_ttl"), nil
	}
	if err := validateClaims(role.Claims); err != nil {
		return logical.ErrorResponse(err.Error()), nil
	}

	if err := b.setRole(ctx, req.Storage, name, role); err != nil {
		return nil, err
	}
	return nil, nil
}

func (b *appTokenBackend) pathRoleDelete(ctx context.Context, req *logical.Request, data *framework.FieldData) (*logical.Response, error) {
	if err := req.Storage.Delete(ctx, roleStoragePrefix+data.Get("name").(string)); err != nil {
		return nil, fmt.Errorf("deleting role: %w", err)
	}
	return nil, nil
}

func (b *appTokenBackend) getRole(ctx context.Context, s logical.Storage, name string) (*appRole, error) {
	if name == "" {
		return nil, fmt.Errorf("missing role name")
	}
	entry, err := s.Get(ctx, roleStoragePrefix+name)
	if err != nil {
		return nil, err
	}
	if entry == nil {
		return nil, nil
	}
	role := &appRole{}
	if err := entry.DecodeJSON(role); err != nil {
		return nil, err
	}
	return role, nil
}

func (b *appTokenBackend) setRole(ctx context.Context, s logical.Storage, name string, role *appRole) error {
	entry, err := logical.StorageEntryJSON(roleStoragePrefix+name, role)
	if err != nil {
		return err
	}
	return s.Put(ctx, entry)
}

// reservedClaims are registered JWT claims the engine sets itself; roles may
// not override them via custom claims.
var reservedClaims = map[string]bool{
	"iss": true, "sub": true, "aud": true, "exp": true,
	"nbf": true, "iat": true, "jti": true,
}

func validateClaims(claims map[string]string) error {
	for k := range claims {
		if reservedClaims[k] {
			return fmt.Errorf("claim %q is reserved and cannot be set on a role", k)
		}
	}
	return nil
}