package digest

import (
	"bytes"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/hex"
	"errors"
	"hash"
	"strings"
)

type Algorithm int

const (
	Undefined Algorithm = iota
	SHA256
	SHA384
	SHA512
)

type Digest struct {
	algorithm Algorithm
	decoded   []byte
}

func NewDigest(algorithm Algorithm, payload []byte) *Digest {
	var sum []byte
	var hasher hash.Hash
	switch algorithm {
	case SHA512:
		hasher = sha512.New()
	case SHA384:
		hasher = sha512.New384()
	default:
		hasher = sha256.New()
	}
	hasher.Write(payload)
	sum = hasher.Sum(nil)
	digobj := &Digest{
		algorithm: algorithm,
		decoded:   sum,
	}
	return digobj
}

func ParseDigest(str string) (*Digest, error) {
	var err error
	digobj := &Digest{}

	str = strings.ToLower(str)
	str = strings.TrimPrefix(str, "sha256:")
	str = strings.TrimPrefix(str, "sha384:")
	str = strings.TrimPrefix(str, "sha512:")

	decoded, err := hex.DecodeString(str)
	if err != nil {
		err := errors.New("Can't decode digobj")
		return digobj, err
	}
	digobj.decoded = decoded
	switch len(decoded) {
	case 32:
		digobj.algorithm = SHA256
	case 48:
		digobj.algorithm = SHA384
	case 64:
		digobj.algorithm = SHA512
	default:
		err = errors.New("Unknown digest type")
		return digobj, err
	}
	return digobj, err
}

func (dig *Digest) Encoded() string {
	return hex.EncodeToString(dig.decoded)
}

func (dig *Digest) Algorithm() Algorithm {
	return dig.algorithm
}

func (dig *Digest) Prefix() string {
	var prefix string
	switch dig.algorithm {
	case SHA256:
		prefix = "sha256"
	case SHA384:
		prefix = "sha384"
	case SHA512:
		prefix = "sha512"
	}
	return prefix
}

func (dig *Digest) String() string {
	var prefix string
	switch dig.algorithm {
	case SHA256:
		prefix = "sha256"
	case SHA384:
		prefix = "sha384"
	case SHA512:
		prefix = "sha512"
	}
	hexx := hex.EncodeToString(dig.decoded)
	return prefix + ":" + hexx
}

func DetectDigestType(digobj string) Algorithm {
	var err error
	var algorithm Algorithm
	digobj = strings.ToLower(digobj)
	digobj = strings.TrimPrefix(digobj, "sha256:")
	digobj = strings.TrimPrefix(digobj, "sha384:")
	digobj = strings.TrimPrefix(digobj, "sha512:")
	decoded, err := hex.DecodeString(digobj)
	if err != nil {
		return Undefined
	}
	switch len(decoded) {
	case 32:
		algorithm = SHA256
	case 48:
		algorithm = SHA384
	case 64:
		algorithm = SHA512
	default:
		algorithm = Undefined
	}
	return algorithm
}

func MakeSHA256Digest(src []byte) string {
	hasher := sha256.New()
	hasher.Write(src)
	sum := hasher.Sum(nil)
	return "sha256:" + hex.EncodeToString(sum)
}

func MakeSHA384Digest(src []byte) string {
	hasher := sha512.New384()
	hasher.Write(src)
	sum := hasher.Sum(nil)
	return "sha256:" + hex.EncodeToString(sum)
}

func MakeSHA512Digest(src []byte) string {
	hasher := sha512.New()
	hasher.Write(src)
	sum := hasher.Sum(nil)
	return "sha512:" + hex.EncodeToString(sum)
}

func (dig Digest) MarshalJSON() ([]byte, error) {
	writer := bytes.NewBuffer(nil)
	writer.WriteString(`"`)
	switch dig.algorithm {
	case SHA256:
		writer.WriteString("sha256")
	case SHA384:
		writer.WriteString("sha384")
	case SHA512:
		writer.WriteString("sha512")
	}
	writer.WriteString(":")
	writer.WriteString(hex.EncodeToString(dig.decoded))
	writer.WriteString(`"`)

	return writer.Bytes(), nil
}

func (dig *Digest) UnmarshalJSON(data []byte) error {
	var err error
	digobj := Digest{}

	str := strings.ToLower(string(data))
	str = strings.TrimPrefix(str, `"`)
	str = strings.TrimSuffix(str, `"`)
	str = strings.TrimPrefix(str, "sha256:")
	str = strings.TrimPrefix(str, "sha384:")
	str = strings.TrimPrefix(str, "sha512:")

	decoded, err := hex.DecodeString(str)
	if err != nil {
		err := errors.New("Can't decode digest")
		return err
	}
	digobj.decoded = decoded
	switch len(decoded) {
	case 32:
		digobj.algorithm = SHA256
	case 48:
		digobj.algorithm = SHA384
	case 64:
		digobj.algorithm = SHA512
	default:
		err = errors.New("Unknown digest type")
		return err
	}
	*dig = digobj
	return err

}
