package archive

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdh"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"fmt"
	"io"

	"github.com/carlosframework/platform/internal/appconfig"
)

// ChunkSize is the AES-GCM chunk (§8.2). Each chunk is a 4-byte big-endian
// length and the ciphertext; its nonce is its index, and its associated
// data flags whether it is the last. Every chunk but the last holds exactly
// ChunkSize bytes of plaintext and the last holds fewer (possibly none), so
// the framing of a plaintext under a key is unique, Open knows from a
// chunk's length which flag it must carry, and a stream cut at a chunk
// boundary never ends on a chunk flagged last.
const ChunkSize = 1 << 20

// maxChunks bounds a sealed object: the largest decoded archive, plus
// zstd's worst-case expansion of incompressible input (under 1/256, and
// a margin for frame headers), in chunks, plus the final chunk.
const maxChunks = uint64((MaxDecoded+MaxDecoded/256+64<<10)/ChunkSize + 2)

// ErrSealed is every refusal to open: a data key that is unbound or sealed
// for another app, a chunk that does not authenticate, a stream cut short
// or carrying data after its final chunk.
var ErrSealed = errors.New("archive: sealed object does not open")

func nonceFor(index uint64, last bool) ([]byte, []byte) {
	nonce := make([]byte, 12)
	binary.BigEndian.PutUint64(nonce[4:], index)
	aad := []byte{0}
	if last {
		aad[0] = 1
	}
	return nonce, aad
}

// Seal encrypts plain into w under a fresh 256-bit data key and returns
// the data key sealed with appconfig.SealBound (a v2 envelope bound to b,
// with its own fresh ephemeral key and nonce) and the sha256 of what it
// wrote. Never deterministic. SealBound, not appconfig.Seal: Seal writes
// unbound v1 envelopes unless the process turned bound sealing on, and the
// host agent never does (F4).
func Seal(plain io.Reader, pub *ecdh.PublicKey, b appconfig.Binding, w io.Writer) (sealedDataKey, sealedSHA string, err error) {
	key := make([]byte, 32)
	if _, err := rand.Read(key); err != nil {
		return "", "", err
	}
	defer clear(key)
	gcm, err := newGCM(key)
	if err != nil {
		return "", "", err
	}
	h := sha256.New()
	out := io.MultiWriter(w, h)
	buf := make([]byte, ChunkSize)
	for index := uint64(0); ; index++ {
		if index >= maxChunks {
			return "", "", fmt.Errorf("%w: more than %d chunks", ErrTooLarge, maxChunks)
		}
		n, rerr := io.ReadFull(plain, buf)
		last := rerr == io.EOF || rerr == io.ErrUnexpectedEOF
		if rerr != nil && !last {
			return "", "", rerr
		}
		nonce, aad := nonceFor(index, last)
		ct := gcm.Seal(nil, nonce, buf[:n], aad)
		var hdr [4]byte
		binary.BigEndian.PutUint32(hdr[:], uint32(len(ct)))
		if _, err := out.Write(hdr[:]); err != nil {
			return "", "", err
		}
		if _, err := out.Write(ct); err != nil {
			return "", "", err
		}
		if last {
			break
		}
	}
	sealedDataKey, err = appconfig.SealBound(pub, b, key)
	if err != nil {
		return "", "", err
	}
	return sealedDataKey, digestString(h), nil
}

func newGCM(key []byte) (cipher.AEAD, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return cipher.NewGCM(block)
}

// boundEnvelope reports whether an appconfig envelope is v2 (bound to an
// account and app). Twin of appconfig's unexported boundPrefix and
// splitVersion (crypto.go): appconfig.Open opens a v1 envelope under any
// binding, so Open checks the version itself.
func boundEnvelope(envelope string) bool {
	raw, err := base64.StdEncoding.DecodeString(envelope)
	return err == nil && len(raw) > 4 && string(raw[:4]) == "CSB2"
}

// Open reverses Seal into w. An unbound (v1) data key, a data key sealed
// for another binding, any tampered chunk, a stream cut short and data
// after the final chunk all return ErrSealed. Each chunk is written to w
// only after it authenticates, but the stream as a whole is authentic only
// when Open returns nil: a reader must treat what it got as provisional
// until then.
func Open(sealed io.Reader, sealedDataKey string, priv *ecdh.PrivateKey, b appconfig.Binding, w io.Writer) error {
	if !boundEnvelope(sealedDataKey) {
		return fmt.Errorf("%w: the data key is not bound to an app", ErrSealed)
	}
	key, err := appconfig.Open(priv, b, sealedDataKey)
	if err != nil {
		return fmt.Errorf("%w: data key: %v", ErrSealed, err)
	}
	defer clear(key)
	if len(key) != 32 {
		return fmt.Errorf("%w: data key of %d bytes", ErrSealed, len(key))
	}
	gcm, err := newGCM(key)
	if err != nil {
		return err
	}
	full := uint32(ChunkSize + gcm.Overhead())
	ct := make([]byte, full)
	pt := make([]byte, 0, ChunkSize)
	for index := uint64(0); ; index++ {
		if index >= maxChunks {
			return fmt.Errorf("%w: more than %d chunks", ErrSealed, maxChunks)
		}
		var hdr [4]byte
		if _, err := io.ReadFull(sealed, hdr[:]); err != nil {
			return truncated(err, index)
		}
		n := binary.BigEndian.Uint32(hdr[:])
		if n < uint32(gcm.Overhead()) || n > full {
			return fmt.Errorf("%w: chunk %d has length %d", ErrSealed, index, n)
		}
		if _, err := io.ReadFull(sealed, ct[:n]); err != nil {
			return truncated(err, index)
		}
		last := n < full
		nonce, aad := nonceFor(index, last)
		// Into pt, never ct[:0]: a failed Open may overwrite its dst.
		out, err := gcm.Open(pt[:0], nonce, ct[:n], aad)
		if err != nil {
			return fmt.Errorf("%w: chunk %d does not authenticate", ErrSealed, index)
		}
		if _, err := w.Write(out); err != nil {
			return err
		}
		if !last {
			continue
		}
		// The final chunk must end the stream: trailing ciphertext would
		// otherwise pass Verify against a manifest whose sealed digest
		// covers only this authenticated prefix (round-6 finding 6).
		var one [1]byte
		m, err := io.ReadFull(sealed, one[:])
		if m > 0 {
			return fmt.Errorf("%w: data after the final chunk", ErrSealed)
		}
		if err != io.EOF {
			return err
		}
		return nil
	}
}

// truncated classifies a short read: the stream ending is ErrSealed, any
// other read error is the reader's own.
func truncated(err error, index uint64) error {
	if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
		return fmt.Errorf("%w: truncated at chunk %d", ErrSealed, index)
	}
	return err
}