package token

import (
	"crypto/ed25519"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/google/uuid"
)

const (
	execTokenPrefix     = "et-dev"
	execTokenVersion    = 1
	defaultExecTokenTTL = 12 * time.Hour
	// RuntimeAuthorizationStore admits at most 512 process authorizations and
	// two overlapping credentials per authorization. Reserve the full overlap
	// so a full store can still issue every pre-expiry rotation.
	maxExecTokens = 1024
)

// ErrExecTokenCapacity indicates that every bounded Exec token slot is still
// active. Issuance fails closed instead of silently invalidating another
// runtime capability.
var ErrExecTokenCapacity = errors.New("exec token capacity exhausted")

// ExecTokenRequest contains daemon-only lifecycle fields. Only AgentID, the
// generated JTI, and expiry are placed in the signed token payload.
type ExecTokenRequest struct {
	AgentID           string
	ServerID          string
	SessionID         string
	RunID             string
	ProcessGeneration string
	TTL               time.Duration
	NotAfter          time.Time
}

// ExecTokenInfo is the validated active token record returned to the daemon.
type ExecTokenInfo struct {
	JTI               string
	AgentID           string
	ServerID          string
	SessionID         string
	RunID             string
	ProcessGeneration string
	ExpiresAt         time.Time
}

type execTokenClaims struct {
	Version int    `json:"version"`
	AgentID string `json:"agent_id"`
	JTI     string `json:"jti"`
	Exp     int64  `json:"exp"`
}

type execTokenRecord struct {
	agentID           string
	serverID          string
	sessionID         string
	runID             string
	processGeneration string
	expiresAt         time.Time
	revokedAt         time.Time
}

// GenerateExecToken creates an Ed25519-signed runtime capability. Its payload
// contains only version/agent_id/jti/exp; server/session/run stay in the active
// daemon record and are checked again during validation.
func (m *KeyManager) GenerateExecToken(request ExecTokenRequest) (string, ExecTokenInfo, error) {
	if m == nil {
		return "", ExecTokenInfo{}, errors.New("api key manager is nil")
	}
	request.AgentID = strings.TrimSpace(request.AgentID)
	request.ServerID = strings.TrimSpace(request.ServerID)
	request.SessionID = strings.TrimSpace(request.SessionID)
	request.RunID = strings.TrimSpace(request.RunID)
	request.ProcessGeneration = strings.TrimSpace(request.ProcessGeneration)
	if request.AgentID == "" || request.ServerID == "" || request.SessionID == "" || request.RunID == "" {
		return "", ExecTokenInfo{}, errors.New("exec token requires agent, server, session, and run identity")
	}
	ttl := request.TTL
	if ttl <= 0 {
		ttl = defaultExecTokenTTL
	}
	now := m.currentTime()
	expiryCandidate := now.Add(ttl)
	if !request.NotAfter.IsZero() && request.NotAfter.Before(expiryCandidate) {
		expiryCandidate = request.NotAfter
	}
	// exp is encoded as whole Unix seconds. Keep the active token record and
	// returned expiry on exactly that same boundary so callers never believe a
	// token is valid after its signed claim has expired.
	expiresAt := time.Unix(expiryCandidate.Unix(), 0)
	if !now.Before(expiresAt) {
		return "", ExecTokenInfo{}, errors.New("exec token TTL does not reach the next signed second")
	}
	claims := execTokenClaims{
		Version: execTokenVersion,
		AgentID: request.AgentID,
		JTI:     uuid.NewString(),
		Exp:     expiresAt.Unix(),
	}
	payload, err := json.Marshal(claims)
	if err != nil {
		return "", ExecTokenInfo{}, fmt.Errorf("marshal exec token claims: %w", err)
	}

	m.mu.Lock()
	defer m.mu.Unlock()
	if m.priKey == nil {
		return "", ExecTokenInfo{}, errors.New("api key manager not initialized")
	}
	m.pruneExecTokensLocked(now)
	if len(m.execTokens) >= maxExecTokens {
		return "", ExecTokenInfo{}, ErrExecTokenCapacity
	}
	signature := ed25519.Sign(m.priKey, payload)
	value := fmt.Sprintf(
		"%s.%s.%s",
		execTokenPrefix,
		base64.RawURLEncoding.EncodeToString(payload),
		base64.RawURLEncoding.EncodeToString(signature),
	)
	if m.execTokens == nil {
		m.execTokens = make(map[string]execTokenRecord)
	}
	m.execTokens[claims.JTI] = execTokenRecord{
		agentID: request.AgentID, serverID: request.ServerID,
		sessionID: request.SessionID, runID: request.RunID,
		processGeneration: request.ProcessGeneration,
		expiresAt:         expiresAt,
	}
	return value, ExecTokenInfo{
		JTI: claims.JTI, AgentID: request.AgentID, ServerID: request.ServerID,
		SessionID: request.SessionID, RunID: request.RunID,
		ProcessGeneration: request.ProcessGeneration,
		ExpiresAt:         expiresAt,
	}, nil
}

