package crypto

import (
	"context"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"

	"golang.org/x/crypto/chacha20"

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

// ShadowEncrypt encrypts plaintext using the binary/decimal shadow channel.
//
// The plaintext is converted to hexagram tokens, each independently
// permuted by an S_3 pair derived from the epoch Phase() at that token's
// position. The permuted tokens are then crystallized into a lattice
// nucleated from the recipient's public key.
//
// The encryption is structurally dependent on the binary/decimal
// incommensurability: the inner trigram rotates by the binary clock
// component of the phase, and the outer trigram rotates by the decimal
// clock component. Neither projection alone recovers the original.
//
// Algorithm:
//  1. Convert plaintext to hexagram tokens (state.EncodeBytes)
//  2. For each token, compute phase-derived S_3 pair and apply
//  3. Bond shadow elements into recipient's lattice
//  4. Run dissolution passes (add noise)
//  5. XOR bonding data with phase-derived keystream
//  6. Extract ciphertext
func ShadowEncrypt(
	pubkey *PublicKey,
	plaintext []byte,
	params Params,
	ep epoch.Epoch,
	nonce []byte,
) (*ShadowCiphertext, error) {
	if !params.Valid() {
		return nil, errors.New("crypto: invalid parameters")
	}
	if pubkey == nil || pubkey.Spore == nil {
		return nil, errors.New("crypto: public key has no spore")
	}
	if len(nonce) == 0 {
		return nil, errors.New("crypto: empty nonce")
	}

	// 1. Hexagram encoding: 3 bytes → 4 tokens (6-bit each).
	tokens := state.EncodeBytes(plaintext)
	if len(tokens) == 0 && len(plaintext) > 0 {
		return nil, errors.New("crypto: hexagram encoding failed")
	}

	// 2. Phase-dependent shadow decomposition.
	encrypted := make([]state.Hexagram, len(tokens))
	for i, tok := range tokens {
		binP, decP := PhasePerms(ep, i)
		encrypted[i] = ShadowDecompose(tok, binP, decP)
	}

	// 3. Build shadow elements for lattice bonding.
	factory := publicConstraintFactory(pubkey.Basis)
	l := pubkey.Spore.Nucleate(params.N, factory)
	if l.Size() == 0 {
		return nil, errors.New("crypto: nucleation produced empty lattice")
	}

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

	solution := make(chan axiom.Element, len(encrypted))
	for i, enc := range encrypted {
		vertex, key := PhaseProjection(ep, i)
		binP, _ := PhasePerms(ep, i)
		solution <- ShadowElement{
			Index:      i,
			Encrypted:  enc,
			ProjVertex: vertex,
			ProjKey:    key,
			ProjPath:   uint16(i),
			PermIdx:    uint8(binP),
			TypeTag:    tags[i%len(tags)],
		}
	}
	close(solution)

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

	// 5. 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,
			})
		}
	}

	// 6. Extract bonding pattern and apply keystream masking.
	sites := snapshot(l)

	// Generate per-message nonce hash.
	var nonceBuf []byte
	nonceBuf = append(nonceBuf, []byte("shadow-keystream-v1")...)
	nonceBuf = append(nonceBuf, []byte(pubkey.SporeHash)...)
	nonceBuf = append(nonceBuf, []byte(ep.String())...)
	nonceBuf = append(nonceBuf, nonce...)
	nonceHash := Hash(nonceBuf)

	// XOR value hashes with keystream for additional masking.
	// The keystream is keyed by the site's lattice node ID, not its
	// position in the slice, so the masking survives reordering.
	for i := range sites {
		if sites[i].Occupied {
			block := keystreamForSite(nonceHash, sites[i].Index)
			sites[i].ValueHash = xorHamadryad(sites[i].ValueHash, block)
		}
	}

	return &ShadowCiphertext{
		Sites:      sites,
		Noise:      allNoise,
		Params:     params,
		Basis:      pubkey.Basis,
		Nonce:      nonceHash,
		EpochDec:   ep.DecExp,
		EpochBin:   ep.BinExp,
		TokenCount: len(tokens),
		OrigLen:    len(plaintext),
	}, nil
}

