package exec

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

	"github.com/google/uuid"

	"github.com/tidbcloud/AgentTeam/desktop/internal/token"
	"github.com/tidbcloud/AgentTeam/desktop/internal/wire/apiclient"
)

const (
	defaultExecCredentialTTL = 12 * time.Hour
	// RuntimeAuthorizationStore bounds process authorizations separately from
	// bearer credentials. Token KeyManager reserves two credential slots per
	// authorization so rotation can always issue before retiring its predecessor.
	maxRuntimeAuthorizations   = 512
	maxCredentialGenerations   = 2
	maxAuthorizationInFlight   = 32
	maxGlobalAuthorizationUses = 1024
)

var errExecAuthorizationCapacity = errors.New("exec runtime authorization request capacity exhausted")

type workflowCompactionObserver interface {
	WorkflowStepCompleted(runID, stepID string)
	WorkflowStepUnsuccessful(runID, stepID string)
	WorkflowRunCanceled(runID string)
}

// runtimeAuthorization is the immutable result of authorizing one Exec request.
// It exposes data/services only; there is no execute or dispatch method.
type runtimeAuthorization struct {
	AuthorizationID    string
	ServerID           string
	AgentID            string
	SessionID          string
	RunID              string
	WorkDir            string
	ProcessGeneration  string
	RuntimeCatalog     *RuntimeCatalog
	WorkflowCompaction workflowCompactionObserver
	Server             *apiclient.Client
	ContextExpiry      time.Time
	revoked            <-chan struct{}
}

type credentialGeneration struct {
	expiresAt time.Time
}

type runtimeAuthorizationRecord struct {
	id          string
	snapshot    runtimeAuthorization
	tokenTTL    time.Duration
	generations map[string]credentialGeneration
	rotating    bool
	revoked     bool
	inFlight    int
	drained     chan struct{}
	drainedOnce sync.Once
	revokedDone chan struct{}
}

type execTokenManager interface {
	GenerateExecToken(token.ExecTokenRequest) (string, token.ExecTokenInfo, error)
	ValidateExecToken(agentID, serverID, value string) (token.ExecTokenInfo, error)
	RevokeExecToken(jti string) bool
}

// RuntimeAuthorizationStore owns process-lifetime Exec authorizations and
// their signed, short-lived bearer credential generations.
type RuntimeAuthorizationStore struct {
	mu             sync.Mutex
	tokens         execTokenManager
	byID           map[string]*runtimeAuthorizationRecord
	byJTI          map[string]*runtimeAuthorizationRecord
	now            func() time.Time
	pendingCreates int
	inFlight       int
	defaultCatalog *RuntimeCatalog
}

// RuntimeAuthorizationConfig defines one daemon-authored process authorization.
type RuntimeAuthorizationConfig struct {
	ServerID          string
	AgentID           string
	SessionID         string
	RunID             string
	WorkDir           string
	ProcessGeneration string
	RuntimeCatalog    *RuntimeCatalog
	// WorkflowCompaction observes only Server-authoritative terminal outcomes;
	// it does not participate in Exec authorization.
	WorkflowCompaction interface {
		WorkflowStepCompleted(runID, stepID string)
		WorkflowStepUnsuccessful(runID, stepID string)
		WorkflowRunCanceled(runID string)
	}
	Server         *apiclient.Client
	ContextExpires time.Time
	TokenTTL       time.Duration
}

// RuntimeCredentialLease identifies one process authorization and the bearer
// credential generation currently published to that process.
type RuntimeCredentialLease struct {
	AuthorizationID string
	CredentialID    string
	Token           string
	ExpiresAt       time.Time
}

// NewRuntimeAuthorizationStore creates an empty authorization store.
func NewRuntimeAuthorizationStore(tokens *token.KeyManager) *RuntimeAuthorizationStore {
	return &RuntimeAuthorizationStore{
		tokens: tokens,
		byID:   map[string]*runtimeAuthorizationRecord{},
		byJTI:  map[string]*runtimeAuthorizationRecord{},
		now:    time.Now,
	}
}

