package suite

import (
	"fmt"

	"github.com/tidbcloud/eks-provider/tests/e2e/dag"
	"github.com/tidbcloud/eks-provider/tests/e2e/label"
)

type EnvMap interface {
	EnvNames() []string
	GetEnv(name string) label.Object[dag.Env]
	AddEnv(env label.Object[dag.Env]) error

	GetCases(name string) []label.Object[dag.Case]
	AddCase(c label.Object[dag.Case]) error

	GetEnvParent(name string, afterFilter bool) []string

	FilterEnvParent(name string, parent ...string)
}

type envMap struct {
	envs []string
	em   map[string]*optionalEnv
	cm   map[string][]label.Object[dag.Case]
}

type optionalEnv struct {
	e label.Object[dag.Env]

	filtered []string
}

func NewEnvMap() EnvMap {
	return &envMap{
		em: map[string]*optionalEnv{},
		cm: map[string][]label.Object[dag.Case]{},
	}
}

func (m *envMap) EnvNames() []string {
	return m.envs
}

func (m *envMap) GetEnv(name string) label.Object[dag.Env] {
	oe, ok := m.em[name]
	if !ok {
		return nil
	}

	return oe.e
}

func (m *envMap) AddEnv(env label.Object[dag.Env]) error {
	name := env.Get().Name()
	if _, ok := m.em[name]; ok {
		return fmt.Errorf("env %s has been registered", name)
	}

	m.em[name] = &optionalEnv{
		e: env,
	}

	m.envs = append(m.envs, name)

	return nil
}

func (m *envMap) AddCase(c label.Object[dag.Case]) error {
	for _, env := range c.Get().Envs() {
		if _, ok := m.em[env]; !ok {
			return fmt.Errorf("env %s for case(%s) is not registered", env, c.GetDesc())
		}
		cases := m.cm[env]
		nc := c.DeepCopy()
		cases = append(cases, nc)
		m.cm[env] = cases
	}

	return nil
}

func (m *envMap) GetCases(name string) []label.Object[dag.Case] {
	cs, ok := m.cm[name]
	if !ok {
		return nil
	}
	return cs
}

func (m *envMap) GetEnvParent(name string, afterFilter bool) []string {
	oe, ok := m.em[name]
	if !ok {
		return nil
	}

	inner := oe.e.Get()
	var parents []string
	parents = append(parents, inner.Envs()...)
	if es, ok := inner.(dag.WithOptionalEnvs); ok {
		parents = append(parents, es.OptionalEnvs()...)
	}

	if afterFilter {
		var filtered []string
		for _, p := range parents {
			found := false
			for _, fp := range oe.filtered {
				if p == fp {
					found = true
				}
			}

			if !found {
				filtered = append(filtered, p)
			}
		}

		return filtered
	}

	return parents
}

func (m *envMap) FilterEnvParent(name string, parent ...string) {
	oe, ok := m.em[name]
	if !ok {
		return
	}

	for _, p := range parent {
		found := false
		for _, ep := range oe.filtered {
			if p == ep {
				found = true
				break
			}
		}

		if !found {
			oe.filtered = append(oe.filtered, p)
		}
	}
}