// ShadowDecrypt recovers plaintext from a shadow ciphertext using the
// private key and the matching epoch.
//
// Algorithm:
//  1. Unmask the bonding pattern using the nonce keystream
//  2. Read bonded values via private key (CVP trapdoor)
//  3. Recompute phase schedule from epoch
//  4. Apply inverse S_3 pairs to recover original hexagrams
//  5. Decode hexagrams to bytes
func ShadowDecrypt(
	privkey *PrivateKey,
	ct *ShadowCiphertext,
	ep epoch.Epoch,
) ([]byte, error) {
	if ct == nil {
		return nil, errors.New("crypto: nil shadow ciphertext")
	}
	if privkey == nil || privkey.ConstraintFactory == nil {
		return nil, errors.New("crypto: private key has no constraint factory")
	}

	// 1. Unmask value hashes.
	// The keystream is keyed by lattice node ID, matching the encrypt side.
	sites := make([]SiteMark, len(ct.Sites))
	copy(sites, ct.Sites)
	for i := range sites {
		if sites[i].Occupied {
			block := keystreamForSite(ct.Nonce, sites[i].Index)
			sites[i].ValueHash = xorHamadryad(sites[i].ValueHash, block)
		}
	}

	// 2. Recover encrypted hexagram values from the bonding pattern.
	// The private key holder can verify each site's occupant by
	// testing all 64 hexagram values against the value hash.
	type indexedHex struct {
		index int
		hex   state.Hexagram
	}
	var recovered []indexedHex

	for _, site := range sites {
		if !site.Occupied {
			continue
		}
		if site.LockIn.Less(ct.Params.SmoothingParam) {
			continue // below noise floor
		}

		// Test all 64 hexagram values against the stored hash.
		for h := range 64 {
			candidate := hashValue(state.Hexagram(h))
			if candidate == site.ValueHash {
				// Use ProjPath as the token stream index.
				// ShadowEncrypt stores the token index as ProjectionPath
				// in each ShadowElement.
				recovered = append(recovered, indexedHex{
					index: int(site.ProjPath),
					hex:   state.Hexagram(h),
				})
				break
			}
		}
	}

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

	// 3. Reconstruct token stream, applying inverse phase permutations.
	tokens := make([]state.Hexagram, ct.TokenCount)
	for _, r := range recovered {
		tokenIdx := r.index
		if tokenIdx >= ct.TokenCount {
			continue
		}

		// Apply inverse phase permutations.
		binP, decP := PhasePerms(ep, tokenIdx)
		tokens[tokenIdx] = ShadowRecompose(r.hex, binP, decP)
	}

	// 4. Decode hexagrams back to bytes.
	return state.DecodeHexagrams(tokens, ct.OrigLen), nil
}

// keystreamForSite generates a Hamadryad-sized keystream block for a
// specific lattice site using ChaCha20.
//
// The ChaCha20 key is derived from the nonce via a single Hamadryad hash
// (one-time key derivation, not used as a PRF). The site index selects
// the ChaCha20 block counter, giving each site an independent keystream
// block without relying on SWIFFT's unproven PRF properties.
func keystreamForSite(nonce Hamadryad, siteIndex uint64) Hamadryad {
	// Derive ChaCha20 key from nonce (one-time hash, proven secure).
	keyHash := Hash(append([]byte("shadow-chacha20-key"), nonce[:]...))
	var key [chacha20.KeySize]byte
	copy(key[:], keyHash[:chacha20.KeySize])

	// Derive ChaCha20 nonce from site index.
	var chachaNonce [chacha20.NonceSize]byte
	binary.LittleEndian.PutUint64(chachaNonce[:8], siteIndex)

	cipher, _ := chacha20.NewUnauthenticatedCipher(key[:], chachaNonce[:])
	var result Hamadryad
	cipher.XORKeyStream(result[:], result[:]) // XOR zeros = raw keystream
	return result
}

// xorHamadryad XORs two Hamadryad values byte-by-byte.
func xorHamadryad(a, b Hamadryad) Hamadryad {
	var result Hamadryad
	for i := range HamBytes {
		result[i] = a[i] ^ b[i]
	}
	return result
}

// generateNonce creates a cryptographically random nonce.
func generateNonce() []byte {
	nonce := make([]byte, 32)
	if _, err := rand.Read(nonce); err != nil {
		panic(fmt.Sprintf("crypto/rand: %v", err))
	}
	return nonce
}