package machine_identity

import (
	"bytes"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"net/url"
	"os"
	"path"
	"regexp"
	"slices"
	"strings"
	"sync"
	"time"
)

var (
	principalPattern = regexp.MustCompile(`^yv:(sense|brain|bell):[a-z0-9][a-z0-9.-]{0,62}$`)
	keyIDPattern     = regexp.MustCompile(`^[A-Za-z0-9._-]{8,64}$`)
	tokenIDPattern   = regexp.MustCompile(`^[A-Za-z0-9_-]{22,64}$`)
	validAudiences   = map[string]bool{"yovision-sense": true, "yovision-brain": true, "yovision-bell": true}
	validScopes      = map[string]bool{"source-config:write": true, "runtime-status:write": true, "events:ingest": true, "evidence:read": true}
)

const (
	Version       = "yovision.machine-identity/v1"
	TokenType     = "YOVISION-MACHINE+JWT"
	MaxLifetime   = 5 * time.Minute
	AllowedSkew   = 30 * time.Second
	MaxKeyOverlap = 24 * time.Hour
)

type Error struct{ Code string }

func (e *Error) Error() string { return e.Code }

func codeError(code string) error { return &Error{Code: code} }

// BearerToken deliberately has no cookie or query fallback.
func BearerToken(authorization string) (string, error) {
	parts := strings.Split(authorization, " ")
	if len(parts) != 2 || !strings.EqualFold(parts[0], "Bearer") || parts[1] == "" || strings.ContainsAny(parts[1], " \t\r\n,") {
		return "", codeError("machine_token_missing")
	}
	return parts[1], nil
}

type Claims struct {
	Version    string   `json:"ver"`
	Issuer     string   `json:"iss"`
	Subject    string   `json:"sub"`
	Audience   string   `json:"aud"`
	Scopes     []string `json:"scope"`
	IssuedAt   int64    `json:"iat"`
	NotBefore  int64    `json:"nbf"`
	ExpiresAt  int64    `json:"exp"`
	TokenID    string   `json:"jti"`
	Method     string   `json:"htm"`
	Path       string   `json:"htu"`
	BodySHA256 string   `json:"body_sha256"`
}

type protectedHeader struct {
	Algorithm string `json:"alg"`
	Type      string `json:"typ"`
	KeyID     string `json:"kid"`
	Version   string `json:"ver"`
}

type KeyRecord struct {
	Principal string
	KeyID     string
	PublicKey ed25519.PublicKey
	Audience  string
	Scopes    []string
	Enabled   bool
	Revoked   bool
}

type Registry struct {
	mu   sync.RWMutex
	keys map[string]KeyRecord
}

func NewRegistry(records ...KeyRecord) (*Registry, error) {
	r := &Registry{keys: make(map[string]KeyRecord, len(records))}
	for _, record := range records {
		if !keyIDPattern.MatchString(record.KeyID) || !principalPattern.MatchString(record.Principal) || !validAudiences[record.Audience] || len(record.PublicKey) != ed25519.PublicKeySize || !validScopeList(record.Scopes) {
			return nil, errors.New("invalid machine key record")
		}
		if _, exists := r.keys[record.KeyID]; exists {
			return nil, errors.New("duplicate machine key id")
		}
		record.PublicKey = slices.Clone(record.PublicKey)
		record.Scopes = slices.Clone(record.Scopes)
		r.keys[record.KeyID] = record
	}
	return r, nil
}

func (r *Registry) Lookup(keyID string) (KeyRecord, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	record, ok := r.keys[keyID]
	record.PublicKey = slices.Clone(record.PublicKey)
	record.Scopes = slices.Clone(record.Scopes)
	return record, ok
}

func (r *Registry) Revoke(keyID string) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	record, ok := r.keys[keyID]
	if !ok {
		return false
	}
	record.Revoked = true
	r.keys[keyID] = record
	return true
}

type ReplayStore struct {
	mu   sync.Mutex
	used map[string]time.Time
}

// ReplayCache must atomically persist accepted (principal, jti) pairs until
// expiry. ReplayStore is process-local and intended for tests or a single
// uninterrupted process; connector implementations inject a durable store.
type ReplayCache interface {
	Consume(principal, tokenID string, expiresAt, now time.Time) bool
}

func NewReplayStore() *ReplayStore { return &ReplayStore{used: map[string]time.Time{}} }

