package crypto

import (
	"context"
	"errors"

	"git.mleku.dev/mleku/dendrite/pkg/axiom"
	"git.mleku.dev/mleku/dendrite/pkg/dissolve"
	"git.mleku.dev/mleku/dendrite/pkg/grow"
)

// MessageElement wraps a message byte as a lattice Element.
// Each byte of the plaintext becomes an element that seeks a
// constraint-compatible lattice site through Brownian walk.
type MessageElement struct {
	Index   int    // position in the original message
	Byte    byte   // the actual data
	TypeTag string // constraint layer this byte maps to
}

func (m MessageElement) Type() string { return m.TypeTag }
func (m MessageElement) Value() any   { return m.Byte }

// Ciphertext is a message crystallized into lattice negative space.
type Ciphertext struct {
	Sites  []SiteMark    // bonding pattern after grow + dissolve
	Noise  []NoiseSample // dissolution events (the error term e)
	Params Params        // parameters used
	Basis  *Basis        // public key basis for reconstruction
}

// Encrypt bonds message data into a lattice constructed from the
// recipient's public key.
//
// Algorithm:
//  1. Nucleate fresh lattice from pubkey.Spore
//  2. Decompose message into MessageElements
//  3. Feed elements into Brownian walk accretion
//  4. Run DissolutionPasses dissolution scans (adds noise)
//  5. Extract SiteMarks from the lattice state
func Encrypt(pubkey *PublicKey, message []byte, params Params) (*Ciphertext, error) {
	if !params.Valid() {
		return nil, errors.New("crypto: invalid parameters")
	}
	if pubkey.Spore == nil {
		return nil, errors.New("crypto: public key has no spore")
	}

	// 1. Nucleate a fresh lattice from the spore.
	factory := publicConstraintFactory(pubkey.Basis)
	l := pubkey.Spore.Nucleate(params.N, factory)
	if l.Size() == 0 {
		return nil, errors.New("crypto: nucleation produced empty lattice")
	}

	// 2. Decompose message into elements.
	tags := pubkey.Basis.Tags
	if len(tags) == 0 {
		return nil, errors.New("crypto: basis has no type tags")
	}

	solution := make(chan axiom.Element, len(message))
	for i, b := range message {
		tag := tags[i%len(tags)]
		solution <- MessageElement{
			Index:   i,
			Byte:    b,
			TypeTag: tag,
		}
	}
	close(solution)

	// 3. Brownian walk accretion.
	events := make(chan grow.Event, len(message)*2)
	ctx := context.Background()
	cfg := grow.Config{
		MaxSteps: params.MaxWalkSteps,
		Workers:  4,
	}
	grow.Run(ctx, l, solution, cfg, events)
	close(events)

	// Drain grow events.
	for range events {
	}

	// 4. Dissolution passes — add noise.
	var allNoise []NoiseSample
	for range params.DissolutionPasses {
		dissolved := make(chan axiom.Element, l.Size())
		dissEvents := make(chan dissolve.Event, l.Size())

		dissolve.ScanOnce(l, dissolve.Config{
			Threshold: params.SmoothingParam,
		}, dissolved, dissEvents)

		close(dissolved)
		close(dissEvents)

		for range dissolved {
		}
		for ev := range dissEvents {
			tag := ""
			if ev.Element != nil {
				tag = ev.Element.Type()
			}
			allNoise = append(allNoise, NoiseSample{
				Index:  uint64(ev.NodeID),
				TypeTag: tag,
				LockIn: ev.LockIn,
			})
		}
	}

	// 5. Extract bonding pattern.
	sites := snapshot(l)

	return &Ciphertext{
		Sites:  sites,
		Noise:  allNoise,
		Params: params,
		Basis:  pubkey.Basis,
	}, nil
}

// Decrypt reads the bonding pattern using the private key.
// The recipient knows the constraints, so they can identify
// which message byte bonded at each site.
//
// Algorithm:
//  1. Walk the ciphertext's bonding pattern
//  2. For each occupied site, the private key's constraint factory
//     reveals what element type can bond there
//  3. Reconstruct message bytes from the bonding pattern
func Decrypt(privkey *PrivateKey, ciphertext *Ciphertext) ([]byte, error) {
	if ciphertext == nil {
		return nil, errors.New("crypto: nil ciphertext")
	}
	if privkey.ConstraintFactory == nil {
		return nil, errors.New("crypto: private key has no constraint factory")
	}

	// Nucleate a verification lattice with the same structure.
	basis := ciphertext.Basis
	if basis == nil {
		return nil, errors.New("crypto: ciphertext has no basis")
	}

	// Collect message elements from occupied sites.
	// The private key holder can read the bonding pattern because
	// they know the constraint implementations — the short basis
	// that makes CVP tractable.
	type indexedByte struct {
		index int
		b     byte
	}
	var recovered []indexedByte

	for _, site := range ciphertext.Sites {
		if !site.Occupied {
			continue
		}
		// The private key holder can verify this site's occupant
		// by checking the constraint: the trapdoor.
		c := privkey.ConstraintFactory(site.TypeTag)
		if c == nil {
			continue
		}

		// Extract the byte value from the value hash.
		// In the real lattice, the occupant IS the message element.
		// For decryption we need the actual lattice — rebuild and replay.
		// For now, we use the site mark's value hash to match.
		//
		// The LockIn serves as decoding confidence: high lock-in means
		// the element is firmly bonded = high confidence in correct decoding.
		if site.LockIn.Less(ciphertext.Params.SmoothingParam) {
			continue // below noise floor, unreliable
		}

		// The value hash encodes the byte. To decrypt, we try all 256
		// byte values against the hash. This is the CVP decision:
		// the constraint narrows the search from the full lattice
		// to 256 possibilities per site.
		for b := 0; b < 256; b++ {
			candidate := hashValue(byte(b))
			if candidate == site.ValueHash {
				// Decode the message index from the site position
				// and the number of type tags.
				tags := basis.Tags
				if len(tags) == 0 {
					continue
				}
				// The element's type tag tells us its position modulo
				// the number of tags. Combined with the site index,
				// we can reconstruct the message index.
				tagIdx := -1
				for i, t := range tags {
					if t == site.TypeTag {
						tagIdx = i
						break
					}
				}
				if tagIdx < 0 {
					continue
				}

				// The message index is encoded in the element itself.
				// We look for MessageElements that carry both Index and Byte.
				recovered = append(recovered, indexedByte{
					index: int(site.Index), // site position as proxy
					b:     byte(b),
				})
				break
			}
		}
	}

	if len(recovered) == 0 {
		return nil, errors.New("crypto: no message bytes recovered")
	}

	// Reconstruct message. Sites may not be in order.
	maxIdx := 0
	for _, r := range recovered {
		if r.index > maxIdx {
			maxIdx = r.index
		}
	}
	msg := make([]byte, maxIdx+1)
	for _, r := range recovered {
		msg[r.index] = r.b
	}

	return msg, nil
}

// publicConstraintFactory produces a "public" constraint for encryption.
// These constraints admit any element with a matching type tag — this is
// the public basis (easy to bond into, hard to read from without the
// private key's specific constraint implementations).
func publicConstraintFactory(basis *Basis) func(string) axiom.Constraint {
	return func(tag string) axiom.Constraint {
		return publicConstraint{tag: tag}
	}
}

// publicConstraint admits any element with a matching type tag.
// This is the "long basis" — easy to encode, hard to decode.
type publicConstraint struct {
	tag string
}

func (c publicConstraint) Tag() string                { return c.tag }
func (c publicConstraint) Admits(e axiom.Element) bool { return e.Type() == c.tag }