package bitwarden

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"strings"
)

// EncString is the Bitwarden cipher-string wire format:
//
//	<type>.<b64 iv>.<b64 ciphertext>[.<b64 mac>]
//
// Type 2 (AesCbc256_HmacSha256_B64) authenticates with HMAC-SHA256 over
// iv||ciphertext under the MAC half of a 64-byte key; type 0
// (AesCbc256_B64) is unauthenticated AES-256-CBC under a 32-byte key.
// Only the read paths Secrets Manager needs are implemented; anything
// else fails closed with a typed error. Cipher text is never included in
// error messages.
type EncString struct {
	Type byte
	IV   []byte
	CT   []byte
	MAC  []byte
}

// ParseEncString decodes the wire format. ok=false means the value is not
// EncString-shaped at all (callers treat such values as plaintext).
func ParseEncString(s string) (EncString, bool) {
	parts := strings.Split(s, ".")
	if len(parts) < 3 || len(parts[0]) != 1 {
		return EncString{}, false
	}
	t := parts[0][0]
	switch t {
	case '0', '2':
	default:
		return EncString{}, false
	}
	iv, err := b64Decode(parts[1])
	if err != nil || len(iv) != 16 {
		return EncString{}, false
	}
	ct, err := b64Decode(parts[2])
	if err != nil || len(ct) == 0 || len(ct)%aes.BlockSize != 0 {
		return EncString{}, false
	}
	e := EncString{Type: t, IV: iv, CT: ct}
	if t == '2' {
		if len(parts) != 4 {
			return EncString{}, false
		}
		mac, err := b64Decode(parts[3])
		if err != nil || len(mac) != sha256.Size {
			return EncString{}, false
		}
		e.MAC = mac
	} else if len(parts) != 3 {
		return EncString{}, false
	}
	return e, true
}

// String re-encodes as the wire format (used by the fake server in tests).
func (e EncString) String() string {
	var b strings.Builder
	fmt.Fprintf(&b, "%c.%s.%s", e.Type, base64.StdEncoding.EncodeToString(e.IV), base64.StdEncoding.EncodeToString(e.CT))
	if len(e.MAC) > 0 {
		b.WriteString("." + base64.StdEncoding.EncodeToString(e.MAC))
	}
	return b.String()
}

// Decrypt unwraps e under key. It verifies the MAC for type-2 strings
// BEFORE decrypting and fails closed on any integrity error.
func (e EncString) Decrypt(key *SymmetricKey) ([]byte, error) {
	if key == nil {
		return nil, fmt.Errorf("%w: no key available", ErrDecrypt)
	}
	switch e.Type {
	case '2':
		if key.MACKey == nil {
			return nil, fmt.Errorf("%w: mac key required", ErrDecrypt)
		}
		mac := hmac.New(sha256.New, key.MACKey)
		mac.Write(e.IV)
		mac.Write(e.CT)
		if !hmac.Equal(mac.Sum(nil), e.MAC) {
			return nil, fmt.Errorf("%w: integrity check failed", ErrDecrypt)
		}
	case '0':
	default:
		return nil, fmt.Errorf("%w: unsupported type", ErrDecrypt)
	}
	block, err := aes.NewCipher(key.EncKey)
	if err != nil {
		return nil, fmt.Errorf("%w: %s", ErrDecrypt, errBadKey)
	}
	pt := make([]byte, len(e.CT))
	cipher.NewCBCDecrypter(block, e.IV).CryptBlocks(pt, e.CT)
	out, err := pkcs7Unpad(pt)
	if err != nil {
		return nil, fmt.Errorf("%w: bad padding", ErrDecrypt)
	}
	return out, nil
}

// Encrypt seals pt under key with a fresh random IV (used by the fake
// server in tests to produce real Bitwarden-shaped payloads).
func Encrypt(key *SymmetricKey, pt []byte) (EncString, error) {
	if key == nil || len(key.EncKey) != 32 {
		return EncString{}, fmt.Errorf("%w: %s", ErrDecrypt, errBadKey)
	}
	iv := make([]byte, aes.BlockSize)
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return EncString{}, fmt.Errorf("%w: entropy source", ErrDecrypt)
	}
	padded := pkcs7Pad(pt)
	ct := make([]byte, len(padded))
	block, err := aes.NewCipher(key.EncKey)
	if err != nil {
		return EncString{}, fmt.Errorf("%w: %s", ErrDecrypt, errBadKey)
	}
	cipher.NewCBCEncrypter(block, iv).CryptBlocks(ct, padded)
	e := EncString{Type: '0', IV: iv, CT: ct}
	if key.MACKey != nil {
		mac := hmac.New(sha256.New, key.MACKey)
		mac.Write(iv)
		mac.Write(ct)
		e.Type = '2'
		e.MAC = mac.Sum(nil)
	}
	return e, nil
}

// SymmetricKey is a Bitwarden symmetric key: a 32-byte AES key, or a
// 64-byte AES||HMAC key. Keys are held as bytes (not strings) so callers
// can zero them; they are never stringified or logged.
type SymmetricKey struct {
	EncKey []byte
	MACKey []byte
}

// NewSymmetricKey builds a key from raw bytes: 64 bytes -> AES+HMAC,
// 32 bytes -> AES only (Bitwarden legacy org-key lengths).
func NewSymmetricKey(raw []byte) (*SymmetricKey, error) {
	switch len(raw) {
	case 64:
		return &SymmetricKey{EncKey: raw[:32], MACKey: raw[32:]}, nil
	case 32:
		return &SymmetricKey{EncKey: raw}, nil
	default:
		return nil, fmt.Errorf("%w: %s", ErrDecrypt, errBadKey)
	}
}

// Zero wipes the key material in place.
func (k *SymmetricKey) Zero() {
	if k == nil {
		return
	}
	for i := range k.EncKey {
		k.EncKey[i] = 0
	}
	for i := range k.MACKey {
		k.MACKey[i] = 0
	}
}

// b64Decode accepts standard base64 with or without padding (the Bitwarden
// servers emit padded values; their own parsers accept both).
func b64Decode(s string) ([]byte, error) {
	if b, err := base64.StdEncoding.DecodeString(s); err == nil {
		return b, nil
	}
	return base64.RawStdEncoding.DecodeString(s)
}

const errBadKey = "invalid key length"

func pkcs7Pad(pt []byte) []byte {
	n := aes.BlockSize - len(pt)%aes.BlockSize
	out := make([]byte, len(pt)+n)
	copy(out, pt)
	for i := len(pt); i < len(out); i++ {
		out[i] = byte(n)
	}
	return out
}

func pkcs7Unpad(pt []byte) ([]byte, error) {
	if len(pt) == 0 || len(pt)%aes.BlockSize != 0 {
		return nil, errors.New("not block aligned")
	}
	n := int(pt[len(pt)-1])
	if n == 0 || n > aes.BlockSize || n > len(pt) {
		return nil, errors.New("invalid padding")
	}
	for _, c := range pt[len(pt)-n:] {
		if int(c) != n {
			return nil, errors.New("invalid padding")
		}
	}
	return pt[:len(pt)-n], nil
}