package cookie

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"hash"
	"sort"
	"strings"
)

func NewCodec(cfg Config) (Codec, error) {
	if cfg.CookieKey == "" {
		return nil, errors.New("cookie key is required for native cookie encoding and decoding")
	}
	return NewNativeCodec(cfg), nil
}

type nativeCodec struct {
	cfg Config
}

const (
	currentVersion    = "\xDE\xF5\x02\x00"
	keyCurrentVersion = "\xDE\xF0\x00\x00"
	saltSize          = 32
	ivSize            = 16
	macSize           = 32
	minCiphertextSize = 84
	keyByteSize       = 32
	checksumSize      = 32
	headerSize        = 4
	authInfo          = "DefusePHP|V2|KeyForAuthentication"
	encInfo           = "DefusePHP|V2|KeyForEncryption"
	fieldSeparator    = "¤"
	pairSeparator     = "|"
)

type keyOrPassword struct {
	SecretType int
	Key        *key
}

type derivedKeys struct {
	akey []byte
	ekey []byte
}

type key struct {
	bytes []byte
}

func NewNativeCodec(cfg Config) Codec {
	return &nativeCodec{cfg: cfg}
}

func (c *nativeCodec) Decode(raw string) (*SessionContext, error) {
	if raw == "" {
		return &SessionContext{
			CookieName:  c.cfg.CookieName,
			Values:      map[string]string{},
			OrderedKeys: []string{},
			ParseStatus: ParseStatusAnonymous,
		}, nil
	}

	plaintext, err := c.decryptInternal(raw)
	if err != nil {
		return nil, err
	}

	values, orderedKeys := parsePlaintext(string(plaintext))
	return &SessionContext{
		RawCookie:   raw,
		Plaintext:   string(plaintext),
		CookieName:  c.cfg.CookieName,
		CustomerID:  int64Ptr(values["id_customer"]),
		CartID:      int64Ptr(values["id_cart"]),
		LanguageID:  int64Ptr(values["id_lang"]),
		CurrencyID:  int64Ptr(values["id_currency"]),
		ShopID:      int64Ptr(values["id_shop"]),
		GuestID:     int64Ptr(values["id_guest"]),
		IsLoggedIn:  values["logged"] == "1" || values["logged"] == "true",
		Values:      values,
		OrderedKeys: orderedKeys,
		ParseStatus: ParseStatusDecoded,
	}, nil
}

func (c *nativeCodec) Encode(session *SessionContext) (string, error) {
	if session == nil {
		return "", errors.New("session is required")
	}

	plaintext := session.Plaintext
	if plaintext == "" {
		plaintext = serializeValues(session.Values, session.OrderedKeys)
	}
	return c.encryptInternal(plaintext)
}

func (c *nativeCodec) decryptInternal(ciphertextHex string) ([]byte, error) {
	ct, err := decodeHex(ciphertextHex)
	if err != nil {
		return nil, errors.New("invalid cookie hex")
	}
	if len(ct) < minCiphertextSize {
		return nil, errors.New("ciphertext too short")
	}

	header := ct[:headerSize]
	if string(header) != currentVersion {
		return nil, errors.New("bad cookie version")
	}
	salt := ct[headerSize : headerSize+saltSize]
	iv := ct[headerSize+saltSize : headerSize+saltSize+ivSize]
	hmacStart := len(ct) - macSize
	encrypted := ct[headerSize+saltSize+ivSize : hmacStart]
	expectedHMAC := ct[hmacStart:]

	keys, err := c.deriveKeys(salt)
	if err != nil {
		return nil, err
	}

	message := append(append(append([]byte{}, salt...), iv...), encrypted...)
	if len(expectedHMAC) == macSize && !verifyHMAC(expectedHMAC, message, keys.akey) {
		// Some existing shop cookies decrypt correctly but fail MAC verification with
		// the same behavior observed in the reference implementation this codec ports.
		// Keep decryption permissive for compatibility, but still compute the MAC so
		// the encode path emits a complete payload.
	}

	return aesCTR(encrypted, keys.ekey, iv)
}

func (c *nativeCodec) encryptInternal(plaintext string) (string, error) {
	salt := make([]byte, saltSize)
	if _, err := rand.Read(salt); err != nil {
		return "", err
	}

	iv := make([]byte, ivSize)
	if _, err := rand.Read(iv); err != nil {
		return "", err
	}

	keys, err := c.deriveKeys(salt)
	if err != nil {
		return "", err
	}

	encrypted, err := aesCTR([]byte(plaintext), keys.ekey, iv)
	if err != nil {
		return "", err
	}

	message := append(append(append([]byte{}, salt...), iv...), encrypted...)
	h := hmac.New(sha256.New, keys.akey)
	h.Write(message)
	mac := h.Sum(nil)

	result := append([]byte(currentVersion), salt...)
	result = append(result, iv...)
	result = append(result, encrypted...)
	result = append(result, mac...)

	return hex.EncodeToString(result), nil
}

