package mcrypto

import (
	"bytes"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"hash"
	"io"
	"strings"
	"time"

	"github.com/mediocregopher/mediocre-go-lib/mlog"
)

var (
	errMalformedSig = errors.New("malformed signature")

	// ErrInvalidSig is returned by Signer related functions when an invalid
	// signature is used, e.g. it is a signature for different data, or uses a
	// different secret key, or has expired
	ErrInvalidSig = errors.New("invalid signature")
)

// Signer is some entity which can generate signatures for arbitrary data and
// can later verify those signatures
type Signer interface {
	sign(io.Reader) (string, error)

	// returns an error if io.Reader returns one ever, or if the signature
	// couldn't be verified
	verify(string, io.Reader) error
}

// Sign reads all data from the io.Reader and signs it using the given Signer
func Sign(s Signer, r io.Reader) (string, error) {
	return s.sign(r)
}

// SignBytes uses the Signer to generate a signature for the given []bytes
func SignBytes(s Signer, b []byte) string {
	sig, err := s.sign(bytes.NewBuffer(b))
	if err != nil {
		panic(err)
	}
	return sig
}

// SignString uses the Signer to generate a signature for the given string
func SignString(s Signer, in string) string {
	return SignBytes(s, []byte(in))
}

// Verify reads all data from the io.Reader and uses the Signer to verify that
// the signature is for that data.
//
// Returns any errors from io.Reader, or ErrInvalidSig (use merry.Is(err,
// mcrypto.ErrInvalidSig) to check).
func Verify(s Signer, sig string, r io.Reader) error {
	return s.verify(sig, r)
}

// VerifyBytes uses the Signer to verify that the signature is for the given
// []bytes.
//
// Returns any errors from io.Reader, or ErrInvalidSig (use merry.Is(err,
// mcrypto.ErrInvalidSig) to check).
func VerifyBytes(s Signer, sig string, b []byte) error {
	return s.verify(sig, bytes.NewBuffer(b))
}

// VerifyString uses the Signer to verify that the signature is for the given
// string.
//
// Returns any errors from io.Reader, or ErrInvalidSig (use merry.Is(err,
// mcrypto.ErrInvalidSig) to check).
func VerifyString(s Signer, sig, in string) error {
	return VerifyBytes(s, sig, []byte(in))
}

////////////////////////////////////////////////////////////////////////////////

type signer struct {
	outSize uint8 // in bytes, shouldn't be more than 32, cause sha256
	secret  []byte
}

// NewSigner returns a Signer instance which will use the given secret to sign
// and verify all signatures. The signatures generated by this Signer have no
// expiration
func NewSigner(secret []byte) Signer {
	return signer{outSize: 20, secret: secret}
}

// NewWeakSigner returns a Signer, similar to how NewSigner does. The signatures
// generated by this Signer will be smaller in text size, and therefore weaker,
// but are still fine for most applications.
//
// The Signers returned by both NewSigner and NewWeakSigner can verify
// each-other's signatures, as long as the secret is the same.
func NewWeakSigner(secret []byte) Signer {
	return signer{outSize: 8, secret: secret}
}

func (s signer) signRaw(r io.Reader) (hash.Hash, error) {
	h := hmac.New(sha256.New, s.secret)
	_, err := io.Copy(h, r)
	return h, err
}

func (s signer) sign(r io.Reader) (string, error) {
	h, err := s.signRaw(r)
	if err != nil {
		return "", err
	}
	b := make([]byte, 1+h.Size())
	b[0] = s.outSize
	h.Sum(b[1:1])
	return sigV0 + hex.EncodeToString(b[:1+int(s.outSize)]), nil
}

func (s signer) verify(sig string, r io.Reader) error {
	sig, ok := stripPrefix(sig, sigV0)
	if !ok || len(sig) < 2 {
		return mlog.ErrWithKV(errMalformedSig, mlog.KV{"sig": sig})
	}
	sig = strings.TrimPrefix(sig, sigV0)

	sizeStr, sig := sig[:2], sig[2:]
	sizeB, err := hex.DecodeString(sizeStr)
	if err != nil {
		return mlog.ErrWithKV(err, mlog.KV{"sig": sig})
	}

	size := sizeB[0]
	if hex.DecodedLen(len(sig)) != int(size) {
		return mlog.ErrWithKV(errMalformedSig, mlog.KV{"sig": sig})
	}

	sigB, err := hex.DecodeString(sig)
	if err != nil {
		return mlog.ErrWithKV(err, mlog.KV{"sig": sig})
	}

	h, err := s.signRaw(r)
	if err != nil {
		return mlog.ErrWithKV(err, mlog.KV{"sig": sig})
	}

	if !hmac.Equal(sigB, h.Sum(nil)[:size]) {
		return mlog.ErrWithKV(ErrInvalidSig, mlog.KV{"sig": sig})
	}

	return nil
}

