package store

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"io"
	"regexp"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ssm"
)

const SSM_MAX_SIZE = 4096

// ((16^4)*4096)/1024/1024
// If we ever need more than 256 MB in parameter store, we've done something
// very wrong.
const SSM_KEY_FORMAT = "%s-%04X" //

var (
	// TrimRegex is used to group the SSM keys inside of the Info struct under
	// ByKey. This will only be used for params that exceed 4KB.
	TrimRegex = regexp.MustCompile("-[0-9A-E][0-9A-E][0-9A-E][0-9A-E]$")

	// Optional, can be set set to utilize a specific KMS key if desired.
	KMS_KEY_ID *string = nil

	Tags = []*ssm.Tag{}
)

type Info struct {
	ByKey     map[string]*Entry
	ByFullKey map[string]*Entry
}

func (i *Info) init() {
	if i.ByKey == nil {
		i.ByKey = map[string]*Entry{}
	}
	if i.ByFullKey == nil {
		i.ByFullKey = map[string]*Entry{}
	}
}

func (i *Info) add(e *ssm.ParameterMetadata) {
	name := TrimRegex.ReplaceAllString(*e.Name, "")
	var entry *Entry

	// Doesn't exist, make a new entry
	if _, ok := i.ByKey[name]; !ok {
		entry = &Entry{
			Name: name,
			Keys: []*ssm.ParameterMetadata{e},
		}

		i.ByKey[name] = entry
		i.ByFullKey[*e.Name] = entry
		return
	}

	// Otherwise let's just update the one that's there
	entry = i.ByKey[name]
	entry.Keys = append(entry.Keys, e)
	i.ByFullKey[*e.Name] = entry
}

type Entry struct {
	Name string
	Keys []*ssm.ParameterMetadata
}

// GetInfo returns a populated Info struct from the SSM
func GetInfo(svc *ssm.SSM) (*Info, error) {
	ret := &Info{
		ByKey:     map[string]*Entry{},
		ByFullKey: map[string]*Entry{},
	}

	var out = &ssm.DescribeParametersOutput{}
	var err error
	for {
		out, err = svc.DescribeParameters(&ssm.DescribeParametersInput{
			NextToken: out.NextToken,
		})
		if err != nil {
			return nil, err
		}
		for _, entry := range out.Parameters {
			ret.add(entry)
		}
		if out.NextToken == nil {
			break
		}
	}

	return ret, nil
}

func InsertParam(svc *ssm.SSM, rdr io.Reader, key string) error {
	buf := &bytes.Buffer{}
	enc := base64.NewEncoder(base64.StdEncoding, buf)

	_, err := io.Copy(enc, rdr)
	if err != nil {
		return fmt.Errorf("Error while reading stdin: %w", err)
	}

	err = enc.Close()
	if err != nil {
		return fmt.Errorf("While closing encoder: %w", err)
	}

	n := 1
	for {
		key := fmt.Sprintf(SSM_KEY_FORMAT, key, n)
		n++

		paramBuf := &bytes.Buffer{}

		n, err := io.CopyN(paramBuf, buf, SSM_MAX_SIZE)
		if err != io.EOF && err != nil {
			return fmt.Errorf("While writing output: %w", err)
		}
		// fmt.Fprintf(os.Stderr, "Wrote '%d' bytes to '%s'\n", n, key)

		_, err = svc.PutParameter(&ssm.PutParameterInput{
			KeyId: KMS_KEY_ID,
			Name:  aws.String(key),
			Type:  aws.String("SecureString"),
			Value: aws.String(paramBuf.String()),
			Tags:  Tags,
		})

		if err != nil {
			return fmt.Errorf("Failed putting prameter: %w", err)
		}

		if err == io.EOF || n < SSM_MAX_SIZE {
			break
		}
	}

	return nil
}

func GetParam(svc *ssm.SSM, wrtr io.Writer, key string) error {
	n := 1
	buf := &bytes.Buffer{}

	for {
		key := fmt.Sprintf(SSM_KEY_FORMAT, key, n)
		n++
		out, err := svc.GetParameter(&ssm.GetParameterInput{
			Name:           aws.String(key),
			WithDecryption: aws.Bool(true),
		})
		if err != nil {
			return fmt.Errorf("While fetching: '%s': %w", key, err)
		}

		w, err := buf.WriteString(*out.Parameter.Value)
		if err != nil {
			return fmt.Errorf("While writing to buffer: %w", err)
		}

		if w != SSM_MAX_SIZE {
			break
		}
	}

	dec := base64.NewDecoder(base64.StdEncoding, buf)
	_, err := io.Copy(wrtr, dec)
	if err != nil {
		return fmt.Errorf("Error writing output: %w", err)
	}

	return nil
}

func RemoveParam(svc *ssm.SSM, key string) error {
	info, err := GetInfo(svc)
	if err != nil {
		return fmt.Errorf("When fetching info: %w", err)
	}

	entry, ok := info.ByKey[key]
	if !ok {
		return fmt.Errorf("Entry '%s' not found in parameter store", key)
	}

	var e error
	for _, key := range entry.Keys {
		_, err := svc.DeleteParameter(&ssm.DeleteParameterInput{
			Name: key.Name,
		})
		if err != nil {
			if e == nil {
				e = err
			} else {
				e = fmt.Errorf("%s, %w", err, e)
			}
		}
	}

	return e
}