func (s *RuntimeAuthorizationStore) defaultRuntimeCatalog() (*RuntimeCatalog, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.defaultCatalog != nil {
		return s.defaultCatalog, nil
	}
	builtin, err := DefaultCatalog()
	if err != nil {
		return nil, err
	}
	s.defaultCatalog, err = NewRuntimeCatalog(builtin, ExtensionSnapshot{})
	return s.defaultCatalog, err
}

// Create records one process authorization and issues its first bearer
// credential. A candidate authorization never retires the active process
// authorization; exact replacement remains the dispatcher's responsibility.
func (s *RuntimeAuthorizationStore) Create(
	request RuntimeAuthorizationConfig,
) (RuntimeCredentialLease, error) {
	if s == nil || s.tokens == nil {
		return RuntimeCredentialLease{}, errors.New("exec credential token manager unavailable")
	}
	request.ServerID = strings.TrimSpace(request.ServerID)
	request.AgentID = strings.TrimSpace(request.AgentID)
	request.SessionID = strings.TrimSpace(request.SessionID)
	request.RunID = strings.TrimSpace(request.RunID)
	request.WorkDir = strings.TrimSpace(request.WorkDir)
	request.ProcessGeneration = strings.TrimSpace(request.ProcessGeneration)
	if request.Server == nil || request.ServerID == "" || request.AgentID == "" ||
		request.SessionID == "" || request.RunID == "" || request.WorkDir == "" {
		return RuntimeCredentialLease{}, errors.New("exec runtime authorization requires server, agent, session, run, workdir, and API client")
	}
	runtimeCatalog, err := s.runtimeCatalogForAuthorization(request.RuntimeCatalog)
	if err != nil {
		return RuntimeCredentialLease{}, fmt.Errorf("exec runtime authorization catalog: %w", err)
	}
	tokenTTL := request.TokenTTL
	if tokenTTL <= 0 {
		tokenTTL = defaultExecCredentialTTL
	}
	now := s.currentTime()
	if !request.ContextExpires.IsZero() && !now.Before(request.ContextExpires) {
		return RuntimeCredentialLease{}, errors.New("exec runtime authorization context is already expired")
	}
	// Reserve the authorization slot before signing so concurrent candidates
	// cannot overrun the process authorization bound.
	s.mu.Lock()
	// Revoked records with acquired requests remain lifecycle debt until those
	// requests drain. Count every retained record so repeated replacement
	// cannot accumulate unbounded immutable authorizations.
	if len(s.byID)+s.pendingCreates >= maxRuntimeAuthorizations {
		s.mu.Unlock()
		return RuntimeCredentialLease{}, errors.New("exec runtime authorization capacity exhausted")
	}
	s.pendingCreates++
	s.mu.Unlock()
	value, info, err := s.tokens.GenerateExecToken(token.ExecTokenRequest{
		AgentID: request.AgentID, ServerID: request.ServerID,
		SessionID: request.SessionID, RunID: request.RunID, TTL: tokenTTL,
		ProcessGeneration: request.ProcessGeneration,
		NotAfter:          request.ContextExpires,
	})
	if err != nil {
		s.mu.Lock()
		s.pendingCreates--
		s.mu.Unlock()
		return RuntimeCredentialLease{}, err
	}

	apiSnapshot := *request.Server
	snapshot := runtimeAuthorization{
		ServerID: request.ServerID, AgentID: request.AgentID,
		SessionID: request.SessionID, RunID: request.RunID, WorkDir: request.WorkDir,
		Server:             &apiSnapshot,
		ProcessGeneration:  request.ProcessGeneration,
		RuntimeCatalog:     runtimeCatalog,
		WorkflowCompaction: request.WorkflowCompaction,
		ContextExpiry:      request.ContextExpires,
	}
	record := &runtimeAuthorizationRecord{
		id: uuid.NewString(), snapshot: snapshot, tokenTTL: tokenTTL,
		generations: map[string]credentialGeneration{
			info.JTI: {expiresAt: info.ExpiresAt},
		},
		drained:     make(chan struct{}),
		revokedDone: make(chan struct{}),
	}

	s.mu.Lock()
	s.pendingCreates--
	s.byID[record.id] = record
	s.byJTI[info.JTI] = record
	s.mu.Unlock()
	return newRuntimeCredentialLease(record.id, value, info), nil
}

