package crypto

import (
	"encoding/binary"
	"errors"

	"git.mleku.dev/mleku/dendrite/pkg/ratio"
	"git.mleku.dev/mleku/dendrite/pkg/state"
)

// Compact binary wire format for signatures.
//
// Layout:
//   [56 bytes]  Challenge (Hamadryad hash of message)
//   [56 bytes]  Fingerprint hash (Hamadryad hash of spore)
//   [56 bytes]  Commitment (Hamadryad hash binding challenge to response)
//   [2 bytes]   Occupied count (uint16 LE)
//   [2 bytes]   Tag count (uint16 LE)
//   [variable]  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
//   Aggregate proof (24 bytes):
//     [2 bytes]   Mean lock-in (uint16 LE, fixed-point 0-65535 → 0.0-1.0)
//     [2 bytes]   Mean neighbor count (uint16 LE, fixed-point)
//     [8 bytes]   Hex state histogram (64 × 1 bit = 8 bytes, presence flags)
//     [12 bytes]  Perm distribution (6 × uint16 LE)

const (
	compactHeaderSize   = HamBytes + HamBytes + HamBytes + 2 + 2 // 172 bytes
	compactPerSiteSize  = 6
	compactAggProofSize = 2 + 2 + 8 + 12 // 24 bytes
)

