package crypto

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"

	"github.com/ProtonMail/go-crypto/openpgp"
	"github.com/ProtonMail/go-crypto/openpgp/packet"
	"github.com/klauspost/compress/zstd"
)

// ContainerMagic identifies a logarchiver object and its container version.
var ContainerMagic = []byte("LARC1\n")

const (
	dekSize         = 32 // AES-256
	noncePrefixSize = 4
	counterSize     = 8
	// maxFrameCiphertext bounds a single frame read to avoid unbounded allocation
	// from a corrupt/hostile length prefix.
	maxFrameCiphertext = 128 << 20
)

// Header is the LARC1 object header (JSON), written after the magic and a
// uint32 big-endian length prefix.
type Header struct {
	Version        int    `json:"v"`
	KeyName        string `json:"key_name"`
	KeyFingerprint string `json:"key_fingerprint"`
	WrappedDEKLen  int    `json:"wrapped_dek_len"`
	NoncePrefix    []byte `json:"nonce_prefix"` // base64 in JSON
	FrameSize      int    `json:"frame_size"`   // plaintext (compressed) bytes per frame
	Compression    string `json:"compression"`  // "zstd"
	Cipher         string `json:"cipher"`       // "AES-256-GCM"
}

// SealResult reports what Seal produced (for the index row).
type SealResult struct {
	Header      Header
	RawBytes    int64 // input NDJSON length
	StoredBytes int64 // full container length
}

// Seal compresses plaintext with zstd, encrypts it under a fresh random DEK
// using framed AES-256-GCM, wraps the DEK to pub with OpenPGP, and writes the
// LARC1 container to w. keyName is recorded in the header for operator context.
func Seal(w io.Writer, plaintext []byte, pub *PublicKey, keyName string, frameSize int) (SealResult, error) {
	if pub == nil {
		return SealResult{}, fmt.Errorf("nil public key")
	}
	if frameSize <= 0 {
		frameSize = 1 << 20
	}

	// 1. Compress.
	enc, err := zstd.NewWriter(nil, zstd.WithEncoderLevel(zstd.SpeedBetterCompression))
	if err != nil {
		return SealResult{}, fmt.Errorf("zstd writer: %w", err)
	}
	compressed := enc.EncodeAll(plaintext, nil)
	_ = enc.Close()

	// 2. DEK + nonce prefix.
	dek := make([]byte, dekSize)
	if _, err := rand.Read(dek); err != nil {
		return SealResult{}, fmt.Errorf("gen dek: %w", err)
	}
	noncePrefix := make([]byte, noncePrefixSize)
	if _, err := rand.Read(noncePrefix); err != nil {
		return SealResult{}, fmt.Errorf("gen nonce prefix: %w", err)
	}

	// 3. Wrap the DEK to the public key (small standard OpenPGP message).
	wrapped, err := wrapDEK(dek, pub)
	if err != nil {
		return SealResult{}, err
	}

	block, err := aes.NewCipher(dek)
	if err != nil {
		return SealResult{}, fmt.Errorf("aes cipher: %w", err)
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return SealResult{}, fmt.Errorf("gcm: %w", err)
	}

	hdr := Header{
		Version:        1,
		KeyName:        keyName,
		KeyFingerprint: pub.Fingerprint,
		WrappedDEKLen:  len(wrapped),
		NoncePrefix:    noncePrefix,
		FrameSize:      frameSize,
		Compression:    "zstd",
		Cipher:         "AES-256-GCM",
	}
	hdrJSON, err := json.Marshal(hdr)
	if err != nil {
		return SealResult{}, fmt.Errorf("marshal header: %w", err)
	}

	cw := &countingWriter{w: w}

	// magic
	if _, err := cw.Write(ContainerMagic); err != nil {
		return SealResult{}, err
	}
	// header length + header
	if err := writeUint32(cw, uint32(len(hdrJSON))); err != nil {
		return SealResult{}, err
	}
	if _, err := cw.Write(hdrJSON); err != nil {
		return SealResult{}, err
	}
	// wrapped DEK
	if _, err := cw.Write(wrapped); err != nil {
		return SealResult{}, err
	}

	// frames
	for i, off := 0, 0; off < len(compressed); i++ {
		end := off + frameSize
		if end > len(compressed) {
			end = len(compressed)
		}
		nonce := frameNonce(noncePrefix, uint64(i))
		aad := aadFor(uint64(i))
		ct := gcm.Seal(nil, nonce, compressed[off:end], aad)
		if err := writeUint32(cw, uint32(len(ct))); err != nil {
			return SealResult{}, err
		}
		if _, err := cw.Write(ct); err != nil {
			return SealResult{}, err
		}
		off = end
	}
	// terminating zero-length frame
	if err := writeUint32(cw, 0); err != nil {
		return SealResult{}, err
	}

	return SealResult{Header: hdr, RawBytes: int64(len(plaintext)), StoredBytes: cw.n}, nil
}

