package dispatcher

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"strings"

	"github.com/tidbcloud/AgentTeam/desktop/internal/api/exec"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
	"github.com/tidbcloud/AgentTeam/desktop/internal/wire/apiclient"
)

type extensionCatalogKey struct {
	serverID string
	agentID  string
}

type extensionResolved struct {
	api      *apiclient.Client
	snapshot exec.ExtensionSnapshot
}

// ObserveAgentExtensionCatalogHints records the latest server revision. The
// revision is only a change notification: content is fetched separately and a
// runtime is replaced only when the authorization-projected response ETag
// actually changes. Server-side invocation authorization is always rechecked,
// so the hint-to-fetch window remains fail-closed without eagerly disrupting
// the local process.
func (d *extensionRuntimeService) ObserveAgentExtensionCatalogHints(
	serverID string,
	agents []types.MachineAgentInfo,
) {
	if d == nil {
		return
	}
	serverID = strings.TrimSpace(serverID)
	if serverID == "" {
		return
	}
	observed := make(map[string]types.ExtensionCatalogHint, len(agents))
	for _, agent := range agents {
		agentID := strings.TrimSpace(agent.ID)
		if agentID == "" {
			continue
		}
		hint := types.ExtensionCatalogHint{}
		if agent.ExtensionCatalog != nil {
			hint = *agent.ExtensionCatalog
		}
		if !validExtensionCatalogHint(hint) {
			hint = types.ExtensionCatalogHint{}
		}
		observed[agentID] = hint
	}

	var removed []extensionCatalogKey
	d.catalogMu.Lock()
	for agentID, hint := range observed {
		key := extensionCatalogKey{serverID: serverID, agentID: agentID}
		previous, existed := d.desired[key]
		if existed && extensionCatalogHintsEqual(previous, hint) {
			continue
		}
		d.desired[key] = hint
	}
	for key := range d.desired {
		if key.serverID != serverID {
			continue
		}
		if _, exists := observed[key.agentID]; exists {
			continue
		}
		delete(d.desired, key)
		delete(d.published, key)
		delete(d.snapshots, key)
		removed = append(removed, key)
	}
	d.catalogMu.Unlock()

	for _, key := range removed {
		d.revokeExecAuthorizationsForAgent(key.serverID, key.agentID)
	}
}

// RefreshAgentExtensionCatalog fetches and verifies one desired snapshot. It is
// safe to run repeatedly; an already-published matching snapshot is a no-op.
func (d *extensionRuntimeService) RefreshAgentExtensionCatalog(
	ctx context.Context,
	serverID string,
	agentID string,
	api *apiclient.Client,
) error {
	if d == nil {
		return errors.New("refresh Extension Catalog: dispatcher unavailable")
	}
	key := extensionCatalogKey{
		serverID: strings.TrimSpace(serverID),
		agentID:  strings.TrimSpace(agentID),
	}
	hint, cached, publishedRevision, hasPublished, ok := d.extensionCatalogState(key)
	if !ok {
		return nil
	}
	if hasPublished && publishedRevision == hint.Version {
		return nil
	}
	snapshot, err := d.fetchExtensionCatalog(ctx, key, cached, api)
	if err != nil {
		return err
	}
	accepted, changed := d.publishExtensionCatalog(key, hint, snapshot)
	if !accepted {
		return errors.New("refresh Extension Catalog: Catalog changed during fetch")
	}
	if changed {
		d.revokeExecAuthorizationsForAgent(key.serverID, key.agentID)
		d.RestartByAgent(key.agentID)
	}
	return nil
}

// prepareExtension resolves the Extension material for one launch. It is the
// hard guarantee that an Extension problem never aborts agent startup: any
// failure downgrades this launch to built-in Exec actions only, with one
// operator-visible Warn carrying the cause. The background refresh
// (RefreshAgentExtensionCatalog) republishes and restarts the agent once the
// Catalog becomes fetchable again.
func (d *extensionRuntimeService) prepareExtension(plc *processLaunchContext) (*extensionResolved, error) {
	if d == nil || plc == nil || plc.sess == nil {
		return nil, errors.New("prepare Extension: launch context is required")
	}
	key := extensionCatalogKey{
		serverID: strings.TrimSpace(plc.binding.channelID),
		agentID:  strings.TrimSpace(plc.sess.agentID),
	}
	hint, snapshot, publishedRevision, hasPublished, ok := d.extensionCatalogState(key)
	if !ok || hint.Version == 0 {
		return nil, nil
	}
	if hasPublished && publishedRevision == hint.Version && snapshot.Empty() {
		return nil, nil
	}
	resolved, err := d.resolveExtensionForLaunch(
		plc,
		key,
		hint,
		snapshot,
		publishedRevision,
		hasPublished,
	)
	if err != nil {
		slog.Warn(
			"AgentTeam: agent starting without Extensions",
			"server", key.serverID,
			"agent_id", key.agentID,
			"error", err,
		)
		return nil, nil
	}
	return resolved, nil
}

