package crypto

import (
	"errors"
	"sort"

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

// SharedSecret is the result of a key exchange between two peers.
type SharedSecret struct {
	// CommonTags are the constraint type tags present in both lattices.
	CommonTags []string

	// Secret is the derived Hamadryad shared key material (56 bytes).
	Secret Hamadryad

	// Confidence measures how much structural overlap exists.
	// Higher confidence = more shared constraint structure = stronger key.
	Confidence ratio.Ratio
}

// Exchange performs a Diffie-Hellman-like key exchange using spores.
//
// The protocol:
//  1. Find the constraint types common to both spores
//  2. For each common type, compute the structural overlap:
//     occupancy rates, connectivity, permutation distributions
//  3. Hash the common structure to derive the shared secret
//
// This works because:
//   - Both parties derive the same set of common tags deterministically
//   - The structural overlap (occupancy * connectivity) is symmetric
//   - An eavesdropper sees only the spores but cannot determine which
//     constraints would admit elements without the factory (private key)
func Exchange(
	ownSpore *spore.Spore,
	peerSpore *spore.Spore,
) (*SharedSecret, error) {
	if ownSpore == nil || peerSpore == nil {
		return nil, errors.New("crypto: nil spore in exchange")
	}

	// Find common tags.
	ownTags := tagSet(ownSpore.TypeSignature)
	peerTags := tagSet(peerSpore.TypeSignature)

	var common []string
	for tag := range ownTags {
		if peerTags[tag] {
			common = append(common, tag)
		}
	}
	sort.Strings(common)

	if len(common) == 0 {
		return &SharedSecret{
			Confidence: ratio.Zero,
		}, nil
	}

	// Compute confidence: |common| / max(|own|, |peer|).
	maxTags := max(len(ownTags), len(peerTags))
	confidence := ratio.New(int64(len(common)), int64(maxTags))

	// Derive shared secret from common structure.
	// The key material is: sorted common tags + their occupancy rates
	// from both spores + connectivity ratios, hashed through Hamadryad.
	var buf []byte
	// Domain separation.
	buf = append(buf, []byte("dendrite-exchange-v1")...)

	for _, tag := range common {
		buf = append(buf, []byte(tag)...)

		// Occupancy rates for this tag from both spores.
		// Use min/max ordering for commutativity — both parties
		// must derive the same hash regardless of who is "own" vs "peer".
		ownCount := tagCountLookup(ownSpore.TypeSignature, tag)
		ownRate := ratio.New(int64(ownCount), int64(ownSpore.TotalNodes))
		peerCount := tagCountLookup(peerSpore.TypeSignature, tag)
		peerRate := ratio.New(int64(peerCount), int64(peerSpore.TotalNodes))

		minRate := ratio.Min(ownRate, peerRate)
		maxRate := ratio.Max(ownRate, peerRate)
		buf = append(buf, []byte(minRate.String())...)
		buf = append(buf, []byte(maxRate.String())...)

		// Connectivity overlap — also commutative via min/max.
		ownConn := tagRatioLookup(ownSpore.Connectivity, tag)
		peerConn := tagRatioLookup(peerSpore.Connectivity, tag)

		minConn := ratio.Min(ownConn, peerConn)
		maxConn := ratio.Max(ownConn, peerConn)
		buf = append(buf, []byte(minConn.String())...)
		buf = append(buf, []byte(maxConn.String())...)
	}

	secret := Hash(buf)

	return &SharedSecret{
		CommonTags: common,
		Secret:     secret,
		Confidence: confidence,
	}, nil
}

// DeriveKey produces a fixed-length key from a SharedSecret.
// Uses Hamadryad with domain separation for key derivation.
func DeriveKey(ss *SharedSecret, context string, keyLen int) []byte {
	var buf []byte
	buf = append(buf, []byte("dendrite-derive-")...)
	buf = append(buf, []byte(context)...)
	buf = append(buf, ss.Secret[:]...)

	derived := Hash(buf)
	result := derived[:]

	// Extend if needed by iterating.
	for len(result) < keyLen {
		buf = buf[:0]
		buf = append(buf, result...)
		buf = append(buf, []byte(context)...)
		next := Hash(buf)
		result = append(result, next[:]...)
	}

	return result[:keyLen]
}

// tagSet returns a set of tags from a TypeSignature.
func tagSet(tc []spore.TagCount) map[string]bool {
	m := make(map[string]bool, len(tc))
	for _, t := range tc {
		m[t.Tag] = true
	}
	return m
}

// tagCountLookup returns the count for a tag in a TypeSignature.
func tagCountLookup(tc []spore.TagCount, tag string) int {
	for _, t := range tc {
		if t.Tag == tag {
			return t.Count
		}
	}
	return 0
}

// tagRatioLookup returns the ratio for a tag in a Connectivity slice.
func tagRatioLookup(tr []spore.TagRatio, tag string) ratio.Ratio {
	for _, t := range tr {
		if t.Tag == tag {
			return t.Value
		}
	}
	return ratio.Zero
}