// Rotate issues another short-lived credential for the same immutable
// authorization. The predecessor remains active until the lifecycle owner
// explicitly retires it after publishing the replacement.
func (s *RuntimeAuthorizationStore) Rotate(
	authorizationID string,
) (RuntimeCredentialLease, error) {
	if s == nil || s.tokens == nil {
		return RuntimeCredentialLease{}, errors.New("exec credential token manager unavailable")
	}
	authorizationID = strings.TrimSpace(authorizationID)
	if authorizationID == "" {
		return RuntimeCredentialLease{}, errors.New("exec runtime authorization id is required")
	}

	s.mu.Lock()
	record := s.byID[authorizationID]
	if record == nil || record.revoked {
		s.mu.Unlock()
		return RuntimeCredentialLease{}, token.ErrAgentTokenInactive
	}
	now := s.currentTime()
	if !record.snapshot.ContextExpiry.IsZero() && !now.Before(record.snapshot.ContextExpiry) {
		s.mu.Unlock()
		return RuntimeCredentialLease{}, token.ErrAgentTokenInactive
	}
	s.pruneExpiredGenerationsLocked(record, now)
	if len(record.generations) >= maxCredentialGenerations {
		s.mu.Unlock()
		return RuntimeCredentialLease{}, errors.New("exec credential generation capacity exhausted")
	}
	if record.rotating {
		s.mu.Unlock()
		return RuntimeCredentialLease{}, errors.New("exec credential rotation already in progress")
	}
	record.rotating = true
	agentID := record.snapshot.AgentID
	serverID := record.snapshot.ServerID
	sessionID := record.snapshot.SessionID
	runID := record.snapshot.RunID
	processGeneration := record.snapshot.ProcessGeneration
	tokenTTL := record.tokenTTL
	contextExpiry := record.snapshot.ContextExpiry
	s.mu.Unlock()

	value, info, err := s.tokens.GenerateExecToken(token.ExecTokenRequest{
		AgentID: agentID, ServerID: serverID,
		SessionID: sessionID, RunID: runID,
		ProcessGeneration: processGeneration,
		TTL:               tokenTTL, NotAfter: contextExpiry,
	})
	s.mu.Lock()
	record.rotating = false
	if err != nil {
		s.mu.Unlock()
		return RuntimeCredentialLease{}, err
	}
	// Revoke may win while signing the candidate. Never publish a generation
	// unless the exact logical record is still active and within its context.
	now = s.currentTime()
	if s.byID[authorizationID] != record || record.revoked ||
		(!record.snapshot.ContextExpiry.IsZero() && !now.Before(record.snapshot.ContextExpiry)) {
		s.mu.Unlock()
		s.tokens.RevokeExecToken(info.JTI)
		return RuntimeCredentialLease{}, token.ErrAgentTokenInactive
	}
	record.generations[info.JTI] = credentialGeneration{expiresAt: info.ExpiresAt}
	s.byJTI[info.JTI] = record
	s.mu.Unlock()
	return newRuntimeCredentialLease(record.id, value, info), nil
}

// RetireCredential revokes one predecessor credential without affecting the
// process authorization or requests that already acquired its immutable data.
func (s *RuntimeAuthorizationStore) RetireCredential(
	authorizationID,
	credentialID string,
) bool {
	if s == nil || s.tokens == nil {
		return false
	}
	authorizationID = strings.TrimSpace(authorizationID)
	credentialID = strings.TrimSpace(credentialID)
	if authorizationID == "" || credentialID == "" {
		return false
	}

	s.mu.Lock()
	defer s.mu.Unlock()
	record := s.byID[authorizationID]
	if record == nil || record.revoked {
		return false
	}
	if _, ok := record.generations[credentialID]; !ok {
		return false
	}
	delete(record.generations, credentialID)
	delete(s.byJTI, credentialID)
	s.tokens.RevokeExecToken(credentialID)
	return true
}

