package factory

import (
	"fmt"

	"github.com/jfraeys/fetch_ml/internal/container"
	"github.com/jfraeys/fetch_ml/internal/logging"
	"github.com/jfraeys/fetch_ml/internal/tracking"
	"github.com/jfraeys/fetch_ml/internal/tracking/plugins"
)

// PluginConfig represents the configuration for a single plugin.
type PluginConfig struct {
	Enabled      bool           `toml:"enabled" yaml:"enabled"`
	Image        string         `toml:"image" yaml:"image"`
	Mode         string         `toml:"mode" yaml:"mode"`
	LogBasePath  string         `toml:"log_base_path" yaml:"log_base_path"`
	ArtifactPath string         `toml:"artifact_path" yaml:"artifact_path"`
	Settings     map[string]any `toml:"settings" yaml:"settings"`
}

// PluginFactory is a function that creates a Plugin instance.
type PluginFactory func(
	logger *logging.Logger,
	podman *container.PodmanManager,
	cfg PluginConfig,
) (tracking.Plugin, error)

// PluginLoader uses dependency injection to load plugins.
type PluginLoader struct {
	logger    *logging.Logger
	podman    *container.PodmanManager
	factories map[string]PluginFactory
}

// NewPluginLoader creates a new PluginLoader.
func NewPluginLoader(logger *logging.Logger, podman *container.PodmanManager) *PluginLoader {
	loader := &PluginLoader{
		logger:    logger,
		podman:    podman,
		factories: make(map[string]PluginFactory),
	}

	// Register default factories
	loader.RegisterFactory("mlflow", createMLflowPlugin)
	loader.RegisterFactory("tensorboard", createTensorBoardPlugin)
	loader.RegisterFactory("wandb", createWandbPlugin)

	return loader
}

// RegisterFactory allows external packages to register new plugin types (Marketplace ready).
func (l *PluginLoader) RegisterFactory(name string, factory PluginFactory) {
	l.factories[name] = factory
}

// LoadPlugins loads plugins from the provided configuration map and registers them.
func (l *PluginLoader) LoadPlugins(
	plugins map[string]PluginConfig,
	registry *tracking.Registry,
) error {
	for name, pluginCfg := range plugins {
		if !pluginCfg.Enabled {
			continue
		}

		factory, ok := l.factories[name]
		if !ok {
			l.logger.Warn("unknown plugin type", "name", name)
			continue
		}

		plugin, err := factory(l.logger, l.podman, pluginCfg)
		if err != nil {
			return fmt.Errorf("failed to create plugin %s: %w", name, err)
		}

		registry.Register(plugin)
		l.logger.Info("plugin loaded", "name", name, "mode", pluginCfg.Mode)
	}

	return nil
}

// Factory Implementations
func createMLflowPlugin(
	logger *logging.Logger,
	podman *container.PodmanManager,
	cfg PluginConfig,
) (tracking.Plugin, error) {
	opts := plugins.MLflowOptions{
		Image:            cfg.Image,
		ArtifactBasePath: cfg.ArtifactPath,
	}
	return plugins.NewMLflowPlugin(logger, podman, opts)
}

func createTensorBoardPlugin(
	logger *logging.Logger,
	podman *container.PodmanManager,
	cfg PluginConfig,
) (tracking.Plugin, error) {
	opts := plugins.TensorBoardOptions{
		Image:       cfg.Image,
		LogBasePath: cfg.LogBasePath,
	}
	return plugins.NewTensorBoardPlugin(logger, podman, opts)
}

func createWandbPlugin(
	logger *logging.Logger,
	_ *container.PodmanManager,
	_ PluginConfig,
) (tracking.Plugin, error) {
	return plugins.NewWandbPlugin(), nil
}