// MarshalWithMetadata takes a sops tree branch and sops metadata and marshals them to json.
func (store Store) MarshalWithMetadata(tree sops.TreeBranch, metadata sops.Metadata) ([]byte, error) {
	tree = append(tree, sops.TreeItem{Key: "sops", Value: metadata.ToMap()})
	out, err := store.jsonFromTreeBranch(tree)
	if err != nil {
		return nil, fmt.Errorf("Error marshaling to json: %s", err)
	}
	return out, nil
}

// UnmarshalMetadata takes a yaml document as a string and extracts sops' metadata from it
func (store *Store) UnmarshalMetadata(in []byte) (sops.Metadata, error) {
	var metadata sops.Metadata
	var ok bool
	data := make(map[interface{}]interface{})
	err := yaml.Unmarshal(in, &data)
	if err != nil {
		return metadata, fmt.Errorf("Error unmarshalling input yaml: %s", err)
	}
	if data, ok = data["sops"].(map[interface{}]interface{}); !ok {
		return metadata, sops.MetadataNotFound
	}
	metadata.MessageAuthenticationCode = data["mac"].(string)
	lastModified, err := time.Parse(time.RFC3339, data["lastmodified"].(string))
	if err != nil {
		return metadata, fmt.Errorf("Could not parse last modified date: %s", err)
	}
	metadata.LastModified = lastModified
	unencryptedSuffix, ok := data["unencrypted_suffix"].(string)
	if !ok {
		unencryptedSuffix = sops.DefaultUnencryptedSuffix
	}
	metadata.UnencryptedSuffix = unencryptedSuffix
	if metadata.Version, ok = data["version"].(string); !ok {
		metadata.Version = strconv.FormatFloat(data["version"].(float64), 'f', -1, 64)
	}
	if k, ok := data["kms"].([]interface{}); ok {
		ks, err := store.kmsEntries(k)
		if err == nil {
			metadata.KeySources = append(metadata.KeySources, ks)
		}

	}

	if pgp, ok := data["pgp"].([]interface{}); ok {
		ks, err := store.pgpEntries(pgp)
		if err == nil {
			metadata.KeySources = append(metadata.KeySources, ks)
		}
	}
	return metadata, nil
}

func encrypt(c *cli.Context, file string, fileBytes []byte, output io.Writer) error {
	store := store(file)
	branch, err := store.Unmarshal(fileBytes)
	if err != nil {
		return cli.NewExitError(fmt.Sprintf("Error loading file: %s", err), exitCouldNotReadInputFile)
	}
	var metadata sops.Metadata
	metadata.UnencryptedSuffix = c.String("unencrypted-suffix")
	metadata.Version = "2.0.0"
	var kmsKeys []sops.MasterKey
	var pgpKeys []sops.MasterKey

	if c.String("kms") != "" {
		for _, k := range kms.MasterKeysFromArnString(c.String("kms")) {
			kmsKeys = append(kmsKeys, &k)
		}
	}
	if c.String("pgp") != "" {
		for _, k := range pgp.MasterKeysFromFingerprintString(c.String("pgp")) {
			pgpKeys = append(pgpKeys, &k)
		}
	}

	if c.String("kms") == "" && c.String("pgp") == "" {
		var confBytes []byte
		if c.String("config") != "" {
			confBytes, err = ioutil.ReadFile(c.String("config"))
			if err != nil {
				return cli.NewExitError(fmt.Sprintf("Error loading config file: %s", err), exitErrorReadingConfig)
			}
		}
		kmsString, pgpString, err := yaml.MasterKeyStringsForFile(file, confBytes)
		if err == nil {
			for _, k := range pgp.MasterKeysFromFingerprintString(pgpString) {
				pgpKeys = append(pgpKeys, &k)
			}
			for _, k := range kms.MasterKeysFromArnString(kmsString) {
				kmsKeys = append(kmsKeys, &k)
			}
		}
	}
	kmsKs := sops.KeySource{Name: "kms", Keys: kmsKeys}
	pgpKs := sops.KeySource{Name: "pgp", Keys: pgpKeys}
	metadata.KeySources = append(metadata.KeySources, kmsKs)
	metadata.KeySources = append(metadata.KeySources, pgpKs)
	tree := sops.Tree{Branch: branch, Metadata: metadata}
	key, err := tree.GenerateDataKey()
	if err != nil {
		return cli.NewExitError(err.Error(), exitCouldNotRetrieveKey)
	}
	cipher := aes.Cipher{}
	mac, err := tree.Encrypt(key, cipher)
	encryptedMac, err := cipher.Encrypt(mac, key, []byte(metadata.LastModified.Format(time.RFC3339)))
	if err != nil {
		return cli.NewExitError(fmt.Sprintf("Could not encrypt MAC: %s", err), exitErrorEncryptingTree)
	}
	metadata.MessageAuthenticationCode = encryptedMac
	out, err := store.MarshalWithMetadata(tree.Branch, metadata)
	_, err = output.Write([]byte(out))
	if err != nil {
		return cli.NewExitError(fmt.Sprintf("Could not write to output stream: %s", err), exitCouldNotWriteOutputFile)
	}
	return nil
}