package ueps

import (
	"bufio"
	"bytes"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/binary"
	"io"
)

// packet, err := ReadAndVerify(bufio.NewReader(conn), secret)
// if err == errMissingHMAC { /* no HMAC tag found in frame */ }
var errMissingHMAC = packetError("UEPS packet missing HMAC signature")

// packet, err := ReadAndVerify(bufio.NewReader(conn), wrongSecret)
// if err == errIntegrityViolation { header.ThreatScore += 100; /* reject packet */ }
var errIntegrityViolation = packetError("integrity violation: HMAC mismatch")

// dispatch(packet.Header.IntentID, packet.Header.ThreatScore, packet.Payload)
type ParsedPacket struct {
	Header  UEPSHeader
	Payload []byte
}

// packet, err := ueps.ReadAndVerify(bufio.NewReader(conn), []byte("my-shared-secret"))
// if err == nil { dispatch(packet.Header.IntentID, packet.Header.ThreatScore, packet.Payload) }
func ReadAndVerify(reader *bufio.Reader, sharedSecret []byte) (*ParsedPacket, error) {
	var signedData bytes.Buffer
	header := UEPSHeader{}
	var hmacSignature []byte
	var payload []byte

	for {
		tagType, err := reader.ReadByte()
		if err != nil {
			return nil, err
		}

		if tagType == TagPayload {
			payload, err = io.ReadAll(reader)
			if err != nil {
				return nil, err
			}
			break
		}

		tagValueLength, err := reader.ReadByte()
		if err != nil {
			return nil, err
		}

		tagValue := make([]byte, tagValueLength)
		if _, err := io.ReadFull(reader, tagValue); err != nil {
			return nil, err
		}

		switch tagType {
		case TagVersion:
			header.Version = tagValue[0]
			signedData.WriteByte(tagType)
			signedData.WriteByte(tagValueLength)
			signedData.Write(tagValue)
		case TagCurrentLayer:
			header.CurrentLayer = tagValue[0]
			signedData.WriteByte(tagType)
			signedData.WriteByte(tagValueLength)
			signedData.Write(tagValue)
		case TagTargetLayer:
			header.TargetLayer = tagValue[0]
			signedData.WriteByte(tagType)
			signedData.WriteByte(tagValueLength)
			signedData.Write(tagValue)
		case TagIntent:
			header.IntentID = tagValue[0]
			signedData.WriteByte(tagType)
			signedData.WriteByte(tagValueLength)
			signedData.Write(tagValue)
		case TagThreatScore:
			header.ThreatScore = binary.BigEndian.Uint16(tagValue)
			signedData.WriteByte(tagType)
			signedData.WriteByte(tagValueLength)
			signedData.Write(tagValue)
		case TagHMAC:
			hmacSignature = tagValue
		default:
			// signedData.Write([]byte{tagType, tagValueLength}); signedData.Write(tagValue) — unknown tags included in HMAC
			signedData.WriteByte(tagType)
			signedData.WriteByte(tagValueLength)
			signedData.Write(tagValue)
		}
	}

	if len(hmacSignature) == 0 {
		return nil, errMissingHMAC
	}

	messageAuthCode := hmac.New(sha256.New, sharedSecret)
	messageAuthCode.Write(signedData.Bytes())
	messageAuthCode.Write(payload)
	expectedMessageAuthCode := messageAuthCode.Sum(nil)

	if !hmac.Equal(hmacSignature, expectedMessageAuthCode) {
		return nil, errIntegrityViolation
	}

	return &ParsedPacket{
		Header:  header,
		Payload: payload,
	}, nil
}