package stores

import (
	"fmt"
	"math"
	"sort"
	"strconv"
	"strings"

	"github.com/getsops/sops/v3"
)

const mapSeparator = "__map_"
const listSeparator = "__list_"

////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Unflatten

type token interface{}

type mapToken struct {
	key string
}

type listToken struct {
	position int
}

// tokenize converts a path generated by Flatten to be used as a key
// in the flattened map, and converts it to a slice of tokens
func tokenize(path string) []token {
	const (
		StateNormal = 0
		StateMap    = iota
		StateList   = iota
	)
	var tokens []token
	state := StateNormal
	lastTokenEnd := 0
	i := 0
	finishPrevToken := func() {
		var t token
		switch state {
		case StateNormal:
			t = mapToken{path[lastTokenEnd:i]}
		case StateMap:
			t = mapToken{path[lastTokenEnd+len(mapSeparator) : i]}
		case StateList:
			pos, _ := strconv.Atoi(path[lastTokenEnd+len(listSeparator) : i])
			t = listToken{pos}
		}
		lastTokenEnd = i
		tokens = append(tokens, t)
	}
	for i < len(path) {
		if strings.HasPrefix(path[i:], mapSeparator) {
			finishPrevToken()
			state = StateMap
			i += len(mapSeparator)
		} else if strings.HasPrefix(path[i:], listSeparator) {
			finishPrevToken()
			state = StateList
			i += len(listSeparator)
		} else {
			i++
		}
	}
	finishPrevToken()
	return tokens
}

type node struct {
	value   *interface{}
	subkeys map[string]*node
	indices map[int]*node
}

func place(value *interface{}, path []token, root *node) error {
	for _, token := range path {
		var next *node
		var ok bool
		switch t := token.(type) {
		case mapToken:
			if root.subkeys == nil {
				root.subkeys = make(map[string]*node)
			}
			next, ok = root.subkeys[t.key]
			if !ok {
				next = &node{}
				root.subkeys[t.key] = next
			}
		case listToken:
			if root.indices == nil {
				root.indices = make(map[int]*node)
			}
			next, ok = root.indices[t.position]
			if !ok {
				next = &node{}
				root.indices[t.position] = next
			}
		}
		root = next
	}
	if root.value != nil {
		return fmt.Errorf("Duplicate value")
	}
	root.value = value
	return nil
}

func b2i(value bool) int {
	if value {
		return 1
	}
	return 0
}

func convert(root *node) (interface{}, error) {
	hasValue := root.value != nil
	hasSubkey := len(root.subkeys) > 0
	hasIndex := len(root.indices) > 0
	if b2i(hasValue)+b2i(hasSubkey)+b2i(hasIndex) > 1 {
		return nil, fmt.Errorf("Type mismatch")
	}
	if !hasValue && !hasSubkey && !hasIndex {
		return sops.TreeBranch{}, nil
	}
	if hasValue {
		return *root.value, nil
	}
	if hasSubkey {
		keys := make([]string, len(root.subkeys))
		index := 0
		for k := range root.subkeys {
			keys[index] = k
			index += 1
		}
		sort.Strings(keys)
		result := make(sops.TreeBranch, len(keys))
		for index, key := range keys {
			value, err := convert(root.subkeys[key])
			if err != nil {
				return nil, err
			}
			result[index] = sops.TreeItem{
				Key:   key,
				Value: value,
			}
		}
		return result, nil
	}
	minValue := math.MaxInt
	maxValue := math.MinInt
	for k := range root.indices {
		if k < minValue {
			minValue = k
		}
		if k > maxValue {
			maxValue = k
		}
	}
	if minValue != 0 || maxValue+1 != len(root.indices) {
		return nil, fmt.Errorf("Incomplete list")
	}
	result := make([]interface{}, maxValue+1)
	for k, v := range root.indices {
		value, err := convert(v)
		if err != nil {
			return nil, err
		}
		result[k] = value
	}
	return result, nil
}

func unflattenTreeBranch(branch sops.TreeBranch) (sops.TreeBranch, error) {
	root := &node{}
	for _, item := range branch {
		if _, ok := item.Key.(sops.Comment); ok {
			continue
		}
		if key, ok := item.Key.(string); ok {
			tokens := tokenize(key)
			err := place(&item.Value, tokens, root)
			if err != nil {
				return nil, fmt.Errorf("Error while unflattening %q: %w", key, err)
			}
		} else {
			return nil, fmt.Errorf("Found non-string key %q when unflattening", item.Key)
		}
	}
	result, err := convert(root)
	if err != nil {
		return nil, fmt.Errorf("Error while unflattening: %w", err)
	}
	if tb, ok := result.(sops.TreeBranch); ok {
		return tb, nil
	}
	return nil, fmt.Errorf("Internal error: cannot find root")
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Flatten

func flattenDescendValue(value interface{}, key string, destination sops.TreeBranch, destinationMap *map[string]bool) (sops.TreeBranch, error) {
	switch value := value.(type) {
	case sops.TreeBranch:
		return flattenDescendMap(value, key+mapSeparator, destination, destinationMap)
	case []interface{}:
		return flattenDescendArray(value, key+listSeparator, destination, destinationMap)
	}
	if _, ok := (*destinationMap)[key]; ok {
		return nil, fmt.Errorf("Found key collision %q while flattening", key)
	}
	destination = append(destination, sops.TreeItem{
		Key:   key,
		Value: value,
	})
	(*destinationMap)[key] = true
	return destination, nil
}

func flattenDescendMap(branch sops.TreeBranch, prefix string, destination sops.TreeBranch, destinationMap *map[string]bool) (sops.TreeBranch, error) {
	for _, item := range branch {
		if _, ok := item.Key.(sops.Comment); ok {
			continue
		}
		if key, ok := item.Key.(string); ok {
			var err error
			destination, err = flattenDescendValue(item.Value, prefix+key, destination, destinationMap)
			if err != nil {
				return nil, err
			}
		} else {
			return nil, fmt.Errorf("Found non-string key %q when flattening", item.Key)
		}
	}
	return destination, nil
}

func flattenDescendArray(array []interface{}, prefix string, destination sops.TreeBranch, destinationMap *map[string]bool) (sops.TreeBranch, error) {
	i := 0
	for _, item := range array {
		if _, ok := item.(sops.Comment); ok {
			continue
		}
		var err error
		destination, err = flattenDescendValue(item, fmt.Sprintf("%s%d", prefix, i), destination, destinationMap)
		if err != nil {
			return nil, err
		}
		i++
	}
	return destination, nil
}

func flattenTreeBranch(branch sops.TreeBranch, prefix string) (sops.TreeBranch, error) {
	destinationMap := map[string]bool{}
	return flattenDescendMap(branch, prefix, nil, &destinationMap)
}
