package dbreads

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

	"golang.org/x/crypto/hkdf"
)

// Sealing is the appconfig v2 construction (spec §5.6): ECDH P-256 with a
// fresh sender key per seal, HKDF-SHA256 with this package's own info string,
// AES-256-GCM with bound associated data. The wire is raw bytes, never
// base64, in the bucket: "DBR1" || sender key (65) || nonce (12) || ciphertext.
const (
	sealPrefix   = "DBR1"
	sealEphLen   = 65
	sealNonceLen = 12
	// SealOverhead is what sealing adds: under 100 bytes (spec §5.7).
	SealOverhead = len(sealPrefix) + sealEphLen + sealNonceLen + 16
)

// The streams a chunk can carry, bound into its associated data.
const (
	StreamRows = "rows" // a query's output
	StreamLTX  = "ltx"  // a pull's planned files (plan 3)
)

// ErrSeal is any sealed object that does not open: wrong key, moved,
// reordered, relabelled, truncated or corrupt.
var ErrSeal = errors.New("dbreads: the sealed bytes do not open")

var errBadReplyKey = errors.New("dbreads: reply_key is not an uncompressed P-256 point")

// randReader is a var so a test can make sealing deterministic.
var randReader io.Reader = rand.Reader

// NewReplyKey is the CLI's ephemeral key pair for one read. The private half
// never leaves the CLI's memory.
func NewReplyKey() (*ecdh.PrivateKey, error) { return ecdh.P256().GenerateKey(randReader) }

func EncodeReplyKey(pub *ecdh.PublicKey) string {
	return base64.StdEncoding.EncodeToString(pub.Bytes())
}

func ParseReplyKey(s string) (*ecdh.PublicKey, error) {
	raw, err := base64.StdEncoding.Strict().DecodeString(s)
	if err != nil || len(raw) != sealEphLen {
		return nil, errBadReplyKey
	}
	pub, err := ecdh.P256().NewPublicKey(raw)
	if err != nil {
		return nil, errBadReplyKey
	}
	return pub, nil
}

// lengthPrefixed encodes fields so no choice of values makes two lists
// encode the same.
func lengthPrefixed(fields ...string) []byte {
	var out []byte
	for _, f := range fields {
		out = binary.BigEndian.AppendUint32(out, uint32(len(f)))
		out = append(out, f...)
	}
	return out
}

func chunkAAD(id string, seq int64, stream string, last bool) []byte {
	l := "0"
	if last {
		l = "1"
	}
	return lengthPrefixed(Purpose, "chunk", id, strconv.FormatInt(seq, 10), stream, l)
}

func summaryAAD(id string) []byte { return lengthPrefixed(Purpose, "summary", id) }

func gcmFor(shared []byte) (cipher.AEAD, error) {
	key := make([]byte, 32)
	if _, err := io.ReadFull(hkdf.New(sha256.New, shared, nil, []byte(Purpose)), key); err != nil {
		return nil, err
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return cipher.NewGCM(block)
}

func seal(pub *ecdh.PublicKey, aad, plain []byte) ([]byte, error) {
	eph, err := ecdh.P256().GenerateKey(randReader)
	if err != nil {
		return nil, err
	}
	nonce := make([]byte, sealNonceLen)
	if _, err := io.ReadFull(randReader, nonce); err != nil {
		return nil, err
	}
	shared, err := eph.ECDH(pub)
	if err != nil {
		return nil, err
	}
	gcm, err := gcmFor(shared)
	if err != nil {
		return nil, err
	}
	out := make([]byte, 0, SealOverhead+len(plain))
	out = append(out, sealPrefix...)
	out = append(out, eph.PublicKey().Bytes()...)
	out = append(out, nonce...)
	return gcm.Seal(out, nonce, plain, aad), nil
}

func open(priv *ecdh.PrivateKey, aad, sealed []byte) ([]byte, error) {
	if len(sealed) < SealOverhead || string(sealed[:len(sealPrefix)]) != sealPrefix {
		return nil, ErrSeal
	}
	body := sealed[len(sealPrefix):]
	ephPub, err := ecdh.P256().NewPublicKey(body[:sealEphLen])
	if err != nil {
		return nil, ErrSeal
	}
	shared, err := priv.ECDH(ephPub)
	if err != nil {
		return nil, ErrSeal
	}
	gcm, err := gcmFor(shared)
	if err != nil {
		return nil, ErrSeal
	}
	plain, err := gcm.Open(nil, body[sealEphLen:sealEphLen+sealNonceLen], body[sealEphLen+sealNonceLen:], aad)
	if err != nil {
		return nil, ErrSeal
	}
	return plain, nil
}

// SealChunk seals one chunk of a read's output: bound to the read, its
// sequence number, its stream and whether it is the last.
func SealChunk(pub *ecdh.PublicKey, id string, seq int64, stream string, last bool, plain []byte) ([]byte, error) {
	return seal(pub, chunkAAD(id, seq, stream, last), plain)
}

func OpenChunk(priv *ecdh.PrivateKey, id string, seq int64, stream string, last bool, sealed []byte) ([]byte, error) {
	return open(priv, chunkAAD(id, seq, stream, last), sealed)
}

// SealSummary seals the result's summary (base64, it rides in result.json).
func SealSummary(pub *ecdh.PublicKey, id string, s Summary) (string, error) {
	raw, err := json.Marshal(s)
	if err != nil {
		return "", err
	}
	sealed, err := seal(pub, summaryAAD(id), raw)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(sealed), nil
}

func OpenSummary(priv *ecdh.PrivateKey, id, sealed string) (Summary, error) {
	raw, err := base64.StdEncoding.Strict().DecodeString(sealed)
	if err != nil {
		return Summary{}, ErrSeal
	}
	plain, err := open(priv, summaryAAD(id), raw)
	if err != nil {
		return Summary{}, err
	}
	var s Summary
	if err := json.Unmarshal(plain, &s); err != nil {
		return Summary{}, ErrSeal
	}
	return s, nil
}