// wrapDEK OpenPGP-encrypts the DEK to pub, producing a compact binary message.
func wrapDEK(dek []byte, pub *PublicKey) ([]byte, error) {
	var buf bytes.Buffer
	cfg := &packet.Config{
		DefaultCipher: packet.CipherAES256,
	}
	wc, err := openpgp.Encrypt(&buf, []*openpgp.Entity{pub.entity}, nil, nil, cfg)
	if err != nil {
		return nil, fmt.Errorf("openpgp encrypt dek: %w", err)
	}
	if _, err := wc.Write(dek); err != nil {
		return nil, fmt.Errorf("write dek: %w", err)
	}
	if err := wc.Close(); err != nil {
		return nil, fmt.Errorf("close openpgp: %w", err)
	}
	return buf.Bytes(), nil
}

// UnwrapFunc recovers the DEK from the wrapped OpenPGP blob. In production this
// calls the Vault GPG engine decrypt endpoint; tests supply a local one.
type UnwrapFunc func(wrappedDEK []byte) (dek []byte, err error)

// ReadHeader reads and validates the LARC1 magic + header and the wrapped DEK,
// leaving r positioned at the first frame. It does not require decryption keys,
// so it is cheap for `search`/metadata inspection.
func ReadHeader(r io.Reader) (Header, []byte, error) {
	magic := make([]byte, len(ContainerMagic))
	if _, err := io.ReadFull(r, magic); err != nil {
		return Header{}, nil, fmt.Errorf("read magic: %w", err)
	}
	if !bytes.Equal(magic, ContainerMagic) {
		return Header{}, nil, fmt.Errorf("bad magic: not a logarchiver (LARC1) object")
	}
	hlen, err := readUint32(r)
	if err != nil {
		return Header{}, nil, fmt.Errorf("read header len: %w", err)
	}
	if hlen == 0 || hlen > 1<<20 {
		return Header{}, nil, fmt.Errorf("implausible header length %d", hlen)
	}
	hdrJSON := make([]byte, hlen)
	if _, err := io.ReadFull(r, hdrJSON); err != nil {
		return Header{}, nil, fmt.Errorf("read header: %w", err)
	}
	var hdr Header
	if err := json.Unmarshal(hdrJSON, &hdr); err != nil {
		return Header{}, nil, fmt.Errorf("parse header: %w", err)
	}
	if hdr.Version != 1 {
		return Header{}, nil, fmt.Errorf("unsupported container version %d", hdr.Version)
	}
	if hdr.WrappedDEKLen <= 0 || hdr.WrappedDEKLen > 1<<20 {
		return Header{}, nil, fmt.Errorf("implausible wrapped dek length %d", hdr.WrappedDEKLen)
	}
	wrapped := make([]byte, hdr.WrappedDEKLen)
	if _, err := io.ReadFull(r, wrapped); err != nil {
		return Header{}, nil, fmt.Errorf("read wrapped dek: %w", err)
	}
	return hdr, wrapped, nil
}

