package store

import (
	"context"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"database/sql"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
)

// At-rest encryption for the sqlite store.
//
// The on-disk file at cfg.DBPath is ALWAYS ciphertext: a fixed magic header +
// a random 12-byte GCM nonce + AES-256-GCM ciphertext of the whole sqlite
// file. On Open the ciphertext is decrypted into a tmpfs (RAM-backed)
// plaintext working copy that modernc sqlite operates on directly; on Close
// the working copy is checkpointed, re-encrypted, and atomically written back,
// then wiped. A crash leaves plaintext only in RAM (gone on reboot) and never
// a half-written ciphertext file (atomic rename).
//
// Wrong key or a tampered file => GCM authentication fails => Open fails
// closed. We NEVER fall back to opening the file as plaintext.
//
// ponytail: chosen over option 2 (github.com/ncruces/go-sqlite3 pure-Go
// page-level encryption VFS = no transient plaintext file). Rejected for now
// because the threat model here is disk-at-rest only — the tmpfs plaintext is
// RAM that clears on reboot, acceptable — and option 2 swaps the sqlite driver
// project-wide. Revisit if the threat model grows to include RAM capture.
//
// KEY: OpenEncrypted takes a raw 32-byte key. Today the operator supplies it
// as base64 in config/env (see config.DBEncryptionKey). No KDF: x/crypto isn't
// a dependency and stdlib has no argon2/scrypt, so a passphrase-derived key
// would only get sha256 — worse than demanding real key material. This same
// []byte seam is where the L3 passkey-derived cold-start key plugs in later
// (auth/tier.go Layer3): produce the 32 bytes from the passkey op and pass them
// here instead of reading config. ponytail: if a passphrase path is ever
// wanted, store a random salt in the header below and derive with argon2id
// (new dep) — do NOT bolt on sha256(passphrase).

const (
	cryptMagic    = "MVNC1\x00" // 6-byte file magic; version in the trailing byte
	cryptMagicLen = len(cryptMagic)
	nonceLen      = 12 // AES-GCM standard nonce
	keyLen        = 32 // AES-256
)

// ErrKeyLen — the supplied encryption key was not exactly 32 bytes.
var ErrKeyLen = errors.New("store: encryption key must be 32 bytes")

// ErrDecrypt — the ciphertext failed to authenticate/decrypt (wrong key or
// tampering). Fail closed: the caller gets no handle and no plaintext.
var ErrDecrypt = errors.New("store: decrypt failed (wrong key or corrupt file)")

type encState struct {
	cipherPath string // on-disk ciphertext (cfg.DBPath)
	plainPath  string // tmpfs working copy
	key        []byte // 32 bytes; zeroed on Close
}

// OpenEncrypted opens the encrypted store whose ciphertext lives at cipherPath.
// The plaintext working copy is created at tmpfsPath (default under /dev/shm if
// empty). key must be exactly 32 bytes.
//
// Flow:
//   - ciphertext file exists (our magic)  -> decrypt to tmpfs, open.
//   - file exists but is legacy plaintext -> first-run upgrade: copy to tmpfs,
//     open; Close will write ciphertext over the original (rename) so the
//     plaintext original is gone after the first clean shutdown.
//   - no file                             -> fresh empty db in tmpfs.
func OpenEncrypted(ctx context.Context, cipherPath, tmpfsPath string, key []byte) (*Store, error) {
	if len(key) != keyLen {
		return nil, ErrKeyLen
	}
	if tmpfsPath == "" {
		tmpfsPath = defaultTmpfsPath(cipherPath)
	}

	// Stale working copy from a previous crash: tmpfs is RAM so this only
	// happens without a reboot in between. Remove it; the ciphertext file is
	// the source of truth.
	removeDBFiles(tmpfsPath)

	raw, err := os.ReadFile(cipherPath)
	switch {
	case err == nil && isCiphertext(raw):
		plain, derr := decrypt(key, raw)
		if derr != nil {
			return nil, derr // ErrDecrypt — fail closed
		}
		if err := writeFileSync(tmpfsPath, plain, 0o600); err != nil {
			return nil, fmt.Errorf("write working copy: %w", err)
		}
	case err == nil:
		// Legacy plaintext sqlite db: first-run upgrade. Copy verbatim to tmpfs
		// and let Close encrypt it back over the original.
		if err := writeFileSync(tmpfsPath, raw, 0o600); err != nil {
			return nil, fmt.Errorf("stage plaintext for upgrade: %w", err)
		}
	case errors.Is(err, os.ErrNotExist):
		// fresh: nothing to stage; sqlite creates tmpfsPath on open.
	default:
		return nil, fmt.Errorf("read ciphertext %s: %w", cipherPath, err)
	}

	db, err := openAt(ctx, tmpfsPath)
	if err != nil {
		removeDBFiles(tmpfsPath)
		return nil, err
	}
	return &Store{db: db, enc: &encState{cipherPath: cipherPath, plainPath: tmpfsPath, key: key}}, nil
}

// closeAndSeal checkpoints the WAL into the main file, closes the handle,
// re-encrypts the working copy to the ciphertext file atomically, then wipes
// the plaintext copy and zeroes the key.
func (e *encState) closeAndSeal(db *sql.DB) error {
	// Fold -wal/-shm into the main file so we encrypt a single complete db.
	// Best-effort: a checkpoint failure still lets us encrypt what's committed.
	_, _ = db.ExecContext(context.Background(), "PRAGMA wal_checkpoint(TRUNCATE)")
	if err := db.Close(); err != nil {
		return fmt.Errorf("close db: %w", err)
	}

	plain, err := os.ReadFile(e.plainPath)
	if err != nil {
		return fmt.Errorf("read working copy: %w", err)
	}
	blob, err := encrypt(e.key, plain)
	if err != nil {
		return fmt.Errorf("encrypt: %w", err)
	}
	if err := atomicWrite(e.cipherPath, blob); err != nil {
		return fmt.Errorf("seal ciphertext: %w", err)
	}

	removeDBFiles(e.plainPath)
	zero(e.key)
	return nil
}

