Cyrene_v2/backend/internal/plugin/registry.go

package plugin

import (
	"context"
	"fmt"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
	"sync"
)

// Registry manages all loaded plugins and provides type-safe access to providers.
type Registry struct {
	mu           sync.RWMutex
	plugins      map[string]Plugin
	iotProviders []IoTProvider
	toolProviders []ToolProvider
	ragProviders []RAGProvider
}

// NewRegistry creates a new empty plugin registry.
func NewRegistry() *Registry {
	return &Registry{
		plugins:      make(map[string]Plugin),
		iotProviders: make([]IoTProvider, 0),
		toolProviders: make([]ToolProvider, 0),
		ragProviders: make([]RAGProvider, 0),
	}
}

// Register adds a plugin to the registry after calling Init.
// If Init fails or the plugin is disabled, it is not added.
func (r *Registry) Register(p Plugin, cfg PluginConfig) error {
	ctx := context.Background()

	if err := p.Init(ctx, cfg); err != nil {
		return fmt.Errorf("plugin %s init failed: %w", p.Name(), err)
	}

	if !p.Enabled() {
		logger.Infof("[plugin] %s is disabled, skipping", p.Name())
		return nil
	}

	r.mu.Lock()
	defer r.mu.Unlock()

	// Check for duplicate
	if _, exists := r.plugins[p.Name()]; exists {
		return fmt.Errorf("plugin %s already registered", p.Name())
	}

	r.plugins[p.Name()] = p

	// Auto-detect provider interfaces
	if iot, ok := p.(IoTProvider); ok {
		r.iotProviders = append(r.iotProviders, iot)
		logger.Infof("[plugin] registered IoT provider: %s v%s", p.Name(), p.Version())
	}
	if tool, ok := p.(ToolProvider); ok {
		r.toolProviders = append(r.toolProviders, tool)
		logger.Infof("[plugin] registered tool provider: %s v%s (%d tools)", p.Name(), p.Version(), len(tool.Tools()))
	}
	if rag, ok := p.(RAGProvider); ok {
		r.ragProviders = append(r.ragProviders, rag)
		logger.Infof("[plugin] registered RAG provider: %s v%s", p.Name(), p.Version())
	}

	logger.Infof("[plugin] %s v%s loaded successfully", p.Name(), p.Version())
	return nil
}

// Get returns a plugin by name, or nil if not found.
func (r *Registry) Get(name string) Plugin {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return r.plugins[name]
}

// ListNames returns the names of all loaded plugins.
func (r *Registry) ListNames() []string {
	r.mu.RLock()
	defer r.mu.RUnlock()
	names := make([]string, 0, len(r.plugins))
	for name := range r.plugins {
		names = append(names, name)
	}
	return names
}

// GetIoTProvider returns the first registered IoT provider, or nil.
// When multiple IoT providers are registered, the first one wins.
func (r *Registry) GetIoTProvider() IoTProvider {
	r.mu.RLock()
	defer r.mu.RUnlock()
	if len(r.iotProviders) > 0 {
		return r.iotProviders[0]
	}
	return nil
}

// GetAllToolProviders returns all registered tool providers.
func (r *Registry) GetAllToolProviders() []ToolProvider {
	r.mu.RLock()
	defer r.mu.RUnlock()
	result := make([]ToolProvider, len(r.toolProviders))
	copy(result, r.toolProviders)
	return result
}

// GetAllToolDefinitions returns a flat list of all tool definitions from all tool providers.
func (r *Registry) GetAllToolDefinitions() []ToolDefinition {
	r.mu.RLock()
	defer r.mu.RUnlock()
	defs := make([]ToolDefinition, 0)
	for _, tp := range r.toolProviders {
		defs = append(defs, tp.Tools()...)
	}
	return defs
}

// ExecuteTool finds a tool by name across all tool providers and executes it.
func (r *Registry) ExecuteTool(ctx context.Context, toolName string, params map[string]any) (string, error) {
	r.mu.RLock()
	defer r.mu.RUnlock()

	for _, tp := range r.toolProviders {
		for _, def := range tp.Tools() {
			if def.Name == toolName {
				return tp.Execute(ctx, toolName, params)
			}
		}
	}
	return "", fmt.Errorf("tool %s not found in any provider", toolName)
}

// GetRAGProviders returns all registered RAG providers.
func (r *Registry) GetRAGProviders() []RAGProvider {
	r.mu.RLock()
	defer r.mu.RUnlock()
	result := make([]RAGProvider, len(r.ragProviders))
	copy(result, r.ragProviders)
	return result
}

// Count returns the total number of loaded plugins.
func (r *Registry) Count() int {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return len(r.plugins)
}

// Shutdown calls Close on all plugins in reverse registration order.
func (r *Registry) Shutdown() {
	r.mu.Lock()
	defer r.mu.Unlock()

	// Close base plugins that are not specialized providers
	for _, p := range r.plugins {
		// Check if it was already closed via a provider list
		if _, ok := p.(IoTProvider); ok {
			continue
		}
		if _, ok := p.(ToolProvider); ok {
			continue
		}
		if _, ok := p.(RAGProvider); ok {
			continue
		}
		if err := p.Close(); err != nil {
			logger.Infof("[plugin] error closing %s: %v", p.Name(), err)
		}
	}

	// Close IoT providers
	for _, p := range r.iotProviders {
		if err := p.Close(); err != nil {
			logger.Infof("[plugin] error closing %s: %v", p.Name(), err)
		}
	}
	// Close tool providers
	for _, p := range r.toolProviders {
		if err := p.Close(); err != nil {
			logger.Infof("[plugin] error closing %s: %v", p.Name(), err)
		}
	}
	// Close RAG providers
	for _, p := range r.ragProviders {
		if err := p.Close(); err != nil {
			logger.Infof("[plugin] error closing %s: %v", p.Name(), err)
		}
	}
	logger.Infof("[plugin] all plugins shut down")
}