package crypto

import (
	"encoding/binary"
	"errors"
)

// Compact binary wire format for EphemeralMessage (V2).
//
// Layout:
//   [2 bytes]   Version (uint16 LE, value=2)
//   [4 bytes]   Signature length (uint32 LE)
//   [variable]  Compact V2 signature bytes (Signature.MarshalV2 output, ~98 bytes)
//   [2 bytes]   Pattern site count (uint16 LE, total including unoccupied)
//   [2 bytes]   Pattern occupied count (uint16 LE)
//   [2 bytes]   Pattern tag count (uint16 LE)
//   [variable]  Pattern tag table: for each tag, 1-byte length + UTF-8 bytes
//   Per occupied site (6 bytes each):
//     [2 bytes]  Site index (uint16 LE)
//     [1 byte]   Tag table index (uint8)
//     [1 byte]   Projection (6-bit) | Perm low 2 bits (high 2 bits)
//     [1 byte]   Lock-in quantized to uint8 (0-255)
//     [1 byte]   Flags: bit 0 = perm bit 2, bits 1-7 reserved
//   Sender fingerprint:
//     [56 bytes]  Hash as Hamadryad (hashed from string)
//     [12 bytes]  PermDist (6 × uint16 LE)

const ephemeralWireVersion = 2

// MarshalEphemeral encodes an EphemeralMessage into compact binary form.
// The caller must supply the lattice dimension N (from Params) for V2
// signature encoding.
func (m *EphemeralMessage) MarshalEphemeral(n int) ([]byte, error) {
	if m == nil {
		return nil, errors.New("crypto: nil ephemeral message")
	}
	if m.Signature == nil {
		return nil, errors.New("crypto: nil signature in ephemeral message")
	}

	// Marshal the signature using compact V2 format (~98 bytes for N=256).
	sigBytes, err := m.Signature.MarshalV2(n)
	if err != nil {
		return nil, err
	}

	// Build pattern tag table.
	tagIndex := make(map[string]uint8)
	var tagTable []string
	var occupied []int
	for i, site := range m.Pattern {
		if site.Occupied {
			if _, ok := tagIndex[site.TypeTag]; !ok {
				if len(tagTable) >= 255 {
					return nil, errors.New("crypto: too many distinct pattern tags")
				}
				tagIndex[site.TypeTag] = uint8(len(tagTable))
				tagTable = append(tagTable, site.TypeTag)
			}
			occupied = append(occupied, i)
		}
	}

	// Calculate tag table size.
	tagTableSize := 0
	for _, tag := range tagTable {
		tagTableSize += 1 + len(tag)
	}

	// Total size.
	headerSize := 2 + 4 + len(sigBytes) + 2 + 2 + 2 + tagTableSize +
		len(occupied)*compactPerSiteSize + HamBytes + 12
	buf := make([]byte, headerSize)
	pos := 0

	// Version.
	binary.LittleEndian.PutUint16(buf[pos:], ephemeralWireVersion)
	pos += 2

	// Signature length + bytes.
	binary.LittleEndian.PutUint32(buf[pos:], uint32(len(sigBytes)))
	pos += 4
	copy(buf[pos:], sigBytes)
	pos += len(sigBytes)

	// Pattern site count (total).
	binary.LittleEndian.PutUint16(buf[pos:], uint16(len(m.Pattern)))
	pos += 2

	// Pattern occupied count.
	binary.LittleEndian.PutUint16(buf[pos:], uint16(len(occupied)))
	pos += 2

	// Pattern tag count.
	binary.LittleEndian.PutUint16(buf[pos:], uint16(len(tagTable)))
	pos += 2

	// Tag table.
	for _, tag := range tagTable {
		buf[pos] = uint8(len(tag))
		pos++
		copy(buf[pos:], tag)
		pos += len(tag)
	}

	// Per occupied site (6 bytes each).
	for _, idx := range occupied {
		site := m.Pattern[idx]

		binary.LittleEndian.PutUint16(buf[pos:], uint16(site.Index))
		pos += 2

		buf[pos] = tagIndex[site.TypeTag]
		pos++

		projPerm := (site.Projection & 0x3F) | ((site.Perm & 0x03) << 6)
		buf[pos] = projPerm
		pos++

		buf[pos] = quantizeLockIn(site.LockIn)
		pos++

		flags := uint8(0)
		if site.Perm&0x04 != 0 {
			flags |= 0x01
		}
		buf[pos] = flags
		pos++
	}

	// Sender fingerprint: hash as Hamadryad.
	fpHash := Hash([]byte(m.SenderFingerprint.Hash))
	copy(buf[pos:], fpHash[:])
	pos += HamBytes

	// PermDist (6 × uint16 LE).
	for i := range 6 {
		binary.LittleEndian.PutUint16(buf[pos:], uint16(m.SenderFingerprint.PermDist[i]))
		pos += 2
	}

	return buf[:pos], nil
}

