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[:])
}

// RoutingTable manages known nodes.
type RoutingTable struct {
	mu    sync.RWMutex
	ownID NodeID
	nodes []Node
}

// NewRoutingTable creates a new routing table.
func NewRoutingTable(ownID NodeID) *RoutingTable {
	return &RoutingTable{
		ownID: ownID,
		nodes: make([]Node, 0),
	}
}

// AddNode adds a node to the routing table or updates it.
func (rt *RoutingTable) AddNode(n Node) {
	rt.mu.Lock()
	defer rt.mu.Unlock()

	if n.ID == rt.ownID {
		return
	}

	for i, existing := range rt.nodes {
		if existing.ID == n.ID {
			rt.nodes[i].Addr = n.Addr
			return
		}
	}

	rt.nodes = append(rt.nodes, n)

	// Sort by distance to our own ID and keep top 1000 nodes for simplicity.
	sort.Slice(rt.nodes, func(i, j int) bool {
		distI := Distance(rt.nodes[i].ID, rt.ownID)
		distJ := Distance(rt.nodes[j].ID, rt.ownID)
		return distI.Cmp(distJ) < 0
	})

	if len(rt.nodes) > 1000 {
		rt.nodes = rt.nodes[:1000]
	}
}

// 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()

	// Since we don't have true k-buckets, we sort the entire list.
	// This is O(n log n) but fine for a small simplified table of 1000 nodes.
	sortedNodes := make([]Node, len(rt.nodes))
	copy(sortedNodes, rt.nodes)

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

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

// Len returns the number of nodes in the routing table.
func (rt *RoutingTable) Len() int {
	rt.mu.RLock()
	defer rt.mu.RUnlock()
	return len(rt.nodes)
}