package creds

import (
	"encoding/json"
	"fmt"
	"gopkg.in/yaml.v3"
)

// ---------- Core interfaces & registry ----------

type ProviderCredential interface {
	GetCloudProviderID() string
}

var providerRegistry = map[string]func() ProviderCredential{}

func RegisterProvider(id string, factory func() ProviderCredential) {
	providerRegistry[id] = factory
}

// ---------- Wrapper type (provider-discriminated union) ----------

type CredentialEntry struct {
	// Persisted "provider" key in files
	Provider string `json:"provider" yaml:"provider"`
	// In-memory typed value; excluded from default serialization
	Value ProviderCredential `json:"-" yaml:"-"`
}

// Common decode from a generic map (used by JSON & YAML)
func (c *CredentialEntry) decodeFromMap(m map[string]any) error {
	rawProv, ok := m["provider"]
	if !ok {
		return fmt.Errorf("missing 'provider'")
	}
	provider, ok := rawProv.(string)
	if !ok || provider == "" {
		return fmt.Errorf("invalid 'provider'")
	}
	factory, ok := providerRegistry[provider]
	if !ok {
		return fmt.Errorf("unknown provider: %s", provider)
	}
	cred := factory()

	// Reuse JSON to map→struct unmarshal; unknown keys are ignored by struct
	b, err := json.Marshal(m)
	if err != nil {
		return err
	}
	if err := json.Unmarshal(b, cred); err != nil {
		return err
	}

	c.Provider = provider
	c.Value = cred
	return nil
}

// Common encode to a generic map (used by JSON & YAML)
func (c CredentialEntry) encodeToMap() (map[string]any, error) {
	if c.Value == nil {
		return nil, fmt.Errorf("nil credential Value")
	}
	b, err := json.Marshal(c.Value) // serialize provider-specific fields
	if err != nil {
		return nil, err
	}
	out := map[string]any{}
	if err := json.Unmarshal(b, &out); err != nil {
		return nil, err
	}
	out["provider"] = c.Value.GetCloudProviderID()
	return out, nil
}

// ---------- JSON round-trip ----------

func (c *CredentialEntry) UnmarshalJSON(b []byte) error {
	m := map[string]any{}
	if err := json.Unmarshal(b, &m); err != nil {
		return err
	}
	return c.decodeFromMap(m)
}

func (c CredentialEntry) MarshalJSON() ([]byte, error) {
	m, err := c.encodeToMap()
	if err != nil {
		return nil, err
	}
	return json.Marshal(m)
}

// ---------- YAML round-trip ----------

func (c *CredentialEntry) UnmarshalYAML(n *yaml.Node) error {
	m := map[string]any{}
	if err := n.Decode(&m); err != nil {
		return err
	}
	return c.decodeFromMap(m)
}

func (c CredentialEntry) MarshalYAML() (interface{}, error) {
	m, err := c.encodeToMap()
	if err != nil {
		return nil, err
	}
	return m, nil // let yaml encode the map
}

// ---------- File container ----------

type CredentialFile struct {
	Credentials map[string]CredentialEntry `json:"credentials" yaml:"credentials"`
}

// ---------- Provider implementations this would be in the respective provider packages ----------

// Lambda Labs
type LambdaLabsCredential struct {
	APIKey          string `json:"api_key"           yaml:"api_key"`
	DefaultLocation string `json:"default_location"  yaml:"default_location"`
}
func (l *LambdaLabsCredential) GetCloudProviderID() string { return "lambdalabs" } // use provider id const

// Nebius
type NebiusCredential struct {
	ServiceAccountJSON string `json:"service_account_json" yaml:"service_account_json"`
	DefaultLocation    string `json:"default_location"     yaml:"default_location"`
	RefID              string `json:"ref_id,omitempty"     yaml:"ref_id,omitempty"`
}
func (n *NebiusCredential) GetCloudProviderID() string { return "nebius" } // use provider id const

// Register at init (or from each provider package's init)
func init() {
	RegisterProvider("lambdalabs", func() ProviderCredential { return &LambdaLabsCredential{} }) // use provider id const
	RegisterProvider("nebius",     func() ProviderCredential { return &NebiusCredential{} }) // use provider id const
}