package main

import (
	"encoding/base64"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/hashicorp/vault/api"
)

// vaultClient wraps the Vault API client and turns a password-store gpg-id (a
// "<mount>/<key>" reference into the vault-plugin-secrets-gpg engine) into
// encrypt/decrypt calls. The private key never leaves Vault; this shim only
// ever ships base64 to and from the engine.
type vaultClient struct {
	api *api.Client
}

func newVault() (*vaultClient, error) {
	cfg := api.DefaultConfig()
	if err := cfg.ReadEnvironment(); err != nil {
		return nil, fmt.Errorf("reading Vault environment: %w", err)
	}
	c, err := api.NewClient(cfg)
	if err != nil {
		return nil, fmt.Errorf("creating Vault client: %w", err)
	}
	if c.Token() == "" {
		if tok := resolveToken(); tok != "" {
			c.SetToken(tok)
		}
	}
	if c.Token() == "" {
		return nil, fmt.Errorf("no Vault token (set VAULT_TOKEN or run `vault login`)")
	}
	return &vaultClient{api: c}, nil
}

// resolveToken falls back to the on-disk token that `vault login` writes.
func resolveToken() string {
	if t := os.Getenv("VAULT_TOKEN"); t != "" {
		return t
	}
	home, err := os.UserHomeDir()
	if err != nil {
		return ""
	}
	b, err := os.ReadFile(filepath.Join(home, ".vault-token"))
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(b))
}

// splitRef parses "<mount>/<key>" (e.g. "gpg/app") into its parts.
func splitRef(ref string) (mount, key string, err error) {
	ref = strings.Trim(strings.TrimSpace(ref), "/")
	i := strings.LastIndex(ref, "/")
	if i <= 0 || i == len(ref)-1 {
		return "", "", fmt.Errorf("invalid key reference %q (want <mount>/<key>, e.g. gpg/app)", ref)
	}
	return ref[:i], ref[i+1:], nil
}

func (v *vaultClient) encrypt(ref string, plaintext []byte) ([]byte, error) {
	mount, key, err := splitRef(ref)
	if err != nil {
		return nil, err
	}
	resp, err := v.api.Logical().Write(mount+"/encrypt/"+key, map[string]interface{}{
		"plaintext": base64.StdEncoding.EncodeToString(plaintext),
		"format":    "base64",
	})
	if err != nil {
		return nil, fmt.Errorf("vault encrypt (%s): %w", ref, err)
	}
	return fieldBytes(resp, "ciphertext")
}

func (v *vaultClient) decrypt(ref string, ciphertext []byte) ([]byte, error) {
	mount, key, err := splitRef(ref)
	if err != nil {
		return nil, err
	}
	resp, err := v.api.Logical().Write(mount+"/decrypt/"+key, map[string]interface{}{
		// The engine auto-detects armored vs raw binary; we send base64 binary.
		"ciphertext": base64.StdEncoding.EncodeToString(ciphertext),
	})
	if err != nil {
		return nil, fmt.Errorf("vault decrypt (%s): %w", ref, err)
	}
	return fieldBytes(resp, "plaintext")
}

// publicKey returns the armored public key for a reference, used by `init` to
// verify the key exists and is reachable before writing the store.
func (v *vaultClient) publicKey(ref string) (string, error) {
	mount, key, err := splitRef(ref)
	if err != nil {
		return "", err
	}
	resp, err := v.api.Logical().Read(mount + "/keys/" + key)
	if err != nil {
		return "", fmt.Errorf("vault read key (%s): %w", ref, err)
	}
	if resp == nil || resp.Data["public_key"] == nil {
		return "", fmt.Errorf("key %q not found in Vault", ref)
	}
	return resp.Data["public_key"].(string), nil
}

// fieldBytes pulls a base64 string field out of a Vault response and decodes it.
func fieldBytes(resp *api.Secret, field string) ([]byte, error) {
	if resp == nil || resp.Data[field] == nil {
		return nil, fmt.Errorf("vault response missing %q", field)
	}
	s, ok := resp.Data[field].(string)
	if !ok {
		return nil, fmt.Errorf("vault field %q is not a string", field)
	}
	return base64.StdEncoding.DecodeString(s)
}