func isCiphertext(b []byte) bool {
	return len(b) >= cryptMagicLen && string(b[:cryptMagicLen]) == cryptMagic
}

// encrypt: magic || nonce || AES-256-GCM(seal). The magic is bound as
// additional data so a truncated/retagged header also fails authentication.
func encrypt(key, plain []byte) ([]byte, error) {
	gcm, err := newGCM(key)
	if err != nil {
		return nil, err
	}
	nonce := make([]byte, nonceLen)
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, fmt.Errorf("nonce: %w", err)
	}
	out := make([]byte, 0, cryptMagicLen+nonceLen+len(plain)+gcm.Overhead())
	out = append(out, cryptMagic...)
	out = append(out, nonce...)
	out = gcm.Seal(out, nonce, plain, []byte(cryptMagic))
	return out, nil
}

func decrypt(key, blob []byte) ([]byte, error) {
	if len(blob) < cryptMagicLen+nonceLen {
		return nil, ErrDecrypt
	}
	gcm, err := newGCM(key)
	if err != nil {
		return nil, err
	}
	nonce := blob[cryptMagicLen : cryptMagicLen+nonceLen]
	ct := blob[cryptMagicLen+nonceLen:]
	plain, err := gcm.Open(nil, nonce, ct, []byte(cryptMagic))
	if err != nil {
		return nil, ErrDecrypt // fail closed, don't leak the GCM error detail
	}
	return plain, nil
}

func newGCM(key []byte) (cipher.AEAD, error) {
	if len(key) != keyLen {
		return nil, ErrKeyLen
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("aes: %w", err)
	}
	return cipher.NewGCM(block)
}

// defaultTmpfsPath: a stable per-db path under /dev/shm so restarts reuse it.
// Hash of the ciphertext path keeps distinct dbs from colliding.
func defaultTmpfsPath(cipherPath string) string {
	sum := sha256.Sum256([]byte(cipherPath))
	name := fmt.Sprintf("maven-%s.db", encodeHex(sum[:6]))
	return filepath.Join("/dev/shm", name)
}

func encodeHex(b []byte) string {
	const hexd = "0123456789abcdef"
	out := make([]byte, len(b)*2)
	for i, c := range b {
		out[i*2] = hexd[c>>4]
		out[i*2+1] = hexd[c&0xf]
	}
	return string(out)
}

// atomicWrite writes to a temp file in the same dir, fsyncs, and renames over
// the target so the ciphertext file is never observed half-written.
func atomicWrite(path string, data []byte) error {
	dir := filepath.Dir(path)
	if err := os.MkdirAll(dir, 0o700); err != nil {
		return err
	}
	tmp, err := os.CreateTemp(dir, ".maven-seal-*")
	if err != nil {
		return err
	}
	tmpName := tmp.Name()
	if _, err := tmp.Write(data); err != nil {
		tmp.Close()
		os.Remove(tmpName)
		return err
	}
	if err := tmp.Sync(); err != nil {
		tmp.Close()
		os.Remove(tmpName)
		return err
	}
	if err := tmp.Close(); err != nil {
		os.Remove(tmpName)
		return err
	}
	if err := os.Chmod(tmpName, 0o600); err != nil {
		os.Remove(tmpName)
		return err
	}
	if err := os.Rename(tmpName, path); err != nil {
		os.Remove(tmpName)
		return err
	}
	return nil
}

func writeFileSync(path string, data []byte, perm os.FileMode) error {
	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
		return err
	}
	return os.WriteFile(path, data, perm)
}

// removeDBFiles removes the db and its -wal/-shm sidecars, best-effort.
func removeDBFiles(path string) {
	for _, p := range []string{path, path + "-wal", path + "-shm"} {
		_ = os.Remove(p)
	}
}

func zero(b []byte) {
	for i := range b {
		b[i] = 0
	}
}

// SealPlaintext encrypts an existing plaintext sqlite file at plainPath and
// writes the ciphertext to cipherPath, atomically. key must be 32 bytes. The
// plaintext file is left alone: this is a recovery path, and deleting the only
// good copy of the data on the strength of a write that just succeeded is not
// a trade worth making here.
//
// It exists for the case closeAndSeal cannot cover: a daemon that was killed
// rather than shut down, leaving a live working copy in tmpfs and a stale
// ciphertext on disk. mavseal folds the WAL in first, so what arrives here is
// a single complete database.
//
// Nothing else should call this. The normal path is Close, which seals and
// then wipes the plaintext and the key.
func SealPlaintext(plainPath, cipherPath string, key []byte) error {
	if len(key) != keyLen {
		return ErrKeyLen
	}
	plain, err := os.ReadFile(plainPath)
	if err != nil {
		return fmt.Errorf("read working copy: %w", err)
	}
	blob, err := encrypt(key, plain)
	if err != nil {
		return fmt.Errorf("encrypt: %w", err)
	}
	if err := atomicWrite(cipherPath, blob); err != nil {
		return fmt.Errorf("seal ciphertext: %w", err)
	}
	return nil
}