package dispatcher

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"net/url"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
	"github.com/tidbcloud/AgentTeam/desktop/internal/api/exec"
	"github.com/tidbcloud/AgentTeam/desktop/internal/execruntimeenv"
	"github.com/tidbcloud/AgentTeam/desktop/internal/runtimecredential"
	"github.com/tidbcloud/AgentTeam/desktop/internal/wire/apiclient"
)

const (
	execGenerationTTL                      = 12 * time.Hour
	maxExecAuthorizationOwners             = 512
	execCredentialCleanupInitialRetryDelay = 250 * time.Millisecond
	execCredentialCleanupMaxRetryDelay     = 30 * time.Second
)

var errExecAuthorizationOwnerCapacity = errors.New("activate Exec: runtime authorization capacity exhausted")

type agentteamAPIClientProvider interface {
	APIClient(agentID ...string) *apiclient.Client
}

type execResolved struct {
	api            *apiclient.Client
	catalog        *exec.Catalog
	runtimeCatalog *exec.RuntimeCatalog
}

type execAuthorizationCarrier interface {
	ExecAuthorizationID() string
}

type execCredential interface {
	Path() string
	Replace(string) error
	Close() error
}

type execAuthorizationOwner struct {
	id                   string
	agentID              string
	serverID             string
	tokenTTL             time.Duration
	credential           execCredential
	stop                 chan struct{}
	done                 chan struct{}
	stopOnce             sync.Once
	cleanupOnce          sync.Once
	cleanupDone          chan struct{}
	authorizationDrained <-chan struct{}
}

type retiringExecCredential struct {
	id       string
	retireAt time.Time
}

func (d *execAuthorizationService) prepareExec(plc *processLaunchContext) (*execResolved, error) {
	if plc == nil || plc.sess == nil {
		return nil, fmt.Errorf("prepare Exec: launch context is required")
	}
	ch, ok := d.serverConnections.ServerConnection(strings.TrimSpace(plc.binding.channelID))
	if !ok {
		return nil, nil
	}
	provider, ok := ch.(agentteamAPIClientProvider)
	if !ok {
		return nil, nil
	}
	api := provider.APIClient(plc.binding.serverAgentID)
	if api == nil || strings.TrimSpace(api.AgentID) == "" || strings.TrimSpace(api.AuthToken) == "" || strings.TrimSpace(api.RouteWorkspaceID) == "" {
		return nil, nil
	}
	api.LocalAgentID = strings.TrimSpace(plc.sess.agentID)

	registry, err := exec.DefaultCatalog()
	if err != nil {
		return nil, err
	}
	extension, err := d.prepareExtension(plc)
	if err != nil {
		return nil, err
	}
	extensionSnapshot := exec.ExtensionSnapshot{}
	if extension != nil {
		extensionSnapshot = extension.snapshot
	}
	runtimeCatalog, err := exec.NewRuntimeCatalog(registry, extensionSnapshot)
	if err != nil {
		return nil, fmt.Errorf("prepare Exec Runtime Catalog: %w", err)
	}
	// Report isolated Extension action diagnostics to operators. Each
	// diagnostic names the extension, the action, the failing schema stage, and
	// a bounded root cause — never raw schema bodies, input, or secrets. The
	// The snapshot cache makes identical publications a no-op, so repeated
	// prepares do not refetch or republish the same response.
	logIsolatedExtensionDiagnostics(runtimeCatalog, plc.binding.channelID, plc.sess.agentID)
	return &execResolved{
		api:            api,
		catalog:        registry,
		runtimeCatalog: runtimeCatalog,
	}, nil
}

// logIsolatedExtensionDiagnostics writes one structured operator-visible Warn
// per isolated Extension action diagnostic. Attrs carry the server and agent
// identity when available; each diagnostic is bounded publish-safe text
// (extension= action= schema= + root cause) and never includes raw schema
// bodies, input, or secrets.
func logIsolatedExtensionDiagnostics(
	runtimeCatalog *exec.RuntimeCatalog,
	serverID string,
	agentID string,
) {
	if runtimeCatalog == nil {
		return
	}
	diagnostics := runtimeCatalog.Diagnostics()
	if len(diagnostics) == 0 {
		return
	}
	var attrs []any
	if serverID = strings.TrimSpace(serverID); serverID != "" {
		attrs = append(attrs, "server", serverID)
	}
	if agentID = strings.TrimSpace(agentID); agentID != "" {
		attrs = append(attrs, "agent", agentID)
	}
	for _, diagnostic := range diagnostics {
		slog.Warn("exec: isolated Extension action (not published)", append(attrs, "detail", diagnostic)...)
	}
}

