package graph

import (
	"container/heap"
	"math/bits"
	"sort"

	"github.com/local/glpi-neural-brain/internal/model"
)

// ClusterSearchStats separates cheap semantic-hash work from expensive exact
// cosine work. CoarseComparisons are Hamming-distance operations; only
// ExactComparisons execute a full embedding dot product.
type ClusterSearchStats struct {
	IndexedNodes      int `json:"indexed_nodes"`
	CoarseComparisons int `json:"coarse_comparisons"`
	ExactComparisons  int `json:"exact_comparisons"`
	CandidatePool     int `json:"candidate_pool"`
	HashBits          int `json:"hash_bits"`
	HashTables        int `json:"hash_tables"`
}

type semanticHashEntry struct {
	node       model.Node
	vector     []float32
	signatures []uint64
}

type coarseCandidate struct {
	index    int
	distance int
}

type coarseMaxHeap []coarseCandidate

func (h coarseMaxHeap) Len() int { return len(h) }
func (h coarseMaxHeap) Less(i, j int) bool {
	if h[i].distance == h[j].distance {
		return h[i].index > h[j].index
	}
	return h[i].distance > h[j].distance
}
func (h coarseMaxHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *coarseMaxHeap) Push(x any)   { *h = append(*h, x.(coarseCandidate)) }
func (h *coarseMaxHeap) Pop() any {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[:n-1]
	return x
}

func normalizeClusterConfig(hashBits, hashTables, candidateLimit int) (int, int, int) {
	if hashBits < 8 {
		hashBits = 8
	}
	if hashBits > 63 {
		hashBits = 63
	}
	if hashTables < 1 {
		hashTables = 1
	}
	if hashTables > 4 {
		hashTables = 4
	}
	if candidateLimit < 16 {
		candidateLimit = 16
	}
	return hashBits, hashTables, candidateLimit
}

// sparseSemanticHash is a deterministic sparse random-projection hash. Each
// bit samples only six embedding dimensions, so building the coarse index is
// orders of magnitude cheaper than N full 768-dimensional cosine products.
func sparseSemanticHash(v []float32, table, hashBits int) uint64 {
	if len(v) == 0 {
		return 0
	}
	var signature uint64
	seedBase := uint64(0x9e3779b97f4a7c15) ^ uint64(table+1)*0xbf58476d1ce4e5b9
	for bit := 0; bit < hashBits; bit++ {
		seed := mix64(seedBase ^ uint64(bit+1)*0x94d049bb133111eb)
		var sum float32
		for sample := 0; sample < 6; sample++ {
			seed = mix64(seed + uint64(sample+1)*0x9e3779b97f4a7c15)
			idx := int(seed % uint64(len(v)))
			if seed&(1<<63) != 0 {
				sum -= v[idx]
			} else {
				sum += v[idx]
			}
		}
		if sum >= 0 {
			signature |= 1 << bit
		}
	}
	return signature
}

func mix64(x uint64) uint64 {
	x ^= x >> 30
	x *= 0xbf58476d1ce4e5b9
	x ^= x >> 27
	x *= 0x94d049bb133111eb
	x ^= x >> 31
	return x
}

func signaturesFor(v []float32, hashBits, hashTables int) []uint64 {
	out := make([]uint64, hashTables)
	for table := 0; table < hashTables; table++ {
		out[table] = sparseSemanticHash(v, table, hashBits)
	}
	return out
}

func signatureDistance(a, b []uint64) int {
	n := len(a)
	if len(b) < n {
		n = len(b)
	}
	distance := 0
	for i := 0; i < n; i++ {
		distance += bits.OnesCount64(a[i] ^ b[i])
	}
	return distance
}

func pushBestCoarse(h *coarseMaxHeap, candidate coarseCandidate, limit int) {
	if h.Len() < limit {
		heap.Push(h, candidate)
		return
	}
	worst := (*h)[0]
	if candidate.distance < worst.distance || (candidate.distance == worst.distance && candidate.index < worst.index) {
		heap.Pop(h)
		heap.Push(h, candidate)
	}
}