// Marshal encodes a Signature into the compact binary wire format.
func (s *Signature) Marshal() ([]byte, error) {
	if s == nil {
		return nil, errors.New("crypto: nil signature")
	}

	// Build tag table from occupied sites.
	tagIndex := make(map[string]uint8)
	var tagTable []string
	occupied := make([]int, 0) // indices into Response
	for i, site := range s.Response {
		if site.Occupied {
			if _, ok := tagIndex[site.TypeTag]; !ok {
				if len(tagTable) >= 255 {
					return nil, errors.New("crypto: too many distinct 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) // length byte + UTF-8
	}

	totalSize := compactHeaderSize + tagTableSize +
		len(occupied)*compactPerSiteSize + compactAggProofSize
	buf := make([]byte, totalSize)
	pos := 0

	// Challenge (Hamadryad, 56 bytes).
	copy(buf[pos:], s.Challenge[:])
	pos += HamBytes

	// Fingerprint hash (Hamadryad, 56 bytes).
	fpHash := hashFingerprint(s.Fingerprint.Hash)
	copy(buf[pos:], fpHash[:])
	pos += HamBytes

	// Commitment (Hamadryad, 56 bytes).
	copy(buf[pos:], s.Commitment[:])
	pos += HamBytes

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

	// 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.
	for _, idx := range occupied {
		site := s.Response[idx]

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

		// Tag table index.
		buf[pos] = tagIndex[site.TypeTag]
		pos++

		// Projection (6-bit) | Perm low 2 bits in high 2 bits.
		projPerm := (site.Projection & 0x3F) | ((site.Perm & 0x03) << 6)
		buf[pos] = projPerm
		pos++

		// Lock-in quantized to uint8.
		buf[pos] = quantizeLockIn(site.LockIn)
		pos++

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

	// Aggregate proof.
	meanLI, meanNC := aggregateProof(s.Proof)
	binary.LittleEndian.PutUint16(buf[pos:], meanLI)
	pos += 2
	binary.LittleEndian.PutUint16(buf[pos:], meanNC)
	pos += 2

	// Hex state histogram: 64 bits = 8 bytes.
	var hexHist [8]byte
	for _, h := range s.Proof.HexTrace {
		if h < 64 {
			hexHist[h/8] |= 1 << (h % 8)
		}
	}
	copy(buf[pos:], hexHist[:])
	pos += 8

	// Perm distribution: 6 × uint16.
	permDist := [6]uint16{}
	for _, site := range s.Response {
		if site.Occupied && site.Perm < 6 {
			permDist[site.Perm]++
		}
	}
	for i := range 6 {
		binary.LittleEndian.PutUint16(buf[pos:], permDist[i])
		pos += 2
	}

	return buf[:pos], nil
}

// UnmarshalSignature decodes a compact binary signature.
func UnmarshalSignature(data []byte) (*Signature, error) {
	if len(data) < compactHeaderSize+compactAggProofSize {
		return nil, errors.New("crypto: compact signature too short")
	}
	pos := 0

	// Challenge (Hamadryad, 56 bytes).
	var challenge Hamadryad
	copy(challenge[:], data[pos:pos+HamBytes])
	pos += HamBytes

	// Fingerprint hash (Hamadryad, 56 bytes).
	var fpHashBytes Hamadryad
	copy(fpHashBytes[:], data[pos:pos+HamBytes])
	pos += HamBytes

	// Commitment (Hamadryad, 56 bytes).
	var commitment Hamadryad
	copy(commitment[:], data[pos:pos+HamBytes])
	pos += HamBytes

	// Occupied count.
	occCount := int(binary.LittleEndian.Uint16(data[pos:]))
	pos += 2

	// Tag count.
	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 tag table")
		}
		tagLen := int(data[pos])
		pos++
		if pos+tagLen > len(data) {
			return nil, errors.New("crypto: truncated tag name")
		}
		tagTable[i] = string(data[pos : pos+tagLen])
		pos += tagLen
	}

	// Check we have enough data for sites + aggregate proof.
	needed := occCount*compactPerSiteSize + compactAggProofSize
	if pos+needed > len(data) {
		return nil, errors.New("crypto: truncated site data")
	}

	// Per occupied site.
	response := make([]SiteMark, occCount)
	lockIns := make([]ratio.Ratio, occCount)
	neighborCounts := make([]int, occCount)
	hexTrace := make([]state.Hexagram, occCount)

	for i := range occCount {
		// Site index.
		siteIdx := binary.LittleEndian.Uint16(data[pos:])
		pos += 2

		// Tag table index.
		tagIdx := data[pos]
		pos++
		tag := ""
		if int(tagIdx) < len(tagTable) {
			tag = tagTable[tagIdx]
		}

		// Projection | Perm.
		projPerm := data[pos]
		pos++
		proj := projPerm & 0x3F
		permLow := (projPerm >> 6) & 0x03

		// Lock-in.
		liQuant := data[pos]
		pos++
		li := dequantizeLockIn(liQuant)

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

		response[i] = SiteMark{
			Index:      uint64(siteIdx),
			Occupied:   true,
			TypeTag:    tag,
			Projection: proj,
			Perm:       perm,
			LockIn:     li,
		}

		lockIns[i] = li
		// Neighbor counts not stored per-site in compact form;
		// set to 1 (minimum valid) — verifier checks >= 1.
		neighborCounts[i] = 1
	}

	// Aggregate proof.
	// Mean lock-in (informational, not used in per-site verification).
	_ = binary.LittleEndian.Uint16(data[pos:])
	pos += 2
	// Mean neighbor count.
	_ = binary.LittleEndian.Uint16(data[pos:])
	pos += 2

	// Hex histogram → reconstruct hex trace.
	var hexHist [8]byte
	copy(hexHist[:], data[pos:pos+8])
	pos += 8
	// Assign hex values from histogram to trace entries round-robin.
	var presentHex []state.Hexagram
	for b := range 64 {
		if hexHist[b/8]&(1<<(b%8)) != 0 {
			presentHex = append(presentHex, state.Hexagram(b))
		}
	}
	if len(presentHex) > 0 {
		for i := range hexTrace {
			hexTrace[i] = presentHex[i%len(presentHex)]
		}
	}

	// Perm distribution (6 × uint16).
	permDist := [6]int{}
	for i := range 6 {
		permDist[i] = int(binary.LittleEndian.Uint16(data[pos:]))
		pos += 2
	}

	// Reconstruct fingerprint with stored hash and perm dist.
	fp := SporeFingerprint{
		Hash:     encodeHamadryadHex(fpHashBytes),
		PermDist: permDist,
	}

	return &Signature{
		Fingerprint: fp,
		Challenge:   challenge,
		Response:    response,
		Proof: SporeProof{
			LockIns:        lockIns,
			NeighborCounts: neighborCounts,
			HexTrace:       hexTrace,
		},
		Commitment: commitment,
	}, nil
}

// quantizeLockIn maps a ratio in [0, ∞) to uint8 [0, 255].
// Uses the mapping: q = min(255, floor(ratio * 255)).
func quantizeLockIn(r ratio.Ratio) uint8 {
	if r.Denom == 0 {
		return 0
	}
	// r.Num / r.Denom * 255, clamped to [0, 255].
	v := r.Num * 255 / r.Denom
	if v > 255 {
		return 255
	}
	if v < 0 {
		return 0
	}
	return uint8(v)
}

// dequantizeLockIn reverses quantization: uint8 → ratio.
func dequantizeLockIn(q uint8) ratio.Ratio {
	return ratio.New(int64(q), 255)
}

// aggregateProof computes mean lock-in and mean neighbor count
// as fixed-point uint16 values.
func aggregateProof(p SporeProof) (meanLI, meanNC uint16) {
	n := len(p.LockIns)
	if n == 0 {
		return 0, 0
	}

	// Mean lock-in: exact rational sum, then scale to 0-65535.
	sumLI := ratio.Zero
	for _, li := range p.LockIns {
		if li.Denom != 0 {
			sumLI = sumLI.Add(li)
		}
	}
	avgLI := sumLI.Div(ratio.FromInt(int64(n)))
	liScaled := avgLI.ScaleInt(65535)
	if liScaled > 65535 {
		liScaled = 65535
	}
	if liScaled < 0 {
		liScaled = 0
	}
	meanLI = uint16(liScaled)

	// Mean neighbor count: integer sum, scale ×256 to fixed-point.
	var sumNC int
	for _, nc := range p.NeighborCounts {
		sumNC += nc
	}
	ncScaled := int64(sumNC) * 256 / int64(n)
	if ncScaled > 65535 {
		ncScaled = 65535
	}
	meanNC = uint16(ncScaled)

	return meanLI, meanNC
}

// encodeHamadryadHex converts a Hamadryad hash to a hex string.
func encodeHamadryadHex(p Hamadryad) string {
	const hex = "0123456789abcdef"
	out := make([]byte, HamBytes*2)
	for i, v := range p {
		out[i*2] = hex[v>>4]
		out[i*2+1] = hex[v&0x0f]
	}
	return string(out)
}