package memory

import (
	"context"
	"strings"

	"github.com/tidbcloud/AgentTeam/internal/identity/trusted"
)

// IdentityAgentOwnershipSource reads the AgentTeam-forwarder-owned agent list from
// the trusted JWT. MeshX stays the authority for memory bindings, while AgentTeam
// remains the authority for which private agents the human can control.
type IdentityAgentOwnershipSource struct{}

func (IdentityAgentOwnershipSource) ListOwnedAgents(ctx context.Context, workspaceID, _ string) ([]AgentOwnership, error) {
	identity, ok := trusted.IdentityFromContext(ctx)
	if !ok {
		return nil, forbidden("agent ownership claims are required")
	}
	claims := identity.RawClaims
	if strings.TrimSpace(stringClaim(claims, "agent_ownership_workspace_id")) != strings.TrimSpace(workspaceID) {
		return nil, forbidden("agent ownership claims are required")
	}
	return agentOwnershipsFromClaims(claims), nil
}

func agentOwnershipsFromClaims(claims map[string]any) []AgentOwnership {
	agents := agentOwnershipsFromStructuredClaim(claims["owned_agents"])
	if len(agents) > 0 {
		return agents
	}
	return agentOwnershipsFromIDClaim(claims["owned_agent_ids"])
}

func agentOwnershipsFromStructuredClaim(raw any) []AgentOwnership {
	values, ok := raw.([]any)
	if !ok {
		return nil
	}
	agents := make([]AgentOwnership, 0, len(values))
	seen := make(map[string]struct{}, len(values))
	for _, value := range values {
		item, ok := value.(map[string]any)
		if !ok {
			continue
		}
		id := strings.TrimSpace(stringClaim(item, "id"))
		if id == "" {
			id = strings.TrimSpace(stringClaim(item, "agent_id"))
		}
		if id == "" {
			continue
		}
		if _, exists := seen[id]; exists {
			continue
		}
		seen[id] = struct{}{}
		agents = append(agents, AgentOwnership{
			ID:          id,
			Name:        strings.TrimSpace(stringClaim(item, "name")),
			DisplayName: strings.TrimSpace(firstStringClaim(item, "display_name", "displayName")),
		})
	}
	return agents
}

func agentOwnershipsFromIDClaim(raw any) []AgentOwnership {
	var ids []string
	switch values := raw.(type) {
	case []string:
		ids = values
	case []any:
		ids = make([]string, 0, len(values))
		for _, value := range values {
			if s, ok := value.(string); ok {
				ids = append(ids, s)
			}
		}
	case string:
		ids = strings.Split(values, ",")
	}
	agents := make([]AgentOwnership, 0, len(ids))
	seen := make(map[string]struct{}, len(ids))
	for _, id := range ids {
		id = strings.TrimSpace(id)
		if id == "" {
			continue
		}
		if _, exists := seen[id]; exists {
			continue
		}
		seen[id] = struct{}{}
		agents = append(agents, AgentOwnership{ID: id})
	}
	return agents
}

func stringClaim(claims map[string]any, key string) string {
	if claims == nil {
		return ""
	}
	value, _ := claims[key].(string)
	return value
}

func firstStringClaim(claims map[string]any, keys ...string) string {
	for _, key := range keys {
		if value := stringClaim(claims, key); value != "" {
			return value
		}
	}
	return ""
}