// ValidateExecToken verifies signature, signed expiry, path agent/server
// binding, and the daemon-side active record.
func (m *KeyManager) ValidateExecToken(agentID, serverID, value string) (ExecTokenInfo, error) {
	if m == nil {
		return ExecTokenInfo{}, errors.New("api key manager is nil")
	}
	agentID = strings.TrimSpace(agentID)
	serverID = strings.TrimSpace(serverID)
	if agentID == "" || serverID == "" {
		return ExecTokenInfo{}, errors.New("agent and server id are required")
	}
	payload, signature, err := parseKeyWithPrefix(strings.TrimSpace(value), execTokenPrefix)
	if err != nil {
		return ExecTokenInfo{}, err
	}
	m.mu.RLock()
	publicKey := append(ed25519.PublicKey(nil), m.pubKey...)
	m.mu.RUnlock()
	if len(publicKey) != ed25519.PublicKeySize {
		return ExecTokenInfo{}, errors.New("api key manager not initialized")
	}
	if !ed25519.Verify(publicKey, payload, signature) {
		return ExecTokenInfo{}, ErrInvalidSignature
	}
	var claims execTokenClaims
	if err := json.Unmarshal(payload, &claims); err != nil {
		return ExecTokenInfo{}, fmt.Errorf("%w: invalid exec token claims", ErrInvalidFormat)
	}
	claims.AgentID = strings.TrimSpace(claims.AgentID)
	claims.JTI = strings.TrimSpace(claims.JTI)
	if claims.Version != execTokenVersion || claims.AgentID == "" || claims.JTI == "" || claims.Exp <= 0 {
		return ExecTokenInfo{}, fmt.Errorf("%w: incomplete exec token claims", ErrInvalidFormat)
	}
	if claims.AgentID != agentID {
		return ExecTokenInfo{}, ErrAgentMismatch
	}
	now := m.currentTime()
	signedExpiry := time.Unix(claims.Exp, 0)
	if !now.Before(signedExpiry) {
		return ExecTokenInfo{}, ErrExpired
	}

	m.mu.Lock()
	defer m.mu.Unlock()
	record, ok := m.execTokens[claims.JTI]
	if !ok {
		return ExecTokenInfo{}, ErrAgentTokenInactive
	}
	if !record.revokedAt.IsZero() {
		return ExecTokenInfo{}, ErrAgentTokenRevoked
	}
	if !now.Before(record.expiresAt) {
		delete(m.execTokens, claims.JTI)
		return ExecTokenInfo{}, ErrExpired
	}
	if record.agentID != claims.AgentID {
		return ExecTokenInfo{}, ErrAgentMismatch
	}
	if record.serverID != serverID {
		return ExecTokenInfo{}, ErrAgentTokenServerMismatch
	}
	if !record.expiresAt.Equal(signedExpiry) {
		return ExecTokenInfo{}, ErrInvalidSignature
	}
	return ExecTokenInfo{
		JTI: claims.JTI, AgentID: record.agentID, ServerID: record.serverID,
		SessionID: record.sessionID, RunID: record.runID,
		ProcessGeneration: record.processGeneration,
		ExpiresAt:         record.expiresAt,
	}, nil
}

// RevokeExecToken immediately makes one JTI inactive for future validation.
func (m *KeyManager) RevokeExecToken(jti string) bool {
	if m == nil {
		return false
	}
	jti = strings.TrimSpace(jti)
	if jti == "" {
		return false
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	record, ok := m.execTokens[jti]
	if !ok {
		return false
	}
	if record.revokedAt.IsZero() {
		record.revokedAt = m.currentTime()
		m.execTokens[jti] = record
	}
	return true
}

func (m *KeyManager) pruneExecTokensLocked(now time.Time) {
	for jti, record := range m.execTokens {
		if !record.revokedAt.IsZero() || !now.Before(record.expiresAt) {
			delete(m.execTokens, jti)
		}
	}
}