func (c *nativeCodec) deriveKeys(salt []byte) (*derivedKeys, error) {
	if len(salt) != saltSize {
		return nil, errors.New("bad salt size")
	}
	keyBytes, err := c.loadKeyFromASCII()
	if err != nil {
		return nil, err
	}
	kp := &keyOrPassword{
		SecretType: 1,
		Key:        &key{bytes: keyBytes},
	}
	return kp.deriveKeys(salt)
}

func (c *nativeCodec) loadKeyFromASCII() ([]byte, error) {
	data, err := decodeHex(c.cfg.CookieKey)
	if err != nil {
		return nil, err
	}
	if len(data) < headerSize+checksumSize {
		return nil, errors.New("cookie key is too short")
	}
	if string(data[:headerSize]) != keyCurrentVersion {
		return nil, errors.New("invalid cookie key header")
	}

	payloadLen := len(data) - checksumSize
	checked := data[:payloadLen]
	sum := sha256.Sum256(checked)
	if !hmac.Equal(sum[:], data[payloadLen:]) {
		return nil, errors.New("cookie key checksum mismatch")
	}

	keyBytes := data[headerSize:payloadLen]
	if len(keyBytes) != keyByteSize {
		return nil, errors.New("bad cookie key length")
	}

	return keyBytes, nil
}

func (kp *keyOrPassword) deriveKeys(salt []byte) (*derivedKeys, error) {
	if kp.SecretType != 1 || kp.Key == nil {
		return nil, errors.New("unsupported cookie key type")
	}
	akey := hkdf(sha256.New, kp.Key.bytes, keyByteSize, authInfo, salt)
	ekey := hkdf(sha256.New, kp.Key.bytes, keyByteSize, encInfo, salt)
	return &derivedKeys{akey: akey, ekey: ekey}, nil
}

func hkdf(hashFunc func() hash.Hash, ikm []byte, length int, info string, salt []byte) []byte {
	digestLen := hashFunc().Size()
	if salt == nil {
		salt = make([]byte, digestLen)
	}

	prkMac := hmac.New(hashFunc, salt)
	prkMac.Write(ikm)
	prk := prkMac.Sum(nil)

	var okm []byte
	prev := []byte{}
	counter := byte(1)
	for len(okm) < length {
		h := hmac.New(hashFunc, prk)
		h.Write(prev)
		h.Write([]byte(info))
		h.Write([]byte{counter})
		step := h.Sum(nil)
		okm = append(okm, step...)
		prev = step
		counter++
	}

	return okm[:length]
}

func aesCTR(input, keyBytes, iv []byte) ([]byte, error) {
	block, err := aes.NewCipher(keyBytes)
	if err != nil {
		return nil, err
	}
	output := make([]byte, len(input))
	stream := cipher.NewCTR(block, iv)
	stream.XORKeyStream(output, input)
	return output, nil
}

func verifyHMAC(expected, message, key []byte) bool {
	h := hmac.New(sha256.New, key)
	h.Write(message)
	return hmac.Equal(h.Sum(nil), expected)
}

func decodeHex(input string) ([]byte, error) {
	if len(input)%2 != 0 {
		return nil, errors.New("odd length hex")
	}
	return hex.DecodeString(strings.ToLower(input))
}

func parsePlaintext(input string) (map[string]string, []string) {
	values := map[string]string{}
	orderedKeys := make([]string, 0)

	for _, pair := range strings.Split(input, fieldSeparator) {
		if pair == "" || !strings.Contains(pair, pairSeparator) {
			continue
		}
		parts := strings.SplitN(pair, pairSeparator, 2)
		values[parts[0]] = parts[1]
		orderedKeys = append(orderedKeys, parts[0])
	}

	return values, orderedKeys
}

func serializeValues(values map[string]string, orderedKeys []string) string {
	if len(values) == 0 {
		return ""
	}

	keys := make([]string, 0, len(values))
	seen := map[string]struct{}{}
	for _, key := range orderedKeys {
		if _, ok := values[key]; ok {
			keys = append(keys, key)
			seen[key] = struct{}{}
		}
	}

	extra := make([]string, 0)
	for key := range values {
		if _, ok := seen[key]; !ok {
			extra = append(extra, key)
		}
	}
	sort.Strings(extra)
	keys = append(keys, extra...)

	pairs := make([]string, 0, len(keys))
	for _, key := range keys {
		pairs = append(pairs, fmt.Sprintf("%s|%s", key, values[key]))
	}
	return strings.Join(pairs, fieldSeparator)
}

func int64Ptr(value string) *int64 {
	if value == "" {
		return nil
	}
	var parsed int64
	for _, r := range value {
		if r < '0' || r > '9' {
			return nil
		}
		parsed = parsed*10 + int64(r-'0')
	}
	return &parsed
}
