package util

import (
	"context"

	"gitea.libretechconsulting.com/rmcguire/ubiquiti-clients/cmd/internal/config"
	"gitea.libretechconsulting.com/rmcguire/ubiquiti-clients/pkg/edgeos"
	"gitea.libretechconsulting.com/rmcguire/ubiquiti-clients/pkg/toughswitch"
	"github.com/rs/zerolog"
)

type ctxKey uint8

const (
	ctxConfig ctxKey = iota
	ctxToughSwitchClient
	ctxEdgeOSClient
)

func ContextWithConfig(baseCtx context.Context, config *config.ClientsConfig) context.Context {
	return context.WithValue(baseCtx, ctxConfig, config)
}

func ConfigFromContext(ctx context.Context) *config.ClientsConfig {
	val := ctx.Value(ctxConfig)
	conf, ok := val.(*config.ClientsConfig)
	if !ok {
		return nil
	}

	return conf
}

// ContextWithLogger stores the logger in context using zerolog's built-in method.
func ContextWithLogger(baseCtx context.Context, logger *zerolog.Logger) context.Context {
	return logger.WithContext(baseCtx)
}

// LoggerFromContext retrieves the logger from context using zerolog's built-in method.
func LoggerFromContext(ctx context.Context) *zerolog.Logger {
	return zerolog.Ctx(ctx)
}

// ContextWithToughSwitchClient stores a toughswitch client in the context.
func ContextWithToughSwitchClient(baseCtx context.Context, client *toughswitch.Client) context.Context {
	return context.WithValue(baseCtx, ctxToughSwitchClient, client)
}

// ToughSwitchClientFromContext retrieves the toughswitch client from context.
func ToughSwitchClientFromContext(ctx context.Context) *toughswitch.Client {
	val := ctx.Value(ctxToughSwitchClient)
	client, ok := val.(*toughswitch.Client)
	if !ok {
		return nil
	}
	return client
}

// ContextWithEdgeOSClient stores an edgeos client in the context.
func ContextWithEdgeOSClient(baseCtx context.Context, client *edgeos.Client) context.Context {
	return context.WithValue(baseCtx, ctxEdgeOSClient, client)
}

// EdgeOSClientFromContext retrieves the edgeos client from context.
func EdgeOSClientFromContext(ctx context.Context) *edgeos.Client {
	val := ctx.Value(ctxEdgeOSClient)
	client, ok := val.(*edgeos.Client)
	if !ok {
		return nil
	}
	return client
}