func (s *ReplayStore) Consume(principal, tokenID string, expiresAt, now time.Time) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	for key, expiry := range s.used {
		if !expiry.After(now) {
			delete(s.used, key)
		}
	}
	key := principal + "\x00" + tokenID
	if _, exists := s.used[key]; exists {
		return false
	}
	s.used[key] = expiresAt
	return true
}

type Signer struct {
	Principal  string
	KeyID      string
	PrivateKey ed25519.PrivateKey
	Now        func() time.Time
}

func LoadPrivateKey(path string) (ed25519.PrivateKey, error) {
	if strings.TrimSpace(path) == "" {
		return nil, errors.New("machine private key path is required")
	}
	raw, err := os.ReadFile(path)
	if err != nil {
		return nil, errors.New("read machine private key")
	}
	block, rest := pem.Decode(raw)
	if block == nil || len(bytes.TrimSpace(rest)) != 0 || block.Type != "PRIVATE KEY" {
		return nil, errors.New("machine private key must be one PKCS#8 PEM block")
	}
	parsed, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		return nil, errors.New("parse machine private key")
	}
	key, ok := parsed.(ed25519.PrivateKey)
	if !ok || len(key) != ed25519.PrivateKeySize {
		return nil, errors.New("machine private key is not Ed25519")
	}
	return slices.Clone(key), nil
}

func (s Signer) Mint(audience string, scopes []string, method, requestPath string, body []byte) (string, error) {
	if !principalPattern.MatchString(s.Principal) || !keyIDPattern.MatchString(s.KeyID) || len(s.PrivateKey) != ed25519.PrivateKeySize || !validAudiences[audience] || !validScopeList(scopes) {
		return "", errors.New("incomplete machine signer configuration")
	}
	normalizedPath, err := normalizePath(requestPath)
	if err != nil {
		return "", err
	}
	method = strings.ToUpper(method)
	if !allowedMethod(method) {
		return "", errors.New("unsupported machine request method")
	}
	now := time.Now().UTC()
	if s.Now != nil {
		now = s.Now().UTC()
	}
	tokenID, err := randomTokenID()
	if err != nil {
		return "", err
	}
	digest := sha256.Sum256(body)
	claims := Claims{Version: Version, Issuer: s.Principal, Subject: s.Principal, Audience: audience,
		Scopes: slices.Clone(scopes), IssuedAt: now.Unix(), NotBefore: now.Unix(), ExpiresAt: now.Add(MaxLifetime).Unix(),
		TokenID: tokenID, Method: method, Path: normalizedPath, BodySHA256: hex.EncodeToString(digest[:])}
	header := protectedHeader{Algorithm: "EdDSA", Type: TokenType, KeyID: s.KeyID, Version: Version}
	headerJSON, _ := json.Marshal(header)
	claimsJSON, _ := json.Marshal(claims)
	signingInput := rawBase64(headerJSON) + "." + rawBase64(claimsJSON)
	signature := ed25519.Sign(s.PrivateKey, []byte(signingInput))
	return signingInput + "." + rawBase64(signature), nil
}

type Verifier struct {
	Registry *Registry
	Replay   ReplayCache
	Now      func() time.Time
}