// acquire validates the bearer and path before retaining the immutable process
// authorization. Body reads and action calls happen only after this succeeds.
func (s *RuntimeAuthorizationStore) acquire(
	_ context.Context,
	tokenValue,
	serverID,
	agentID string,
) (runtimeAuthorization, func(), error) {
	if s == nil || s.tokens == nil {
		return runtimeAuthorization{}, nil, errors.New("exec runtime authorization unavailable")
	}
	info, err := s.tokens.ValidateExecToken(agentID, serverID, tokenValue)
	if err != nil {
		return runtimeAuthorization{}, nil, err
	}
	s.mu.Lock()
	now := s.currentTime()
	record := s.byJTI[info.JTI]
	if record == nil || record.revoked ||
		record.snapshot.ServerID != info.ServerID ||
		record.snapshot.AgentID != info.AgentID ||
		record.snapshot.SessionID != info.SessionID ||
		record.snapshot.RunID != info.RunID ||
		record.snapshot.ProcessGeneration != info.ProcessGeneration ||
		record.snapshot.RuntimeCatalog == nil ||
		(!record.snapshot.ContextExpiry.IsZero() && !now.Before(record.snapshot.ContextExpiry)) {
		s.mu.Unlock()
		return runtimeAuthorization{}, nil, token.ErrAgentTokenInactive
	}
	if record.inFlight >= maxAuthorizationInFlight ||
		s.inFlight >= maxGlobalAuthorizationUses {
		s.mu.Unlock()
		return runtimeAuthorization{}, nil, errExecAuthorizationCapacity
	}
	record.inFlight++
	s.inFlight++
	snapshot := cloneRuntimeAuthorization(record.snapshot)
	snapshot.AuthorizationID = record.id
	snapshot.revoked = record.revokedDone
	s.mu.Unlock()

	var once sync.Once
	release := func() {
		once.Do(func() {
			s.mu.Lock()
			defer s.mu.Unlock()
			if record.inFlight <= 0 {
				return
			}
			record.inFlight--
			if s.inFlight > 0 {
				s.inFlight--
			}
			if record.revoked && record.inFlight == 0 {
				record.drainedOnce.Do(func() { close(record.drained) })
				if s.byID[record.id] == record {
					delete(s.byID, record.id)
				}
			}
		})
	}
	return snapshot, release, nil
}

// Revoke retires every bearer credential for one process authorization.
func (s *RuntimeAuthorizationStore) Revoke(authorizationID string) <-chan struct{} {
	authorizationID = strings.TrimSpace(authorizationID)
	if s == nil || authorizationID == "" {
		return closedChannel()
	}
	s.mu.Lock()
	record := s.byID[authorizationID]
	if record == nil {
		s.mu.Unlock()
		return closedChannel()
	}
	if !record.revoked {
		record.revoked = true
		close(record.revokedDone)
		for jti := range record.generations {
			delete(s.byJTI, jti)
			s.tokens.RevokeExecToken(jti)
		}
		clear(record.generations)
		if record.inFlight == 0 {
			record.drainedOnce.Do(func() { close(record.drained) })
			delete(s.byID, authorizationID)
		}
	}
	drained := record.drained
	s.mu.Unlock()
	return drained
}

func (s *RuntimeAuthorizationStore) currentTime() time.Time {
	if s != nil && s.now != nil {
		return s.now()
	}
	return time.Now()
}

func (s *RuntimeAuthorizationStore) pruneExpiredGenerationsLocked(
	record *runtimeAuthorizationRecord,
	now time.Time,
) {
	if record == nil {
		return
	}
	for jti, generation := range record.generations {
		if !now.Before(generation.expiresAt) {
			delete(record.generations, jti)
			delete(s.byJTI, jti)
		}
	}
}