// NextPairClusteredScopedDepth is the low-resource alternative to
// NextPairScopedDepth. It keeps the rotating anchor behaviour but ranks the
// corpus with cheap semantic hashes and computes exact cosine only on top-K.
func (s *Store) NextPairClusteredScopedDepth(min float64, anchorLimit int, filter NodeFilter, maxAIDepth, hashBits, hashTables, candidateLimit int) (model.Node, model.Node, float64, bool, ClusterSearchStats) {
	hashBits, hashTables, candidateLimit = normalizeClusterConfig(hashBits, hashTables, candidateLimit)
	stats := ClusterSearchStats{HashBits: hashBits, HashTables: hashTables}

	s.mu.Lock()
	defer s.mu.Unlock()

	entries := make([]semanticHashEntry, 0, len(s.nodes))
	for _, n := range s.nodes {
		if n.Kind != "knowledge" && n.Kind != "ai-think" {
			continue
		}
		if !filter.Matches(n) {
			continue
		}
		if n.Kind == "ai-think" && maxAIDepth > 0 && graphNodeGenerationDepth(n) >= maxAIDepth {
			continue
		}
		v, ok := s.vectors[n.ID]
		if !ok || len(v) == 0 {
			continue
		}
		entries = append(entries, semanticHashEntry{node: n, vector: v})
	}
	if len(entries) < 2 {
		return model.Node{}, model.Node{}, 0, false, stats
	}
	sort.Slice(entries, func(i, j int) bool { return entries[i].node.ID < entries[j].node.ID })
	for i := range entries {
		entries[i].signatures = signaturesFor(entries[i].vector, hashBits, hashTables)
	}
	stats.IndexedNodes = len(entries)
	if anchorLimit <= 0 || anchorLimit > len(entries) {
		anchorLimit = len(entries)
	}
	start := s.pairCursor % len(entries)

	anchorIDs := make(map[string]struct{}, anchorLimit)
	for step := 0; step < anchorLimit; step++ {
		anchorIDs[entries[(start+step)%len(entries)].node.ID] = struct{}{}
	}
	// Only materialise blocked neighbours for current anchors. This avoids a
	// full pair-key allocation for every one of the ~180k graph edges.
	blocked := make(map[string]map[string]struct{}, anchorLimit)
	for _, edge := range s.edges {
		if _, ok := anchorIDs[edge.Source]; ok {
			if blocked[edge.Source] == nil {
				blocked[edge.Source] = map[string]struct{}{}
			}
			blocked[edge.Source][edge.Target] = struct{}{}
		}
		if _, ok := anchorIDs[edge.Target]; ok {
			if blocked[edge.Target] == nil {
				blocked[edge.Target] = map[string]struct{}{}
			}
			blocked[edge.Target][edge.Source] = struct{}{}
		}
	}

	best := -1.0
	var bestA, bestB model.Node
	evaluated := make(map[string]struct{}, anchorLimit*candidateLimit)
	for step := 0; step < anchorLimit; step++ {
		i := (start + step) % len(entries)
		left := entries[i]
		h := &coarseMaxHeap{}
		heap.Init(h)
		for j := range entries {
			if i == j {
				continue
			}
			right := entries[j]
			if left.node.Kind == "ai-think" && right.node.Kind == "ai-think" {
				continue
			}
			if neighbors := blocked[left.node.ID]; neighbors != nil {
				if _, exists := neighbors[right.node.ID]; exists {
					continue
				}
			}
			stats.CoarseComparisons++
			pushBestCoarse(h, coarseCandidate{index: j, distance: signatureDistance(left.signatures, right.signatures)}, candidateLimit)
		}
		candidates := make([]coarseCandidate, h.Len())
		for k := len(candidates) - 1; k >= 0; k-- {
			candidates[k] = heap.Pop(h).(coarseCandidate)
		}
		stats.CandidatePool += len(candidates)
		for _, candidate := range candidates {
			right := entries[candidate.index]
			key := pairKey(left.node.ID, right.node.ID)
			if _, seen := evaluated[key]; seen {
				continue
			}
			evaluated[key] = struct{}{}
			if len(left.vector) != len(right.vector) {
				continue
			}
			stats.ExactComparisons++
			score := cosine32(left.vector, right.vector)
			if score >= min && score > best {
				best, bestA, bestB = score, left.node, right.node
			}
		}
	}
	s.pairCursor = (start + anchorLimit) % len(entries)
	return bestA, bestB, best, best >= 0, stats
}

