package utils

import (
	"bytes"
	"encoding/json"
	"reflect"
)

func DictKeysMan(data interface{}, keysToRemove []string, keysToKeep []string) interface{} {
	if data == nil || (keysToKeep == nil && keysToRemove == nil) {
		return data
	}

	var keysToRemoveMap map[string]int
	var keysToKeepMap map[string]int
	if keysToKeep != nil {
		keysToKeepMap = make(map[string]int)
		for _, v := range keysToKeep {
			keysToKeepMap[v] = 1
		}
	} else if keysToRemove != nil {
		keysToRemoveMap = make(map[string]int)
		for _, v := range keysToRemove {
			keysToRemoveMap[v] = 1
		}
	}

	dataIsSlice := true
	valueOfData := reflect.ValueOf(data)
	if valueOfData.Kind() != reflect.Slice {
		dataIsSlice = false
		valueOfData = reflect.ValueOf([]interface{}{data})
	} else if valueOfData.Len() == 0 {
		return data
	}

	retVal := make([]interface{}, valueOfData.Len())

	for index := 0; index < valueOfData.Len(); index++ {
		realV := valueOfData.Index(index)
		if realV.Kind() == reflect.Interface || realV.Kind() == reflect.Ptr {
			realV = realV.Elem()
		}
		if keysToRemoveMap != nil || keysToKeepMap != nil {
			mapV := make(map[string]interface{})
			for _, key := range realV.MapKeys() {
				fieldName := key.String()
				wantThisField := true
				if keysToKeepMap != nil {
					if _, ok := keysToKeepMap[fieldName]; !ok {
						wantThisField = false
					}
				} else if keysToRemoveMap != nil {
					if _, ok := keysToRemoveMap[fieldName]; ok {
						wantThisField = false
					}
				}
				if wantThisField {
					mapV[fieldName] = realV.MapIndex(key)
				}
			}
			retVal[index] = mapV
		} else {
			retVal[index] = realV.Interface()
		}
	}

	if !dataIsSlice {
		return retVal[0]
	}
	return retVal
}

func UnmarshalUseNumber(data []byte, result interface{}) error {
	d := json.NewDecoder(bytes.NewReader(data))
	d.UseNumber()
	return d.Decode(result)
}