func (v Verifier) Verify(token, audience, requiredScope, method, requestPath string, body []byte) (Claims, error) {
	if v.Registry == nil || v.Replay == nil {
		return Claims{}, codeError("machine_token_invalid")
	}
	parts := strings.Split(token, ".")
	if len(parts) != 3 || strings.Contains(token, "=") {
		return Claims{}, codeError("machine_token_invalid")
	}
	headerBytes, err := decodeRaw(parts[0])
	if err != nil {
		return Claims{}, codeError("machine_token_invalid")
	}
	var header protectedHeader
	if err = decodeClosed(headerBytes, &header); err != nil || header.Algorithm != "EdDSA" || header.Type != TokenType || header.Version != Version || !keyIDPattern.MatchString(header.KeyID) {
		return Claims{}, codeError("machine_token_invalid")
	}
	record, ok := v.Registry.Lookup(header.KeyID)
	if !ok {
		return Claims{}, codeError("machine_token_invalid")
	}
	signature, err := decodeRaw(parts[2])
	if err != nil || len(signature) != ed25519.SignatureSize || !ed25519.Verify(record.PublicKey, []byte(parts[0]+"."+parts[1]), signature) {
		return Claims{}, codeError("machine_token_invalid")
	}
	if !record.Enabled || record.Revoked {
		return Claims{}, codeError("machine_identity_revoked")
	}
	claimsBytes, err := decodeRaw(parts[1])
	if err != nil {
		return Claims{}, codeError("machine_token_invalid")
	}
	var claims Claims
	if err = decodeClosed(claimsBytes, &claims); err != nil || !validClaimsShape(claims) || claims.Issuer != record.Principal || claims.Subject != record.Principal {
		return Claims{}, codeError("machine_token_invalid")
	}
	now := time.Now().UTC()
	if v.Now != nil {
		now = v.Now().UTC()
	}
	nowUnix := now.Unix()
	if claims.ExpiresAt-claims.IssuedAt <= 0 || claims.ExpiresAt-claims.IssuedAt > int64(MaxLifetime/time.Second) ||
		claims.NotBefore < claims.IssuedAt || claims.NotBefore > claims.ExpiresAt || claims.IssuedAt > nowUnix+int64(AllowedSkew/time.Second) {
		return Claims{}, codeError("machine_token_invalid")
	}
	if claims.NotBefore > nowUnix+int64(AllowedSkew/time.Second) || claims.ExpiresAt < nowUnix-int64(AllowedSkew/time.Second) {
		return Claims{}, codeError("machine_token_expired")
	}
	if claims.Audience != audience || record.Audience != audience {
		return Claims{}, codeError("machine_audience_denied")
	}
	if !slices.Contains(claims.Scopes, requiredScope) || !slices.Contains(record.Scopes, requiredScope) {
		return Claims{}, codeError("machine_scope_denied")
	}
	normalizedPath, err := normalizePath(requestPath)
	digest := sha256.Sum256(body)
	if err != nil || claims.Method != strings.ToUpper(method) || claims.Path != normalizedPath || claims.BodySHA256 != hex.EncodeToString(digest[:]) {
		return Claims{}, codeError("machine_token_invalid")
	}
	if !v.Replay.Consume(claims.Issuer, claims.TokenID, time.Unix(claims.ExpiresAt, 0).Add(AllowedSkew), now) {
		return Claims{}, codeError("machine_token_replayed")
	}
	return claims, nil
}

func decodeClosed(raw []byte, target any) error {
	decoder := json.NewDecoder(bytes.NewReader(raw))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(target); err != nil {
		return err
	}
	if err := decoder.Decode(&struct{}{}); !errors.Is(err, io.EOF) {
		if err == nil {
			return errors.New("trailing JSON value")
		}
		return err
	}
	return nil
}

func validClaimsShape(claims Claims) bool {
	if claims.Version != Version || !principalPattern.MatchString(claims.Issuer) || claims.Subject != claims.Issuer || !validAudiences[claims.Audience] || !tokenIDPattern.MatchString(claims.TokenID) ||
		len(claims.Scopes) == 0 || len(claims.Scopes) > 4 || !allowedMethod(claims.Method) || claims.Path == "" || len(claims.BodySHA256) != 64 {
		return false
	}
	if !validScopeList(claims.Scopes) {
		return false
	}
	_, err := hex.DecodeString(claims.BodySHA256)
	return err == nil
}

func validScopeList(scopes []string) bool {
	if len(scopes) == 0 || len(scopes) > 4 {
		return false
	}
	seen := map[string]bool{}
	for _, scope := range scopes {
		if !validScopes[scope] || seen[scope] {
			return false
		}
		seen[scope] = true
	}
	return true
}

func normalizePath(value string) (string, error) {
	parsed, err := url.ParseRequestURI(value)
	if err != nil || parsed.IsAbs() || parsed.Host != "" || parsed.RawQuery != "" || parsed.Fragment != "" || parsed.Path == "" || !strings.HasPrefix(parsed.Path, "/") || strings.Contains(parsed.Path, "\\") || strings.Contains(parsed.Path, "//") || path.Clean(parsed.Path) != parsed.Path {
		return "", errors.New("machine request path must be a normalized absolute path without query or fragment")
	}
	return parsed.EscapedPath(), nil
}

func allowedMethod(method string) bool {
	switch method {
	case "GET", "POST", "PUT", "PATCH", "DELETE":
		return true
	default:
		return false
	}
}

func randomTokenID() (string, error) {
	raw := make([]byte, 16)
	if _, err := rand.Read(raw); err != nil {
		return "", fmt.Errorf("generate machine token id: %w", err)
	}
	return rawBase64(raw), nil
}

func rawBase64(value []byte) string { return base64.RawURLEncoding.EncodeToString(value) }

func decodeRaw(value string) ([]byte, error) {
	return base64.RawURLEncoding.Strict().DecodeString(value)
}