package crypto

import (
	"errors"
	"math/big"
	"sort"

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

// Compact V2 binary wire format for signatures.
//
// Eliminates redundant fields (full Challenge recomputed by verifier,
// Fingerprint known to verifier, aggregate proof recomputed from per-site
// data). Uses combinatorial number system for occupancy encoding and
// bitpacking for per-site data.
//
// A 16-byte truncated challenge hash is included for message binding:
// without it the signature would be message-independent and trivially
// replayable against different messages.
//
// Layout:
//   [1 byte]    Magic (0xD2)
//   [1 byte]    Flags: security_level[2] | tag_bits[1] | reserved[5]
//   [1 byte]    Occupied count K
//   [1 byte]    Tag count T
//   [variable]  Tag table: per tag, 1-byte length + UTF-8 string
//   [16 bytes]  Truncated challenge: first 16 bytes of Hamadryad(message)
//   [variable]  Occupancy set: combinatorial encoding of K-subset of [0,N)
//   [variable]  Per-site data: K × (tag_bits + 8 + 3) bits, padded
//
// Target size: ~144 bytes for N=256, K=56, 3 tags.

const compactV2Magic = 0xD2

// securityN maps security level bits to lattice dimension.
func securityN(level uint8) int {
	switch level {
	case 1:
		return 384
	case 2:
		return 512
	default:
		return 256
	}
}

// securityLevel returns the 2-bit security level for a given N.
func securityLevel(n int) uint8 {
	switch {
	case n >= 512:
		return 2
	case n >= 384:
		return 1
	default:
		return 0
	}
}

// --- Combinatorial number system ---
//
// Encodes a K-element subset of [0,N) as an integer in [0, C(N,K)).
// Rank of sorted indices {c_0 < c_1 < ... < c_{K-1}} = sum of C(c_i, i+1).
// This is the most compact possible encoding for a fixed-size subset.

// bigBinom computes C(n, k) as a *big.Int.
func bigBinom(n, k int) *big.Int {
	if k < 0 || k > n {
		return big.NewInt(0)
	}
	// Use the smaller of k and n-k for efficiency.
	if k > n-k {
		k = n - k
	}
	result := big.NewInt(1)
	for i := 0; i < k; i++ {
		result.Mul(result, big.NewInt(int64(n-i)))
		result.Div(result, big.NewInt(int64(i+1)))
	}
	return result
}

// subsetBytes returns the number of bytes needed to encode a K-subset of [0,N).
func subsetBytes(n, k int) int {
	b := bigBinom(n, k)
	bits := b.BitLen()
	if bits == 0 {
		return 1
	}
	return (bits + 7) / 8
}

// encodeSubset encodes sorted indices into a combinatorial rank (little-endian bytes).
func encodeSubset(indices []int, n int) []byte {
	k := len(indices)
	rank := new(big.Int)
	for i, c := range indices {
		rank.Add(rank, bigBinom(c, i+1))
	}
	nbytes := subsetBytes(n, k)
	buf := make([]byte, nbytes)
	rankBytes := rank.Bytes() // big-endian
	// Reverse to little-endian and copy into fixed-size buf.
	for i, j := 0, len(rankBytes)-1; i < len(rankBytes) && j >= 0; i, j = i+1, j-1 {
		buf[i] = rankBytes[j]
	}
	return buf
}

// decodeSubset decodes a combinatorial rank (little-endian bytes) back to sorted indices.
func decodeSubset(data []byte, n, k int) []int {
	// Convert little-endian to big.Int.
	reversed := make([]byte, len(data))
	for i, j := 0, len(data)-1; j >= 0; i, j = i+1, j-1 {
		reversed[i] = data[j]
	}
	rank := new(big.Int).SetBytes(reversed)

	indices := make([]int, k)
	// Greedy decode: from the highest index position downward.
	for i := k - 1; i >= 0; i-- {
		// Find the largest c such that C(c, i+1) <= rank.
		c := i // minimum possible value for position i
		for c+1 < n {
			b := bigBinom(c+1, i+1)
			if b.Cmp(rank) > 0 {
				break
			}
			c++
		}
		indices[i] = c
		rank.Sub(rank, bigBinom(c, i+1))
	}
	return indices
}

// --- Bitpacking ---

type bitWriter struct {
	buf    []byte
	bitPos int
}

func newBitWriter(capacity int) *bitWriter {
	return &bitWriter{buf: make([]byte, 0, capacity)}
}

func (w *bitWriter) writeBits(val uint32, nbits int) {
	for i := 0; i < nbits; i++ {
		byteIdx := w.bitPos / 8
		bitIdx := uint(w.bitPos % 8)
		for byteIdx >= len(w.buf) {
			w.buf = append(w.buf, 0)
		}
		if val&(1<<uint(i)) != 0 {
			w.buf[byteIdx] |= 1 << bitIdx
		}
		w.bitPos++
	}
}

func (w *bitWriter) bytes() []byte {
	// Ensure final partial byte is included.
	totalBytes := (w.bitPos + 7) / 8
	for len(w.buf) < totalBytes {
		w.buf = append(w.buf, 0)
	}
	return w.buf[:totalBytes]
}

type bitReader struct {
	buf    []byte
	bitPos int
}

func newBitReader(data []byte) *bitReader {
	return &bitReader{buf: data}
}

func (r *bitReader) readBits(nbits int) (uint32, error) {
	var val uint32
	for i := 0; i < nbits; i++ {
		byteIdx := r.bitPos / 8
		bitIdx := uint(r.bitPos % 8)
		if byteIdx >= len(r.buf) {
			return 0, errors.New("crypto: bitstream truncated")
		}
		if r.buf[byteIdx]&(1<<bitIdx) != 0 {
			val |= 1 << uint(i)
		}
		r.bitPos++
	}
	return val, nil
}

// --- MarshalV2 / UnmarshalSignatureV2 ---

// MarshalV2 encodes a Signature into the compact V2 binary wire format.
// The caller must supply the lattice dimension N (from Params).
func (s *Signature) MarshalV2(n int) ([]byte, error) {
	if s == nil {
		return nil, errors.New("crypto: nil signature")
	}

	// Collect occupied sites in index order.
	type siteData struct {
		index int
		tag   string
		proj  uint8
		perm  uint8
	}
	var sites []siteData
	for _, site := range s.Response {
		if site.Occupied {
			sites = append(sites, siteData{
				index: int(site.Index),
				tag:   site.TypeTag,
				proj:  site.Projection,
				perm:  site.Perm,
			})
		}
	}
	sort.Slice(sites, func(i, j int) bool { return sites[i].index < sites[j].index })

	k := len(sites)
	if k == 0 {
		return nil, errors.New("crypto: no occupied sites")
	}
	if k > 255 {
		return nil, errors.New("crypto: too many occupied sites for V2")
	}

	// Build tag table.
	tagIndex := make(map[string]uint8)
	var tagTable []string
	for _, sd := range sites {
		if _, ok := tagIndex[sd.tag]; !ok {
			if len(tagTable) >= 8 {
				return nil, errors.New("crypto: too many distinct tags for V2")
			}
			tagIndex[sd.tag] = uint8(len(tagTable))
			tagTable = append(tagTable, sd.tag)
		}
	}

	tagCount := len(tagTable)
	tagBitsFlag := uint8(0) // 2-bit tags
	tagBits := 2
	if tagCount > 4 {
		tagBitsFlag = 1 // 3-bit tags
		tagBits = 3
	}

	// Header: magic + flags + K + tagcount + tag table.
	flags := securityLevel(n) | (tagBitsFlag << 2)
	var header []byte
	header = append(header, compactV2Magic)
	header = append(header, flags)
	header = append(header, uint8(k))
	header = append(header, uint8(tagCount))
	for _, tag := range tagTable {
		if len(tag) > 255 {
			return nil, errors.New("crypto: tag name too long")
		}
		header = append(header, uint8(len(tag)))
		header = append(header, []byte(tag)...)
	}

	// Truncated challenge: first 16 bytes of Hamadryad(message).
	// Provides 128-bit message binding.
	header = append(header, s.Challenge[:16]...)

	// Combinatorial occupancy encoding.
	indices := make([]int, k)
	for i, sd := range sites {
		indices[i] = sd.index
	}
	occBytes := encodeSubset(indices, n)
	nbytes := subsetBytes(n, k)
	// Prefix with 1-byte length so decoder knows how many bytes to read.
	var occSection []byte
	occSection = append(occSection, uint8(nbytes))
	occSection = append(occSection, occBytes...)

	// Bitpacked per-site data.
	bitsPerSite := tagBits + 8 + 3
	bw := newBitWriter((k*bitsPerSite + 7) / 8)
	for _, sd := range sites {
		bw.writeBits(uint32(tagIndex[sd.tag]), tagBits)
		bw.writeBits(uint32(sd.proj), 8)
		bw.writeBits(uint32(sd.perm), 3)
	}
	siteBytes := bw.bytes()

	// Assemble.
	total := len(header) + len(occSection) + len(siteBytes)
	buf := make([]byte, 0, total)
	buf = append(buf, header...)
	buf = append(buf, occSection...)
	buf = append(buf, siteBytes...)

	return buf, nil
}

// UnmarshalSignatureV2 decodes a compact V2 binary signature.
// Returns the reconstructed Signature with synthetic proof fields.
func UnmarshalSignatureV2(data []byte) (*Signature, int, error) {
	if len(data) < 4 {
		return nil, 0, errors.New("crypto: V2 signature too short")
	}
	pos := 0

	// Magic.
	if data[pos] != compactV2Magic {
		return nil, 0, errors.New("crypto: bad V2 magic")
	}
	pos++

	// Flags.
	flags := data[pos]
	pos++
	level := flags & 0x03
	tagBitsFlag := (flags >> 2) & 0x01
	n := securityN(level)

	tagBits := 2
	if tagBitsFlag == 1 {
		tagBits = 3
	}

	// Occupied count.
	k := int(data[pos])
	pos++
	if k == 0 {
		return nil, 0, errors.New("crypto: zero occupied sites")
	}

	// Tag count.
	if pos >= len(data) {
		return nil, 0, errors.New("crypto: truncated header")
	}
	tagCount := int(data[pos])
	pos++

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

	// Truncated challenge (16 bytes).
	const challengeTruncLen = 16
	if pos+challengeTruncLen > len(data) {
		return nil, 0, errors.New("crypto: truncated challenge")
	}
	var challengeTrunc [challengeTruncLen]byte
	copy(challengeTrunc[:], data[pos:pos+challengeTruncLen])
	pos += challengeTruncLen

	// Occupancy set.
	if pos >= len(data) {
		return nil, 0, errors.New("crypto: truncated occupancy")
	}
	occLen := int(data[pos])
	pos++
	if pos+occLen > len(data) {
		return nil, 0, errors.New("crypto: truncated occupancy data")
	}
	indices := decodeSubset(data[pos:pos+occLen], n, k)
	pos += occLen

	// Per-site bitpacked data.
	bitsPerSite := tagBits + 8 + 3
	totalBits := k * bitsPerSite
	siteDataBytes := (totalBits + 7) / 8
	if pos+siteDataBytes > len(data) {
		return nil, 0, errors.New("crypto: truncated site data")
	}

	br := newBitReader(data[pos : pos+siteDataBytes])
	response := make([]SiteMark, k)
	lockIns := make([]ratio.Ratio, k)
	neighborCounts := make([]int, k)
	hexTrace := make([]state.Hexagram, k)

	for i := range k {
		tagIdx, err := br.readBits(tagBits)
		if err != nil {
			return nil, 0, err
		}
		proj, err := br.readBits(8)
		if err != nil {
			return nil, 0, err
		}
		perm, err := br.readBits(3)
		if err != nil {
			return nil, 0, err
		}

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

		response[i] = SiteMark{
			Index:      uint64(indices[i]),
			Occupied:   true,
			TypeTag:    tag,
			Projection: uint8(proj),
			Perm:       uint8(perm),
			LockIn:     ratio.New(1, 1),
		}
		lockIns[i] = ratio.New(1, 1)
		neighborCounts[i] = 1
		hexTrace[i] = 0
	}
	pos += siteDataBytes

	// Reconstruct permutation distribution for fingerprint.
	permDist := [6]int{}
	for _, site := range response {
		if site.Perm < 6 {
			permDist[site.Perm]++
		}
	}

	// Store truncated challenge in the Challenge field (first 16 bytes).
	var challenge Hamadryad
	copy(challenge[:], challengeTrunc[:])

	sig := &Signature{
		Fingerprint: SporeFingerprint{
			PermDist: permDist,
		},
		Challenge: challenge,
		Response:  response,
		Proof: SporeProof{
			LockIns:        lockIns,
			NeighborCounts: neighborCounts,
			HexTrace:       hexTrace,
		},
	}

	return sig, pos, nil
}

// V2Size returns the wire size of a V2 encoding without performing the full marshal.
func (s *Signature) V2Size(n int) int {
	if s == nil {
		return 0
	}

	var k int
	tags := make(map[string]bool)
	for _, site := range s.Response {
		if site.Occupied {
			k++
			tags[site.TypeTag] = true
		}
	}

	tagCount := len(tags)
	tagBits := 2
	if tagCount > 4 {
		tagBits = 3
	}

	// Header: magic(1) + flags(1) + K(1) + tagcount(1) = 4
	headerSize := 4
	tagTableSize := 0
	for tag := range tags {
		tagTableSize += 1 + len(tag)
	}

	// Occupancy: 1-byte length prefix + combinatorial bytes.
	occSize := 1 + subsetBytes(n, k)

	// Per-site data.
	bitsPerSite := tagBits + 8 + 3
	siteDataSize := (k*bitsPerSite + 7) / 8

	// 16-byte truncated challenge.
	challengeSize := 16

	return headerSize + tagTableSize + challengeSize + occSize + siteDataSize
}