// Open reads a LARC1 container from r, recovers the DEK via unwrap, and streams
// the decrypted, decompressed NDJSON to w.
func Open(r io.Reader, w io.Writer, unwrap UnwrapFunc) error {
	hdr, wrapped, err := ReadHeader(r)
	if err != nil {
		return err
	}
	dek, err := unwrap(wrapped)
	if err != nil {
		return fmt.Errorf("unwrap dek: %w", err)
	}
	if len(dek) != dekSize {
		return fmt.Errorf("unwrapped dek has wrong length %d", len(dek))
	}
	block, err := aes.NewCipher(dek)
	if err != nil {
		return fmt.Errorf("aes cipher: %w", err)
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return fmt.Errorf("gcm: %w", err)
	}

	fr := &frameReader{r: r, gcm: gcm, noncePrefix: hdr.NoncePrefix}
	zr, err := zstd.NewReader(fr)
	if err != nil {
		return fmt.Errorf("zstd reader: %w", err)
	}
	defer zr.Close()
	if _, err := io.Copy(w, zr); err != nil {
		return fmt.Errorf("decompress: %w", err)
	}
	return nil
}

// frameReader decrypts LARC1 frames on demand, presenting the decrypted
// (compressed) bytes as an io.Reader for the zstd decoder.
type frameReader struct {
	r           io.Reader
	gcm         cipher.AEAD
	noncePrefix []byte
	counter     uint64
	buf         []byte // leftover decrypted plaintext not yet consumed
	done        bool
}

func (f *frameReader) Read(p []byte) (int, error) {
	if len(f.buf) == 0 && !f.done {
		if err := f.next(); err != nil {
			return 0, err
		}
	}
	if len(f.buf) == 0 {
		return 0, io.EOF
	}
	n := copy(p, f.buf)
	f.buf = f.buf[n:]
	return n, nil
}

func (f *frameReader) next() error {
	ln, err := readUint32(f.r)
	if err != nil {
		return fmt.Errorf("read frame len: %w", err)
	}
	if ln == 0 { // terminator
		f.done = true
		return nil
	}
	if ln > maxFrameCiphertext {
		return fmt.Errorf("frame length %d exceeds max", ln)
	}
	ct := make([]byte, ln)
	if _, err := io.ReadFull(f.r, ct); err != nil {
		return fmt.Errorf("read frame: %w", err)
	}
	nonce := frameNonce(f.noncePrefix, f.counter)
	aad := aadFor(f.counter)
	pt, err := f.gcm.Open(nil, nonce, ct, aad)
	if err != nil {
		return fmt.Errorf("decrypt frame %d: %w", f.counter, err)
	}
	f.counter++
	f.buf = pt
	return nil
}

func frameNonce(prefix []byte, counter uint64) []byte {
	nonce := make([]byte, noncePrefixSize+counterSize)
	copy(nonce, prefix)
	binary.BigEndian.PutUint64(nonce[noncePrefixSize:], counter)
	return nonce
}

func aadFor(counter uint64) []byte {
	aad := make([]byte, counterSize)
	binary.BigEndian.PutUint64(aad, counter)
	return aad
}

func writeUint32(w io.Writer, v uint32) error {
	var b [4]byte
	binary.BigEndian.PutUint32(b[:], v)
	_, err := w.Write(b[:])
	return err
}

func readUint32(r io.Reader) (uint32, error) {
	var b [4]byte
	if _, err := io.ReadFull(r, b[:]); err != nil {
		return 0, err
	}
	return binary.BigEndian.Uint32(b[:]), nil
}

type countingWriter struct {
	w io.Writer
	n int64
}

func (c *countingWriter) Write(p []byte) (int, error) {
	n, err := c.w.Write(p)
	c.n += int64(n)
	return n, err
}