func (d *extensionRuntimeService) resolveExtensionForLaunch(
	plc *processLaunchContext,
	key extensionCatalogKey,
	hint types.ExtensionCatalogHint,
	snapshot exec.ExtensionSnapshot,
	publishedRevision int64,
	hasPublished bool,
) (*extensionResolved, error) {
	ch, connected := d.serverConnections.ServerConnection(key.serverID)
	if !connected {
		return nil, errors.New("prepare Extension: AgentTeam server connection unavailable")
	}
	provider, ok := ch.(agentteamAPIClientProvider)
	if !ok {
		return nil, errors.New("prepare Extension: AgentTeam API client unavailable")
	}
	api := provider.APIClient(plc.binding.serverAgentID)
	if api == nil {
		return nil, errors.New("prepare Extension: AgentTeam API client unavailable")
	}
	api.LocalAgentID = key.agentID
	if !hasPublished || publishedRevision != hint.Version {
		var err error
		snapshot, err = d.fetchExtensionCatalog(plc.requestContext, key, snapshot, api)
		if err != nil {
			return nil, err
		}
		accepted, changed := d.publishExtensionCatalog(key, hint, snapshot)
		if !accepted {
			return nil, errors.New("prepare Extension: Catalog changed during fetch")
		}
		if changed {
			d.revokeExecAuthorizationsForAgent(key.serverID, key.agentID)
		}
	}
	if snapshot.Empty() {
		return nil, nil
	}
	return &extensionResolved{api: api, snapshot: snapshot}, nil
}

func (d *extensionRuntimeService) fetchExtensionCatalog(
	ctx context.Context,
	key extensionCatalogKey,
	cached exec.ExtensionSnapshot,
	api *apiclient.Client,
) (exec.ExtensionSnapshot, error) {
	if api == nil {
		return exec.ExtensionSnapshot{}, errors.New("fetch Extension Catalog: API client unavailable")
	}
	if ctx == nil {
		ctx = context.Background()
	}
	snapshot, diagnostics, _, err := exec.FetchExtensionSnapshot(ctx, api, cached)
	if len(diagnostics) > 0 {
		attrs := []any{
			"server", key.serverID,
			"agent_id", key.agentID,
		}
		for _, diagnostic := range diagnostics {
			slog.Warn(
				"AgentTeam: Extension unit skipped in Catalog fetch",
				append(attrs, "detail", diagnostic)...,
			)
		}
	}
	if err != nil {
		return exec.ExtensionSnapshot{}, fmt.Errorf(
			"fetch Extension Catalog for agent %q on server %q: %w",
			key.agentID,
			key.serverID,
			err,
		)
	}
	return snapshot, nil
}

func (d *extensionRuntimeService) extensionCatalogState(
	key extensionCatalogKey,
) (
	types.ExtensionCatalogHint,
	exec.ExtensionSnapshot,
	int64,
	bool,
	bool,
) {
	d.catalogMu.RLock()
	defer d.catalogMu.RUnlock()
	hint, ok := d.desired[key]
	publishedRevision, hasPublished := d.published[key]
	return hint, d.snapshots[key], publishedRevision, hasPublished, ok
}

func (d *extensionRuntimeService) publishExtensionCatalog(
	key extensionCatalogKey,
	hint types.ExtensionCatalogHint,
	snapshot exec.ExtensionSnapshot,
) (accepted bool, changed bool) {
	d.catalogMu.Lock()
	defer d.catalogMu.Unlock()
	current, ok := d.desired[key]
	if !ok || !extensionCatalogHintsEqual(current, hint) {
		return false, false
	}
	previous, existed := d.snapshots[key]
	changed = (!existed && !snapshot.Empty()) ||
		(existed && previous.ETag() != snapshot.ETag())
	d.snapshots[key] = snapshot
	d.published[key] = hint.Version
	return true, changed
}

func validExtensionCatalogHint(hint types.ExtensionCatalogHint) bool {
	return hint.Version >= 0
}

func extensionCatalogHintsEqual(left, right types.ExtensionCatalogHint) bool {
	return left.Version == right.Version
}