func (d *execAuthorizationService) activateExec(_ context.Context, plc *processLaunchContext, sessionID string, resolved *execResolved) (*drivercontract.ExecConfig, error) {
	if resolved == nil {
		return nil, nil
	}
	if err := d.reserveExecAuthorizationActivation(); err != nil {
		return nil, err
	}
	reservationHeld := true
	defer func() {
		if reservationHeld {
			d.releaseExecAuthorizationActivation()
		}
	}()
	serverID := strings.TrimSpace(plc.binding.channelID)
	agentID := strings.TrimSpace(plc.sess.agentID)
	if strings.Contains(serverID, "/") || strings.Contains(agentID, "/") {
		return nil, errors.New("activate Exec: server and agent ids must be URL path segments")
	}
	sessionID = runtimeAuthorizationSessionID(agentID, sessionID)
	runID := execRunID(plc, sessionID)
	lease, err := d.execAuthorizations.Create(execRuntimeAuthorizationConfig(plc, sessionID, runID, resolved))
	if err != nil {
		return nil, err
	}
	credential, err := d.createExecCredential(lease.Token)
	if err != nil {
		d.execAuthorizations.Revoke(lease.AuthorizationID)
		slog.Warn(
			"agentteam exec credential creation failed",
			"phase", "create", "agent_id", agentID,
			"server_id", serverID,
		)
		return nil, errors.New("create Exec runtime credential: unavailable")
	}
	credentialPath := strings.TrimSpace(credential.Path())
	if credentialPath == "" || !filepath.IsAbs(credentialPath) {
		d.discardUnpublishedExecAuthorization(lease.AuthorizationID, credential, agentID, serverID)
		return nil, errors.New("create Exec runtime credential: invalid path")
	}
	authorizationOwner := &execAuthorizationOwner{
		id: lease.AuthorizationID, agentID: agentID,
		serverID: serverID, tokenTTL: execGenerationTTL,
		credential: credential, stop: make(chan struct{}), done: make(chan struct{}),
		cleanupDone: make(chan struct{}),
	}
	if !d.execCatalogStillCurrent(serverID, agentID, resolved) {
		d.discardUnpublishedExecAuthorization(lease.AuthorizationID, credential, agentID, serverID)
		return nil, errors.New("activate Exec: Extension Catalog changed during activation")
	}
	d.catalogMu.RLock()
	if !d.execCatalogStillCurrentLocked(serverID, agentID, resolved) {
		d.catalogMu.RUnlock()
		d.discardUnpublishedExecAuthorization(lease.AuthorizationID, credential, agentID, serverID)
		return nil, errors.New("activate Exec: Extension Catalog changed during activation")
	}
	d.mu.Lock()
	if d.admissionClosed {
		d.mu.Unlock()
		d.catalogMu.RUnlock()
		d.discardUnpublishedExecAuthorization(lease.AuthorizationID, credential, agentID, serverID)
		return nil, errors.New("activate Exec: dispatcher is shutting down")
	}
	if d.authorizations == nil {
		d.authorizations = make(map[string]*execAuthorizationOwner)
	}
	d.pendingActivations--
	reservationHeld = false
	d.authorizations[lease.AuthorizationID] = authorizationOwner
	d.mu.Unlock()
	d.catalogMu.RUnlock()
	d.warnIfExecExtensionDegraded(serverID, agentID, resolved)
	go d.runExecCredentialRotation(authorizationOwner, lease)

	return &drivercontract.ExecConfig{
		BaseURL: execEndpoint(
			d.serverURL,
			"/api/internal/exec/servers/"+url.PathEscape(serverID)+
				"/agents/"+url.PathEscape(agentID),
		),
		AuthorizationID: lease.AuthorizationID, TokenFile: credentialPath,
	}, nil
}