// SimilarClusteredFiltered performs approximate nearest-neighbour retrieval by
// semantic hash followed by exact cosine on a bounded shortlist.
func (s *Store) SimilarClusteredFiltered(query []float64, limit, candidateLimit int, filter NodeFilter, maxAIDepth, hashBits, hashTables int) ([]model.Hit, ClusterSearchStats) {
	hashBits, hashTables, candidateLimit = normalizeClusterConfig(hashBits, hashTables, candidateLimit)
	if limit < 1 {
		limit = 1
	}
	if candidateLimit < limit {
		candidateLimit = limit
	}
	stats := ClusterSearchStats{HashBits: hashBits, HashTables: hashTables}
	q := make([]float32, len(query))
	for i, value := range query {
		q[i] = float32(value)
	}
	qsig := signaturesFor(q, hashBits, hashTables)

	s.mu.RLock()
	defer s.mu.RUnlock()
	entries := make([]semanticHashEntry, 0, len(s.nodes))
	for _, n := range s.nodes {
		if n.Kind != "knowledge" && n.Kind != "ai-think" {
			continue
		}
		if !filter.Matches(n) {
			continue
		}
		if n.Kind == "ai-think" && maxAIDepth > 0 && graphNodeGenerationDepth(n) >= maxAIDepth {
			continue
		}
		v, ok := s.vectors[n.ID]
		if !ok || len(v) != len(q) {
			continue
		}
		entries = append(entries, semanticHashEntry{node: n, vector: v})
	}
	sort.Slice(entries, func(i, j int) bool { return entries[i].node.ID < entries[j].node.ID })
	stats.IndexedNodes = len(entries)
	h := &coarseMaxHeap{}
	heap.Init(h)
	for i := range entries {
		entries[i].signatures = signaturesFor(entries[i].vector, hashBits, hashTables)
		stats.CoarseComparisons++
		pushBestCoarse(h, coarseCandidate{index: i, distance: signatureDistance(qsig, entries[i].signatures)}, candidateLimit)
	}
	candidates := make([]coarseCandidate, h.Len())
	for i := len(candidates) - 1; i >= 0; i-- {
		candidates[i] = heap.Pop(h).(coarseCandidate)
	}
	stats.CandidatePool = len(candidates)
	hits := make([]model.Hit, 0, len(candidates))
	for _, candidate := range candidates {
		entry := entries[candidate.index]
		stats.ExactComparisons++
		hits = append(hits, model.Hit{NodeID: entry.node.ID, Label: entry.node.Label, Score: cosine32(q, entry.vector), Kind: entry.node.Kind, Status: entry.node.Status})
	}
	sort.Slice(hits, func(i, j int) bool {
		if hits[i].Score == hits[j].Score {
			return hits[i].NodeID < hits[j].NodeID
		}
		return hits[i].Score > hits[j].Score
	})
	if len(hits) > limit {
		hits = hits[:limit]
	}
	return hits, stats
}

// NeighborScores returns only non-taxonomy graph links incident to the seeds,
// without allocating a full graph Snapshot.
func (s *Store) NeighborScores(seedIDs map[string]bool) map[string]float64 {
	s.mu.RLock()
	defer s.mu.RUnlock()
	out := map[string]float64{}
	for _, edge := range s.edges {
		if edge.Status == "rejected" || clusterTaxonomyEdge(edge.Type) {
			continue
		}
		weight := edge.Confidence
		if edge.Weight > weight {
			weight = edge.Weight
		}
		if weight < .2 {
			weight = .2
		}
		if seedIDs[edge.Source] {
			out[edge.Target] += weight
		}
		if seedIDs[edge.Target] {
			out[edge.Source] += weight
		}
	}
	return out
}

func clusterTaxonomyEdge(edgeType string) bool {
	switch edgeType {
	case "categorized_as", "mentions", "derived_from":
		return true
	default:
		return false
	}
}