package auth

import (
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"math/big"
	"strings"
	"sync"
	"time"

	"upspin.io/factotum"
	"upspin.io/upspin"
)

// Vault SSO token format (agreed with the vault issuer):
//
//	v1.<payload>.<r>.<s>
//
// where <payload> is base64url(JSON claims) and <r>,<s> are base64url of the
// big-endian bytes of the ECDSA signature over the ASCII string
// "v1.<payload>" made with the issuer's upspin key. Splitting the token on "."
// yields exactly four fields. Keeping R and S as separate segments makes the
// encoding independent of the signing curve (P-256 vs P-384).
const vaultTokenPrefix = "v1."

// vaultClaims are the token's JSON payload.
type vaultClaims struct {
	Iss string `json:"iss"` // issuer upspin user, e.g. "vault@cosmocr.at"
	Sub string `json:"sub"` // authenticated upspin user
	Aud string `json:"aud"` // intended audience, e.g. "gitweb"
	Iat int64  `json:"iat"` // issued-at (unix seconds)
	Exp int64  `json:"exp"` // expiry (unix seconds)
}

// LooksLikeVaultToken reports whether s has the vault token shape, so the
// authenticator can route it to VaultVerifier instead of the PAT path.
func LooksLikeVaultToken(s string) bool {
	return strings.HasPrefix(s, vaultTokenPrefix)
}

// KeyLookup returns the upspin public key of a user (typically via the
// keyserver). It is injected so tests can supply a fixed key.
type KeyLookup func(upspin.UserName) (upspin.PublicKey, error)

// VaultVerifier validates vault-issued SSO bearer tokens. It trusts tokens
// signed by a single pinned issuer for a single audience, verifying the
// signature against the issuer's upspin public key.
type VaultVerifier struct {
	issuer   upspin.UserName
	audience string
	lookup   KeyLookup
	now      func() time.Time

	mu  sync.Mutex
	key upspin.PublicKey // cached issuer key
}

// NewVaultVerifier returns a verifier that accepts tokens whose iss equals
// issuer and aud equals audience, checking signatures with keys from lookup.
func NewVaultVerifier(issuer upspin.UserName, audience string, lookup KeyLookup) *VaultVerifier {
	return &VaultVerifier{issuer: issuer, audience: audience, lookup: lookup, now: time.Now}
}

// Verify checks a token and returns the authenticated subject (upspin user).
func (v *VaultVerifier) Verify(token string) (string, error) {
	if !strings.HasPrefix(token, vaultTokenPrefix) {
		return "", fmt.Errorf("not a vault token")
	}
	parts := strings.Split(token, ".")
	if len(parts) != 4 || parts[0] != "v1" {
		return "", fmt.Errorf("malformed vault token")
	}
	payloadB64, rB64, sB64 := parts[1], parts[2], parts[3]

	payload, err := base64.RawURLEncoding.DecodeString(payloadB64)
	if err != nil {
		return "", fmt.Errorf("bad token payload: %w", err)
	}
	var claims vaultClaims
	if err := json.Unmarshal(payload, &claims); err != nil {
		return "", fmt.Errorf("bad token claims: %w", err)
	}

	if claims.Iss != string(v.issuer) {
		return "", fmt.Errorf("token issuer %q not trusted", claims.Iss)
	}
	if claims.Aud != v.audience {
		return "", fmt.Errorf("token audience %q not for this service", claims.Aud)
	}
	if claims.Sub == "" {
		return "", fmt.Errorf("token has no subject")
	}
	if v.now().Unix() >= claims.Exp {
		return "", fmt.Errorf("token expired")
	}

	sig, err := decodeSignature(rB64, sB64)
	if err != nil {
		return "", err
	}
	key, err := v.issuerKey()
	if err != nil {
		return "", err
	}
	hash := sha256.Sum256([]byte(vaultTokenPrefix + payloadB64))
	if err := factotum.Verify(hash[:], sig, key); err != nil {
		return "", fmt.Errorf("token signature invalid: %w", err)
	}
	return claims.Sub, nil
}

func (v *VaultVerifier) issuerKey() (upspin.PublicKey, error) {
	v.mu.Lock()
	key := v.key
	v.mu.Unlock()
	if key != "" {
		return key, nil
	}
	key, err := v.lookup(v.issuer)
	if err != nil {
		return "", fmt.Errorf("look up issuer key: %w", err)
	}
	v.mu.Lock()
	v.key = key
	v.mu.Unlock()
	return key, nil
}

func decodeSignature(rB64, sB64 string) (upspin.Signature, error) {
	rb, err := base64.RawURLEncoding.DecodeString(rB64)
	if err != nil {
		return upspin.Signature{}, fmt.Errorf("bad signature R: %w", err)
	}
	sb, err := base64.RawURLEncoding.DecodeString(sB64)
	if err != nil {
		return upspin.Signature{}, fmt.Errorf("bad signature S: %w", err)
	}
	return upspin.Signature{
		R: new(big.Int).SetBytes(rb),
		S: new(big.Int).SetBytes(sb),
	}, nil
}