func (d *execAuthorizationService) execCatalogStillCurrent(
	serverID string,
	agentID string,
	resolved *execResolved,
) bool {
	if d == nil {
		return false
	}
	d.catalogMu.RLock()
	defer d.catalogMu.RUnlock()
	return d.execCatalogStillCurrentLocked(serverID, agentID, resolved)
}

// warnIfExecExtensionDegraded makes a degraded activation operator-visible:
// the desired Catalog is enabled but this authorization carries zero Extension
// actions (the fetch failed or every unit was skipped at prepare time). The
// prepare-time Warn names the cause; this one marks the committed state.
func (d *execAuthorizationService) warnIfExecExtensionDegraded(
	serverID string,
	agentID string,
	resolved *execResolved,
) {
	if resolved == nil || resolved.runtimeCatalog == nil {
		return
	}
	if !resolved.runtimeCatalog.ExtensionSnapshot().Empty() {
		return
	}
	d.catalogMu.RLock()
	hint, observed := d.desired[extensionCatalogKey{
		serverID: strings.TrimSpace(serverID),
		agentID:  strings.TrimSpace(agentID),
	}]
	d.catalogMu.RUnlock()
	if observed && hint.Version > 0 {
		slog.Warn(
			"agentteam exec activation without Extensions (degraded)",
			"agent_id", agentID,
			"server_id", serverID,
		)
	}
}

func (d *execAuthorizationService) execCatalogStillCurrentLocked(
	serverID string,
	agentID string,
	resolved *execResolved,
) bool {
	if resolved == nil || resolved.runtimeCatalog == nil {
		return true
	}
	snapshot := resolved.runtimeCatalog.ExtensionSnapshot()
	key := extensionCatalogKey{
		serverID: strings.TrimSpace(serverID),
		agentID:  strings.TrimSpace(agentID),
	}
	hint, observed := d.desired[key]
	if snapshot.Empty() {
		// An empty snapshot authorizes zero Extension actions, so there is
		// nothing stale to fail-close on: this guard only protects against a
		// NON-empty stale snapshot staying callable after the hint moved.
		// Returning true here is also what lets a degraded launch (Extension
		// fetch failed; the agent started with built-ins only) activate
		// instead of aborting. Revocation of a previous process's
		// authorization is owned by ObserveAgentExtensionCatalogHints.
		return true
	}
	cached, cachedExists := d.snapshots[key]
	publishedRevision, published := d.published[key]
	return observed && cachedExists && published &&
		publishedRevision == hint.Version &&
		cached.ETag() == snapshot.ETag()
}

func execRuntimeAuthorizationConfig(
	plc *processLaunchContext,
	sessionID string,
	runID string,
	resolved *execResolved,
) exec.RuntimeAuthorizationConfig {
	serverID := strings.TrimSpace(plc.binding.channelID)
	agentID := strings.TrimSpace(plc.sess.agentID)
	return exec.RuntimeAuthorizationConfig{
		ServerID:           serverID,
		AgentID:            agentID,
		SessionID:          sessionID,
		RunID:              runID,
		WorkDir:            plc.sess.workDir,
		RuntimeCatalog:     resolved.runtimeCatalog,
		WorkflowCompaction: plc.workflowCompaction,
		Server:             resolved.api,
		ContextExpires:     time.Time{},
		TokenTTL:           execGenerationTTL,
	}
}

