package atspace

import (
	"crypto/hkdf"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"fmt"

	"github.com/bluesky-social/indigo/atproto/atcrypto"
	"github.com/bluesky-social/indigo/atproto/syntax"
	"tangled.org/bnewbold.net/cobalt/atspace/xsyntax"
)

// TODO(review): the current structure of the constructor here constructs the context string twice when generating and signing a space commit. Also need to pass space and author args repeatedly. Maybe a wrapper helper would be good? or different API shape?
// TODO(review): is it good to used fixed-size byte arrays like [32]byte? I assume these can be stack-allocated nicely, but maybe this is resulting in copies
// TODO(review): should we use atdata.Bytes instead of [32]byte or []byte, so that JSON representation is mashalled as valid atproto data?
// TODO(review): the stdlib version of hkdf.Expand expects 'info' to be an actual 'string' as opposed to '[]byte'. but we are putting big-endian length prefixes in there, which might not be valid UTF-8. probably "fine" depending on how string([]byte{}) works, but feels a bit off.

const ATPROTO_SPACE_COMMIT_VERSION int64 = 1
const ATPROTO_SPACE_PROTOCOL_TAG string = "atproto-space-v1"

type SpaceCommit struct {
	// commit format version (ATPROTO_SPACE_COMMIT_VERSION)
	Version int64 `json:"ver" cborgen:"ver"`
	// SHA-256 of the LtHash state (32 bytes)
	Hash [32]byte `json:"hash" cborgen:"hash"`
	// Per-signature nonce (32 random bytes)
	IKM [32]byte `json:"ikm" cborgen:"ikm"`
	// Signature of context string bytes
	Signature []byte `json:"sig" cborgen:"sig"`
	// HMAC over context string, ikm, and hash: HMAC-SHA256(HKDF-Expand(ikm, ctx, 32), hash)
	MAC []byte `json:"mac" cborgen:"mac"`
	// Commit revision (TID), also bound into context string
	Revision syntax.TID `json:"rev" cborgen:"rev"`
}

// Constructs an unsigned space commit struct with 'mac' and secure random 'ikm'
func NewSpaceCommit(hash [32]byte, rev syntax.TID, space xsyntax.SpaceRef, author syntax.DID) SpaceCommit {
	ikm := new([32]byte)
	rand.Read(ikm[:])

	ctxstr := BuildContextString(space, author, rev, *ikm)
	mac := ComputeMAC(*ikm, ctxstr, hash)

	return SpaceCommit{
		Version:  ATPROTO_SPACE_COMMIT_VERSION,
		Hash:     hash,
		IKM:      *ikm,
		MAC:      mac,
		Revision: rev,
	}
}

func (c *SpaceCommit) Sign(space xsyntax.SpaceRef, author syntax.DID, privkey atcrypto.PrivateKey) error {
	ctxstr := BuildContextString(space, author, c.Revision, c.IKM)
	sig, err := privkey.HashAndSign(ctxstr)
	if err != nil {
		return err
	}
	c.Signature = sig
	return nil
}

func (c *SpaceCommit) VerifySignature(space xsyntax.SpaceRef, author syntax.DID, pubkey atcrypto.PublicKey) error {
	if len(c.Signature) == 0 {
		return fmt.Errorf("can not verify unsigned space commit")
	}
	ctxstr := BuildContextString(space, author, c.Revision, c.IKM)
	return pubkey.HashAndVerify(ctxstr, c.Signature)
}

func (c *SpaceCommit) VerifyMAC(space xsyntax.SpaceRef, author syntax.DID) error {
	ctxstr := BuildContextString(space, author, c.Revision, c.IKM)
	mac := ComputeMAC(c.IKM, ctxstr, c.Hash)
	if !hmac.Equal(c.MAC, mac) {
		return fmt.Errorf("space commit MAC validation mismatch")
	}
	return nil
}

func BuildContextString(space xsyntax.SpaceRef, author syntax.DID, rev syntax.TID, ikm [32]byte) []byte {
	/*
		"atproto-space-v1"            // fixed protocol tag
		   || uint16be(len(space))  || space  // space URI (at://authority/space/type/skey)
		   || uint16be(len(author)) || author // author DID of the repo
		   || uint16be(len(rev))    || rev    // commit revision (TID)
		   || uint16be(len(ikm))    || ikm    // per-signature nonce (below)
	*/

	// pre-allocate buffer
	out := make([]byte, 0, len(ATPROTO_SPACE_PROTOCOL_TAG)+len(space)+len(author)+len(rev)+len(ikm)+2*4)
	out = append(out, []byte(ATPROTO_SPACE_PROTOCOL_TAG)...)
	out = binary.BigEndian.AppendUint16(out, uint16(len(space)))
	out = append(out, []byte(space)...)
	out = binary.BigEndian.AppendUint16(out, uint16(len(author)))
	out = append(out, []byte(author)...)
	out = binary.BigEndian.AppendUint16(out, uint16(len(rev)))
	out = append(out, []byte(rev)...)
	out = binary.BigEndian.AppendUint16(out, uint16(len(ikm)))
	out = append(out, ikm[:]...)
	return out
}

func ComputeMAC(ikm [32]byte, ctxstr []byte, hash [32]byte) []byte {

	key, err := hkdf.Expand(sha256.New, ikm[:], string(ctxstr), 32)
	if err != nil || len(key) != 32 {
		panic("unexpected HKDF expansion error")
	}
	hmac := hmac.New(sha256.New, key)
	return hmac.Sum(hash[:])
}