// Package callerauth authenticates protected systems before flex-auth evaluates
// the authorization request they submit.
package callerauth

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"time"
)

type Mode string

const (
	ModeDisabled Mode = "disabled"
	ModeWarn     Mode = "warn"
	ModeEnforce  Mode = "enforce"
)

var (
	ErrUnauthenticated = errors.New("caller is not authenticated")
	ErrForbidden       = errors.New("caller is not allowed to represent the requested system")
	ErrUnavailable     = errors.New("caller identity service is unavailable")
)

type Identity struct {
	Username  string
	Audiences []string
	NotAfter  time.Time
}

// Record is the caller provenance stamped onto a decision envelope.
// Mode is always set. Principal, audience and expiry are present only when a
// token was actually reviewed.
type Record struct {
	Mode      Mode
	Principal string
	Audience  string
	NotAfter  time.Time
}

type recordContextKey struct{}

func WithRecord(ctx context.Context, record Record) context.Context {
	return context.WithValue(ctx, recordContextKey{}, record)
}

func RecordFromContext(ctx context.Context) (Record, bool) {
	record, ok := ctx.Value(recordContextKey{}).(Record)
	return record, ok
}

type TokenReviewer interface {
	Review(context.Context, string) (Identity, error)
}

type WarningFunc func(string, ...any)

type Authenticator struct {
	mode     Mode
	reviewer TokenReviewer
	audience string
	bindings map[string]string
	warnf    WarningFunc
}

func New(mode Mode, reviewer TokenReviewer, audience string, bindings map[string]string, warnf WarningFunc) (*Authenticator, error) {
	switch mode {
	case ModeDisabled:
		return &Authenticator{mode: mode}, nil
	case ModeWarn, ModeEnforce:
	default:
		return nil, fmt.Errorf("unsupported caller-auth mode %q", mode)
	}
	if reviewer == nil {
		return nil, fmt.Errorf("token reviewer is required in %s mode", mode)
	}
	if strings.TrimSpace(audience) == "" {
		return nil, fmt.Errorf("caller audience is required in %s mode", mode)
	}
	if len(bindings) == 0 {
		return nil, fmt.Errorf("at least one caller binding is required in %s mode", mode)
	}
	copyBindings := make(map[string]string, len(bindings))
	for system, principal := range bindings {
		if strings.TrimSpace(system) == "" || strings.TrimSpace(principal) == "" {
			return nil, fmt.Errorf("caller bindings require non-empty system and principal")
		}
		copyBindings[system] = principal
	}
	return &Authenticator{mode: mode, reviewer: reviewer, audience: audience, bindings: copyBindings, warnf: warnf}, nil
}

func Disabled() *Authenticator {
	authenticator, _ := New(ModeDisabled, nil, "", nil, nil)
	return authenticator
}

// Authorize verifies the bearer token and binds every resource.system value to
// the authenticated workload principal. Warn mode records the same failures but
// permits the request so callers can be migrated before enforcement is enabled.
// The returned identity is populated whenever a token was reviewed, including
// warn-mode binding failures, so provenance can name an observed principal.
func (a *Authenticator) Authorize(ctx context.Context, authorization string, systems []string) (Identity, error) {
	if a == nil || a.mode == ModeDisabled {
		return Identity{}, nil
	}
	identity, err := a.authorize(ctx, authorization, systems)
	if err != nil && a.mode == ModeWarn {
		if a.warnf != nil {
			a.warnf("caller authentication warning: %v", err)
		}
		return identity, nil
	}
	return identity, err
}

// Record returns the provenance object for a reviewed identity. Disabled mode
// is stated as {"mode":"disabled"} with no principal.
func (a *Authenticator) Record(identity Identity) Record {
	mode := ModeDisabled
	audience := ""
	if a != nil {
		mode = a.mode
		audience = a.audience
	}
	record := Record{Mode: mode}
	if mode == ModeDisabled || strings.TrimSpace(identity.Username) == "" {
		return record
	}
	record.Principal = identity.Username
	record.Audience = audience
	record.NotAfter = identity.NotAfter
	return record
}

func (a *Authenticator) authorize(ctx context.Context, authorization string, systems []string) (Identity, error) {
	token, ok := strings.CutPrefix(authorization, "Bearer ")
	if !ok || strings.TrimSpace(token) == "" || strings.ContainsAny(strings.TrimSpace(token), " \t\r\n") {
		return Identity{}, ErrUnauthenticated
	}
	token = strings.TrimSpace(token)
	identity, err := a.reviewer.Review(ctx, token)
	if err != nil {
		if errors.Is(err, ErrUnauthenticated) {
			return Identity{}, err
		}
		return Identity{}, fmt.Errorf("%w: %v", ErrUnavailable, err)
	}
	if identity.NotAfter.IsZero() {
		if exp, ok := tokenExpiry(token); ok {
			identity.NotAfter = exp
		}
	}
	if strings.TrimSpace(identity.Username) == "" || !contains(identity.Audiences, a.audience) {
		return Identity{}, ErrUnauthenticated
	}
	if len(systems) == 0 {
		return identity, fmt.Errorf("%w: request has no resources", ErrForbidden)
	}
	for _, system := range systems {
		expected, found := a.bindings[system]
		if !found || expected != identity.Username {
			return identity, fmt.Errorf("%w: principal %q cannot represent system %q", ErrForbidden, identity.Username, system)
		}
	}
	return identity, nil
}

func contains(values []string, wanted string) bool {
	for _, value := range values {
		if value == wanted {
			return true
		}
	}
	return false
}