func (d *execAuthorizationService) reserveExecAuthorizationActivation() error {
	if d == nil {
		return errors.New("activate Exec: dispatcher is unavailable")
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.admissionClosed {
		return errors.New("activate Exec: dispatcher is shutting down")
	}
	// Active authorization owners and retrying cleanup debt share one bound. Reserve before
	// RuntimeAuthorizationStore.Create so concurrent launch attempts cannot create an
	// unbounded number of credential owners outside the tracked authorization map.
	if len(d.authorizations)+d.pendingActivations >= maxExecAuthorizationOwners {
		return errExecAuthorizationOwnerCapacity
	}
	d.pendingActivations++
	return nil
}

func (d *execAuthorizationService) releaseExecAuthorizationActivation() {
	if d == nil {
		return
	}
	d.mu.Lock()
	if d.pendingActivations > 0 {
		d.pendingActivations--
	}
	d.mu.Unlock()
}

func (d *execAuthorizationService) closeExecAuthorizationAdmission() {
	if d == nil {
		return
	}
	d.mu.Lock()
	d.admissionClosed = true
	d.mu.Unlock()
}

func (d *execAuthorizationService) createExecCredential(value string) (execCredential, error) {
	if d.credentialFactory != nil {
		return d.credentialFactory(value)
	}
	return runtimecredential.Create(value)
}

func (d *execAuthorizationService) runExecCredentialRotation(authorizationOwner *execAuthorizationOwner, current exec.RuntimeCredentialLease) {
	defer close(authorizationOwner.done)
	nextRenewal := current.ExpiresAt.Add(-authorizationOwner.tokenTTL / 2)
	var retiring *retiringExecCredential
	for {
		now := d.execNow()
		if retiring != nil && !now.Before(retiring.retireAt) {
			d.execAuthorizations.RetireCredential(authorizationOwner.id, retiring.id)
			retiring = nil
		}
		if !now.Before(nextRenewal) {
			next, err := d.execAuthorizations.Rotate(authorizationOwner.id)
			if err != nil {
				completedAt := d.execNow()
				slog.Warn(
					"agentteam Exec authorization rotation failed",
					"phase", "issue", "agent_id", authorizationOwner.agentID,
					"server_id", authorizationOwner.serverID,
				)
				nextRenewal = completedAt.Add(d.execRetryDelay(authorizationOwner.tokenTTL))
			} else if err := authorizationOwner.credential.Replace(next.Token); err != nil {
				d.execAuthorizations.RetireCredential(authorizationOwner.id, next.CredentialID)
				completedAt := d.execNow()
				slog.Warn(
					"agentteam Exec authorization rotation failed",
					"phase", "publish", "agent_id", authorizationOwner.agentID,
					"server_id", authorizationOwner.serverID,
				)
				nextRenewal = completedAt.Add(d.execRetryDelay(authorizationOwner.tokenTTL))
			} else {
				if retiring != nil {
					d.execAuthorizations.RetireCredential(authorizationOwner.id, retiring.id)
				}
				// Never extend a predecessor, but keep it authorized through its
				// already-signed expiry. A CLI suspended after reading the old file
				// can therefore resume without a rotation-boundary auth failure.
				retiring = &retiringExecCredential{id: current.CredentialID, retireAt: current.ExpiresAt}
				current = next
				nextRenewal = current.ExpiresAt.Add(-authorizationOwner.tokenTTL / 2)
			}
			continue
		}

		nextWake := nextRenewal
		if retiring != nil && retiring.retireAt.Before(nextWake) {
			nextWake = retiring.retireAt
		}
		delay := nextWake.Sub(now)
		if delay < 0 {
			delay = 0
		}
		if !d.execWait(authorizationOwner.stop, delay) {
			return
		}
	}
}

func (d *execAuthorizationService) execNow() time.Time {
	if d.now != nil {
		return d.now()
	}
	return time.Now()
}

func (d *execAuthorizationService) execWait(stop <-chan struct{}, duration time.Duration) bool {
	if d.wait != nil {
		return d.wait(stop, duration)
	}
	timer := time.NewTimer(duration)
	defer timer.Stop()
	select {
	case <-timer.C:
		return true
	case <-stop:
		return false
	}
}

func (d *execAuthorizationService) execRetryDelay(ttl time.Duration) time.Duration {
	delay := ttl / 100
	if delay < time.Second {
		delay = time.Second
	}
	if delay > time.Minute {
		delay = time.Minute
	}
	return delay
}

func (d *execAuthorizationService) revokeExecProcessAuthorization(
	process execAuthorizationCarrier,
) <-chan struct{} {
	if d == nil || process == nil {
		return closedExecDrain()
	}
	return d.revokeExecRuntimeAuthorization(process.ExecAuthorizationID())
}

func (d *execAuthorizationService) revokeExecAuthorizationConfig(cfg *drivercontract.ExecConfig) <-chan struct{} {
	if d == nil || cfg == nil {
		return closedExecDrain()
	}
	return d.revokeExecRuntimeAuthorization(cfg.AuthorizationID)
}

func (d *execAuthorizationService) revokeExecRuntimeAuthorization(authorizationID string) <-chan struct{} {
	authorizationID = strings.TrimSpace(authorizationID)
	if d == nil || d.execAuthorizations == nil || authorizationID == "" {
		return closedExecDrain()
	}
	drained := d.revokeExecAuthorization(authorizationID)
	d.mu.Lock()
	authorizationOwner := d.authorizations[authorizationID]
	if authorizationOwner == nil {
		d.mu.Unlock()
		return drained
	}
	authorizationOwner.stopOnce.Do(func() { close(authorizationOwner.stop) })
	authorizationOwner.cleanupOnce.Do(func() {
		authorizationOwner.authorizationDrained = drained
		go d.finishExecAuthorizationCleanup(authorizationOwner)
	})
	cleanupDone := authorizationOwner.cleanupDone
	d.mu.Unlock()
	return joinExecDrains(drained, cleanupDone)
}

func (d *execAuthorizationService) revokeExecAuthorization(authorizationID string) <-chan struct{} {
	if d.authorizationRevoke != nil {
		return d.authorizationRevoke(authorizationID)
	}
	return d.execAuthorizations.Revoke(authorizationID)
}

func (d *execAuthorizationService) revokeExecAuthorizationsForAgent(serverID, agentID string) {
	if d == nil {
		return
	}
	serverID = strings.TrimSpace(serverID)
	agentID = strings.TrimSpace(agentID)
	d.mu.Lock()
	authorizationIDs := make([]string, 0)
	for authorizationID, owner := range d.authorizations {
		if owner != nil && owner.serverID == serverID && owner.agentID == agentID {
			authorizationIDs = append(authorizationIDs, authorizationID)
		}
	}
	d.mu.Unlock()
	for _, authorizationID := range authorizationIDs {
		d.revokeExecRuntimeAuthorization(authorizationID)
	}
}

func (d *execAuthorizationService) finishExecAuthorizationCleanup(authorizationOwner *execAuthorizationOwner) {
	<-authorizationOwner.done
	attempts := d.closeExecCredentialEventually(authorizationOwner.credential, authorizationOwner.agentID, authorizationOwner.serverID)
	if authorizationOwner.authorizationDrained != nil {
		<-authorizationOwner.authorizationDrained
	}
	if attempts > 1 {
		slog.Debug(
			"agentteam exec credential cleanup recovered",
			"agent_id", authorizationOwner.agentID, "server_id", authorizationOwner.serverID,
			"attempts", attempts,
		)
	}
	d.mu.Lock()
	// Delete the discoverable debt and publish completion under the same mutex.
	// A repeated revoke therefore either takes the mutex first and observes the
	// pending authorization, or takes it after both mutations are complete. Deleting
	// before closing also makes cleanupDone the final no-late-mutation boundary.
	if d.authorizations[authorizationOwner.id] == authorizationOwner {
		delete(d.authorizations, authorizationOwner.id)
	}
	close(authorizationOwner.cleanupDone)
	d.mu.Unlock()
}

func (d *execAuthorizationService) discardUnpublishedExecAuthorization(
	authorizationID string,
	credential execCredential,
	agentID string,
	serverID string,
) {
	authorizationDrained := d.execAuthorizations.Revoke(authorizationID)
	d.closeExecCredentialEventually(credential, agentID, serverID)
	if authorizationDrained != nil {
		<-authorizationDrained
	}
}

func (d *execAuthorizationService) closeExecCredentialEventually(
	credential execCredential,
	agentID string,
	serverID string,
) int {
	attempts := 0
	for {
		attempts++
		if err := credential.Close(); err == nil {
			return attempts
		}
		if attempts == 1 {
			slog.Warn(
				"agentteam exec credential cleanup pending",
				"agent_id", agentID, "server_id", serverID,
			)
		}
		d.waitExecCredentialCleanupRetry(execCredentialCleanupRetryDelay(attempts))
	}
}

func execCredentialCleanupRetryDelay(attempts int) time.Duration {
	if attempts < 1 {
		attempts = 1
	}
	delay := execCredentialCleanupInitialRetryDelay
	for step := 1; step < attempts && delay < execCredentialCleanupMaxRetryDelay; step++ {
		delay *= 2
		if delay > execCredentialCleanupMaxRetryDelay {
			delay = execCredentialCleanupMaxRetryDelay
		}
	}
	return delay
}

func (d *execAuthorizationService) waitExecCredentialCleanupRetry(delay time.Duration) {
	if d.cleanupRetryWait != nil {
		d.cleanupRetryWait(delay)
		return
	}
	// Cleanup debt deliberately survives a bounded shutdown timeout, so this
	// owner has no cancellation signal: it must keep retrying until the private
	// credential tree is gone. Own and stop the timer explicitly rather than
	// parking the goroutine in time.Sleep.
	timer := time.NewTimer(delay)
	defer timer.Stop()
	select {
	case <-timer.C:
	}
}

// revokeAllExecAuthorizations is the final dispatcher shutdown sweep. Callers first
// close Exec admission, so an in-progress launch must either already be in this
// snapshot or reject its final publication and self-clean. Individual process
// teardown may race this sweep; authorization cleanup is idempotent.
func (d *execAuthorizationService) revokeAllExecAuthorizations(ctx context.Context) {
	if d == nil {
		return
	}
	if ctx == nil {
		ctx = context.Background()
	}
	d.mu.Lock()
	authorizationIDs := make([]string, 0, len(d.authorizations))
	for authorizationID := range d.authorizations {
		authorizationIDs = append(authorizationIDs, authorizationID)
	}
	d.mu.Unlock()
	drains := make([]<-chan struct{}, 0, len(authorizationIDs))
	for _, authorizationID := range authorizationIDs {
		drains = append(drains, d.revokeExecRuntimeAuthorization(authorizationID))
	}
	for _, drained := range drains {
		select {
		case <-drained:
		case <-ctx.Done():
			return
		}
	}
}

func closedExecDrain() <-chan struct{} {
	done := make(chan struct{})
	close(done)
	return done
}

func joinExecDrains(first, second <-chan struct{}) <-chan struct{} {
	if first == nil {
		first = closedExecDrain()
	}
	if second == nil {
		second = closedExecDrain()
	}
	done := make(chan struct{})
	go func() {
		<-first
		<-second
		close(done)
	}()
	return done
}

func runtimeAuthorizationSessionID(agentID, effectiveSessionID string) string {
	sessionID := strings.TrimSpace(effectiveSessionID)
	if sessionID == "" {
		sessionID = "pending:" + strings.TrimSpace(agentID)
	}
	return sessionID
}

func execRunID(plc *processLaunchContext, sessionID string) string {
	for _, value := range []string{plc.msg.MessageID, plc.startCfg.InitialDeliveryMessageID, sessionID, plc.sess.agentID} {
		if trimmed := strings.TrimSpace(value); trimmed != "" {
			return trimmed
		}
	}
	return "exec-run"
}

func execEndpoint(baseURL, endpointPath string) string {
	parsed, err := url.Parse(strings.TrimRight(strings.TrimSpace(baseURL), "/"))
	if err != nil || parsed.Scheme == "" || parsed.Host == "" {
		return strings.TrimRight(baseURL, "/") + endpointPath
	}
	escapedPath := strings.TrimRight(parsed.EscapedPath(), "/") + "/" + strings.TrimLeft(endpointPath, "/")
	decodedPath, err := url.PathUnescape(escapedPath)
	if err != nil {
		return strings.TrimRight(baseURL, "/") + endpointPath
	}
	parsed.Path = decodedPath
	parsed.RawPath = escapedPath
	return parsed.String()
}

func clearExecRuntimeEnv(env map[string]string) map[string]string {
	if env == nil {
		return nil
	}
	for key := range env {
		if execruntimeenv.IsDaemonOwnedKey(key) {
			delete(env, key)
		}
	}
	return env
}

func withExecRuntimeEnv(env map[string]string, cfg *drivercontract.ExecConfig) map[string]string {
	out := clearExecRuntimeEnv(cloneStringMapForWrite(env))
	if cfg == nil {
		return out
	}
	out[exec.EnvBaseURL] = cfg.BaseURL
	out[exec.EnvTokenFile] = cfg.TokenFile
	return out
}