// UnmarshalEphemeralMessage decodes a compact binary EphemeralMessage.
func UnmarshalEphemeralMessage(data []byte) (*EphemeralMessage, error) {
	if len(data) < 2 {
		return nil, errors.New("crypto: ephemeral message too short")
	}
	pos := 0

	// Version.
	version := binary.LittleEndian.Uint16(data[pos:])
	if version != ephemeralWireVersion {
		return nil, errors.New("crypto: unsupported ephemeral wire version")
	}
	pos += 2

	// Signature length.
	if pos+4 > len(data) {
		return nil, errors.New("crypto: truncated signature length")
	}
	sigLen := int(binary.LittleEndian.Uint32(data[pos:]))
	pos += 4

	if pos+sigLen > len(data) {
		return nil, errors.New("crypto: truncated signature data")
	}
	sig, consumed, err := UnmarshalSignatureV2(data[pos : pos+sigLen])
	if err != nil {
		return nil, err
	}
	_ = consumed // sigLen already bounds the slice
	pos += sigLen

	// Pattern counts.
	if pos+6 > len(data) {
		return nil, errors.New("crypto: truncated pattern header")
	}
	totalSites := int(binary.LittleEndian.Uint16(data[pos:]))
	pos += 2
	occCount := int(binary.LittleEndian.Uint16(data[pos:]))
	pos += 2
	tagCount := int(binary.LittleEndian.Uint16(data[pos:]))
	pos += 2

	// Tag table.
	tagTable := make([]string, tagCount)
	for i := range tagCount {
		if pos >= len(data) {
			return nil, errors.New("crypto: truncated pattern tag table")
		}
		tagLen := int(data[pos])
		pos++
		if pos+tagLen > len(data) {
			return nil, errors.New("crypto: truncated pattern tag name")
		}
		tagTable[i] = string(data[pos : pos+tagLen])
		pos += tagLen
	}

	// Per occupied site.
	needed := occCount*compactPerSiteSize + HamBytes + 12
	if pos+needed > len(data) {
		return nil, errors.New("crypto: truncated site data or fingerprint")
	}

	// Build pattern: start with all unoccupied, fill in occupied sites.
	pattern := make([]SiteMark, totalSites)
	for i := range occCount {
		siteIdx := binary.LittleEndian.Uint16(data[pos:])
		pos += 2

		tagIdx := data[pos]
		pos++
		tag := ""
		if int(tagIdx) < len(tagTable) {
			tag = tagTable[tagIdx]
		}

		projPerm := data[pos]
		pos++
		proj := projPerm & 0x3F
		permLow := (projPerm >> 6) & 0x03

		liQuant := data[pos]
		pos++
		li := dequantizeLockIn(liQuant)

		flags := data[pos]
		pos++
		perm := permLow
		if flags&0x01 != 0 {
			perm |= 0x04
		}

		if int(siteIdx) < totalSites {
			pattern[siteIdx] = SiteMark{
				Index:      uint64(siteIdx),
				Occupied:   true,
				TypeTag:    tag,
				Projection: proj,
				Perm:       perm,
				LockIn:     li,
			}
		} else {
			// Site index out of range: append.
			_ = i // consumed
			pattern = append(pattern, SiteMark{
				Index:      uint64(siteIdx),
				Occupied:   true,
				TypeTag:    tag,
				Projection: proj,
				Perm:       perm,
				LockIn:     li,
			})
		}
	}

	// Sender fingerprint.
	var fpHash Hamadryad
	copy(fpHash[:], data[pos:pos+HamBytes])
	pos += HamBytes

	var permDist [6]int
	for i := range 6 {
		permDist[i] = int(binary.LittleEndian.Uint16(data[pos:]))
		pos += 2
	}

	fp := SporeFingerprint{
		Hash:     encodeHamadryadHex(fpHash),
		PermDist: permDist,
	}

	return &EphemeralMessage{
		Pattern:           pattern,
		Signature:         sig,
		SenderFingerprint: fp,
	}, nil
}