package dht

import (
	"crypto/rand"
	"math/big"
	"net"
	"sort"
	"sync"
)

// NodeID is a 160-bit Kademlia ID.
type NodeID [20]byte

// RandomNodeID generates a new random NodeID.
func RandomNodeID() NodeID {
	var id NodeID
	_, _ = rand.Read(id[:])
	return id
}

// Node represents a contact in the DHT network.
type Node struct {
	ID   NodeID
	Addr *net.UDPAddr
}

// Distance calculates the XOR distance between two NodeIDs.
func Distance(a, b NodeID) *big.Int {
	var xor [20]byte
	for i := 0; i < 20; i++ {
		xor[i] = a[i] ^ b[i]
	}
	return new(big.Int).SetBytes(xor[:])
}

// Bucket represents a Kademlia k-bucket.
type Bucket struct {
	nodes []Node
}

// RoutingTable manages known nodes using Kademlia k-buckets.
type RoutingTable struct {
	mu      sync.RWMutex
	ownID   NodeID
	buckets [160]*Bucket
}

// NewRoutingTable creates a new routing table.
func NewRoutingTable(ownID NodeID) *RoutingTable {
	rt := &RoutingTable{
		ownID: ownID,
	}
	for i := 0; i < 160; i++ {
		rt.buckets[i] = &Bucket{nodes: make([]Node, 0, MaxNodes)}
	}
	return rt
}

// bucketIndex calculates the appropriate bucket index for a given node ID.
// Returns an index from 0 to 159, or -1 if the ID is our own.
func (rt *RoutingTable) bucketIndex(target NodeID) int {
	for i := 0; i < 20; i++ {
		xor := rt.ownID[i] ^ target[i]
		if xor != 0 {
			// Find the most significant bit set in the byte
			for j := 7; j >= 0; j-- {
				if (xor & (1 << j)) != 0 {
					return 159 - (i*8 + (7 - j))
				}
			}
		}
	}
	return -1
}

// AddNode adds a node to the routing table or updates it.
func (rt *RoutingTable) AddNode(n Node) {
	idx := rt.bucketIndex(n.ID)
	if idx == -1 {
		return // Do not add ourselves
	}

	rt.mu.Lock()
	defer rt.mu.Unlock()

	bucket := rt.buckets[idx]

	// Check if already exists and update
	for i, existing := range bucket.nodes {
		if existing.ID == n.ID {
			bucket.nodes[i].Addr = n.Addr
			// Move to end (most recently seen)
			node := bucket.nodes[i]
			bucket.nodes = append(bucket.nodes[:i], bucket.nodes[i+1:]...)
			bucket.nodes = append(bucket.nodes, node)
			return
		}
	}

	// Add new node if bucket is not full
	if len(bucket.nodes) < MaxNodes {
		bucket.nodes = append(bucket.nodes, n)
	} else {
		// In a full Kademlia implementation, we would ping the oldest node
		// and replace it if it doesn't respond. For now, just drop the new one.
	}
}

// ClosestNodes returns the closest nodes to a given target ID.
func (rt *RoutingTable) ClosestNodes(target NodeID, count int) []Node {
	rt.mu.RLock()
	defer rt.mu.RUnlock()

	var allNodes []Node
	
	// Fast path: find the target's bucket
	idx := rt.bucketIndex(target)
	if idx != -1 {
		allNodes = append(allNodes, rt.buckets[idx].nodes...)
	}

	// Collect nodes from neighboring buckets if we don't have enough
	if len(allNodes) < count {
		// We just collect all nodes and sort them for simplicity. 
		// Optimizing this is possible but not strictly necessary for < 1280 nodes total.
		allNodes = nil
		for i := 0; i < 160; i++ {
			allNodes = append(allNodes, rt.buckets[i].nodes...)
		}
	}

	sort.Slice(allNodes, func(i, j int) bool {
		distI := Distance(allNodes[i].ID, target)
		distJ := Distance(allNodes[j].ID, target)
		return distI.Cmp(distJ) < 0
	})

	if len(allNodes) > count {
		return allNodes[:count]
	}
	return allNodes
}

// Len returns the total number of nodes in the routing table.
func (rt *RoutingTable) Len() int {
	rt.mu.RLock()
	defer rt.mu.RUnlock()
	count := 0
	for i := 0; i < 160; i++ {
		count += len(rt.buckets[i].nodes)
	}
	return count
}