////////////////////////////////////////////////////////////////////////////////

type expireSigner struct {
	s       Signer
	timeout time.Duration

	// only used during tests
	testNow time.Time
}

// ExpireSigner wraps a Signer so that the signatures produced include timestamp
// information about when the signature was made. That information is then used
// during verifying to ensure the signature isn't older than the timeout.
//
// It is allowed to change the timeout ExpireSigner is initialized with.
// Previously generated signatures will be verified (or rejected) using the new
// timeout.
func ExpireSigner(s Signer, timeout time.Duration) Signer {
	return expireSigner{s: s, timeout: timeout}
}

func (es expireSigner) now() time.Time {
	if !es.testNow.IsZero() {
		return es.testNow
	}
	return time.Now()
}

func (es expireSigner) sign(r io.Reader) (string, error) {
	b := make([]byte, 8)
	binary.BigEndian.PutUint64(b, uint64(es.now().UnixNano()))
	sig, err := es.s.sign(prefixReader(r, b))
	return exSigV0 + hex.EncodeToString(b) + sig, err
}

var exSigTimeLen = hex.EncodedLen(8)

func (es expireSigner) verify(sig string, r io.Reader) error {
	sig, ok := stripPrefix(sig, exSigV0)
	if !ok || len(sig) < exSigTimeLen {
		return mlog.ErrWithKV(errMalformedSig, mlog.KV{"sig": sig})
	}

	tStr, sig := sig[:exSigTimeLen], sig[exSigTimeLen:]
	tB, err := hex.DecodeString(tStr)
	if err != nil {
		return mlog.ErrWithKV(err, mlog.KV{"sig": sig})
	}

	t := time.Unix(0, int64(binary.BigEndian.Uint64(tB)))
	if es.now().Sub(t) > es.timeout {
		return mlog.ErrWithKV(ErrInvalidSig, mlog.KV{"sig": sig})
	}

	return es.s.verify(sig, prefixReader(r, tB))
}

////////////////////////////////////////////////////////////////////////////////

type uniqueSigner struct {
	s        Signer
	randSize uint8 // in bytes
}

// UniqueSigner wraps a Signer so that when data is signed some random data is
// included in the signed data, and that random data is included in the
// signature as well. This ensures that even for the same input data signatures
// produced are all unique.
func UniqueSigner(s Signer) Signer {
	return uniqueSigner{s: s, randSize: 10}
}

func (us uniqueSigner) sign(r io.Reader) (string, error) {
	b := make([]byte, 1+us.randSize)
	b[0] = us.randSize
	if _, err := rand.Read(b[1:]); err != nil {
		panic(err)
	}
	sig, err := us.s.sign(prefixReader(r, b[1:]))
	return uniqueSigV0 + hex.EncodeToString(b) + sig, err
}

func (us uniqueSigner) verify(sig string, r io.Reader) error {
	sig, ok := stripPrefix(sig, uniqueSigV0)
	if !ok || len(sig) < 2 {
		return mlog.ErrWithKV(errMalformedSig, mlog.KV{"sig": sig})
	}

	sizeStr, sig := sig[:2], sig[2:]
	sizeB, err := hex.DecodeString(sizeStr)
	if err != nil {
		return mlog.ErrWithKV(err, mlog.KV{"sig": sig})
	}

	size := sizeB[0]
	sizeEnc := hex.EncodedLen(int(size))
	if len(sig) < sizeEnc {
		return mlog.ErrWithKV(errMalformedSig, mlog.KV{"sig": sig})
	}

	bStr, sig := sig[:sizeEnc], sig[sizeEnc:]
	b, err := hex.DecodeString(bStr)
	if err != nil {
		return mlog.ErrWithKV(err, mlog.KV{"sig": sig})
	}

	return us.s.verify(sig, prefixReader(r, b))
}