/*
Volt Mesh — WireGuard-based encrypted overlay network.

Provides peer-to-peer encrypted tunnels between Volt nodes using WireGuard
(kernel module). Each node gets a unique IP from the mesh CIDR, and peers
are discovered via the control plane or a shared cluster token.

Architecture:
  - WireGuard interface: voltmesh0 (configurable)
  - Mesh CIDR: 10.200.0.0/16 (default, supports ~65K nodes)
  - Each node: /32 address within the mesh CIDR
  - Key management: auto-generated WireGuard keypairs per node
  - Peer discovery: token-based join → control plane registration
  - Config persistence: /etc/volt/mesh/

Token format (base64-encoded JSON):
  {
    "mesh_cidr": "10.200.0.0/16",
    "control_endpoint": "198.58.96.144:51820",
    "control_pubkey": "...",
    "join_secret": "...",
    "mesh_id": "..."
  }

Copyright (c) Armored Gates LLC. All rights reserved.
AGPSL v5 — Source-available. Anti-competition clauses apply.
*/
package mesh

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

// ── Constants ────────────────────────────────────────────────────────────────

const (
	DefaultMeshCIDR    = "10.200.0.0/16"
	DefaultMeshPort    = 51820
	DefaultInterface   = "voltmesh0"
	MeshConfigDir      = "/etc/volt/mesh"
	MeshStateFile      = "/etc/volt/mesh/state.json"
	MeshPeersFile      = "/etc/volt/mesh/peers.json"
	WireGuardConfigDir = "/etc/wireguard"
	KeepAliveInterval  = 25 // seconds
)

// ── Token ────────────────────────────────────────────────────────────────────

// ClusterToken is the join token exchanged out-of-band to bootstrap mesh membership.
type ClusterToken struct {
	MeshCIDR         string `json:"mesh_cidr"`
	ControlEndpoint  string `json:"control_endpoint"`
	ControlPublicKey string `json:"control_pubkey"`
	JoinSecret       string `json:"join_secret"`
	MeshID           string `json:"mesh_id"`
}

// EncodeToken serializes and base64-encodes a cluster token.
func EncodeToken(t *ClusterToken) (string, error) {
	data, err := json.Marshal(t)
	if err != nil {
		return "", fmt.Errorf("failed to encode token: %w", err)
	}
	return base64.URLEncoding.EncodeToString(data), nil
}

// DecodeToken base64-decodes and deserializes a cluster token.
func DecodeToken(s string) (*ClusterToken, error) {
	data, err := base64.URLEncoding.DecodeString(s)
	if err != nil {
		return nil, fmt.Errorf("invalid token encoding: %w", err)
	}
	var t ClusterToken
	if err := json.Unmarshal(data, &t); err != nil {
		return nil, fmt.Errorf("invalid token format: %w", err)
	}
	if t.MeshCIDR == "" || t.MeshID == "" {
		return nil, fmt.Errorf("token missing required fields (mesh_cidr, mesh_id)")
	}
	return &t, nil
}

// ── Peer ─────────────────────────────────────────────────────────────────────

// Peer represents a node in the mesh network.
type Peer struct {
	NodeID     string    `json:"node_id"`
	PublicKey  string    `json:"public_key"`
	Endpoint   string    `json:"endpoint"`       // host:port (public IP + WireGuard port)
	MeshIP     string    `json:"mesh_ip"`         // 10.200.x.x/32
	AllowedIPs []string  `json:"allowed_ips"`     // CIDRs routed through this peer
	LastSeen   time.Time `json:"last_seen"`
	Latency    float64   `json:"latency_ms"`      // last measured RTT in ms
	Region     string    `json:"region,omitempty"` // optional region label
	Online     bool      `json:"online"`
}

// ── Mesh State ───────────────────────────────────────────────────────────────

// MeshState is the persistent on-disk state for this node's mesh membership.
type MeshState struct {
	NodeID      string `json:"node_id"`
	MeshID      string `json:"mesh_id"`
	MeshCIDR    string `json:"mesh_cidr"`
	MeshIP      string `json:"mesh_ip"`       // this node's mesh IP (e.g., 10.200.0.2)
	PrivateKey  string `json:"private_key"`
	PublicKey   string `json:"public_key"`
	ListenPort  int    `json:"listen_port"`
	Interface   string `json:"interface"`
	JoinedAt    time.Time `json:"joined_at"`
	IsControl   bool   `json:"is_control"`     // true if this node is the control plane
}

// ── Manager ──────────────────────────────────────────────────────────────────

// Manager handles mesh lifecycle operations.
type Manager struct {
	state *MeshState
	peers []*Peer
	mu    sync.RWMutex
}

// NewManager creates a mesh manager, loading state from disk if available.
func NewManager() *Manager {
	m := &Manager{}
	m.loadState()
	m.loadPeers()
	return m
}

// IsJoined returns true if this node is part of a mesh.
func (m *Manager) IsJoined() bool {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.state != nil && m.state.MeshID != ""
}

// State returns a copy of the current mesh state (nil if not joined).
func (m *Manager) State() *MeshState {
	m.mu.RLock()
	defer m.mu.RUnlock()
	if m.state == nil {
		return nil
	}
	copy := *m.state
	return &copy
}

// Peers returns a copy of the current peer list.
func (m *Manager) Peers() []*Peer {
	m.mu.RLock()
	defer m.mu.RUnlock()
	result := make([]*Peer, len(m.peers))
	for i, p := range m.peers {
		copy := *p
		result[i] = &copy
	}
	return result
}

// ── Init (Create a new mesh) ────────────────────────────────────────────────

// InitMesh creates a new mesh network and makes this node the control plane.
// Returns the cluster token for other nodes to join.
func (m *Manager) InitMesh(meshCIDR string, listenPort int, publicEndpoint string) (*ClusterToken, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	if m.state != nil && m.state.MeshID != "" {
		return nil, fmt.Errorf("already part of mesh %q — run 'volt mesh leave' first", m.state.MeshID)
	}

	if meshCIDR == "" {
		meshCIDR = DefaultMeshCIDR
	}
	if listenPort == 0 {
		listenPort = DefaultMeshPort
	}

	// Generate WireGuard keypair
	privKey, pubKey, err := generateWireGuardKeys()
	if err != nil {
		return nil, fmt.Errorf("failed to generate WireGuard keys: %w", err)
	}

	// Generate mesh ID
	meshID := generateMeshID()

	// Allocate first IP in mesh CIDR for control plane
	meshIP, err := allocateFirstIP(meshCIDR)
	if err != nil {
		return nil, fmt.Errorf("failed to allocate mesh IP: %w", err)
	}

	// Generate join secret
	joinSecret, err := generateSecret(32)
	if err != nil {
		return nil, fmt.Errorf("failed to generate join secret: %w", err)
	}

	// Generate node ID
	nodeID, err := generateNodeID()
	if err != nil {
		return nil, fmt.Errorf("failed to generate node ID: %w", err)
	}

	m.state = &MeshState{
		NodeID:     nodeID,
		MeshID:     meshID,
		MeshCIDR:   meshCIDR,
		MeshIP:     meshIP,
		PrivateKey: privKey,
		PublicKey:  pubKey,
		ListenPort: listenPort,
		Interface:  DefaultInterface,
		JoinedAt:   time.Now().UTC(),
		IsControl:  true,
	}

	// Configure WireGuard interface
	if err := m.configureInterface(); err != nil {
		m.state = nil
		return nil, fmt.Errorf("failed to configure WireGuard interface: %w", err)
	}

	// Save state
	if err := m.saveState(); err != nil {
		return nil, fmt.Errorf("failed to save mesh state: %w", err)
	}

	// Build cluster token
	token := &ClusterToken{
		MeshCIDR:         meshCIDR,
		ControlEndpoint:  publicEndpoint,
		ControlPublicKey: pubKey,
		JoinSecret:       joinSecret,
		MeshID:           meshID,
	}

	return token, nil
}

// ── Join ─────────────────────────────────────────────────────────────────────

// JoinMesh joins this node to an existing mesh using a cluster token.
func (m *Manager) JoinMesh(tokenStr string, listenPort int, publicEndpoint string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if m.state != nil && m.state.MeshID != "" {
		return fmt.Errorf("already part of mesh %q — run 'volt mesh leave' first", m.state.MeshID)
	}

	token, err := DecodeToken(tokenStr)
	if err != nil {
		return fmt.Errorf("invalid cluster token: %w", err)
	}

	if listenPort == 0 {
		listenPort = DefaultMeshPort
	}

	// Generate WireGuard keypair
	privKey, pubKey, err := generateWireGuardKeys()
	if err != nil {
		return fmt.Errorf("failed to generate WireGuard keys: %w", err)
	}

	// Generate node ID
	nodeID, err := generateNodeID()
	if err != nil {
		return fmt.Errorf("failed to generate node ID: %w", err)
	}

	// Allocate a mesh IP (in production, the control plane would assign this;
	// for now, derive from node ID hash to avoid collisions)
	meshIP, err := allocateIPFromNodeID(token.MeshCIDR, nodeID)
	if err != nil {
		return fmt.Errorf("failed to allocate mesh IP: %w", err)
	}

	m.state = &MeshState{
		NodeID:     nodeID,
		MeshID:     token.MeshID,
		MeshCIDR:   token.MeshCIDR,
		MeshIP:     meshIP,
		PrivateKey: privKey,
		PublicKey:  pubKey,
		ListenPort: listenPort,
		Interface:  DefaultInterface,
		JoinedAt:   time.Now().UTC(),
		IsControl:  false,
	}

	// Configure WireGuard interface
	if err := m.configureInterface(); err != nil {
		m.state = nil
		return fmt.Errorf("failed to configure WireGuard interface: %w", err)
	}

	// Add control plane as first peer
	controlPeer := &Peer{
		NodeID:     "control",
		PublicKey:  token.ControlPublicKey,
		Endpoint:   token.ControlEndpoint,
		MeshIP:     "", // resolved dynamically
		AllowedIPs: []string{token.MeshCIDR},
		LastSeen:   time.Now().UTC(),
		Online:     true,
	}
	m.peers = []*Peer{controlPeer}

	// Add control plane peer to WireGuard
	if err := m.addWireGuardPeer(controlPeer); err != nil {
		return fmt.Errorf("failed to add control plane peer: %w", err)
	}

	// Save state
	if err := m.saveState(); err != nil {
		return fmt.Errorf("failed to save mesh state: %w", err)
	}
	if err := m.savePeers(); err != nil {
		return fmt.Errorf("failed to save peer list: %w", err)
	}

	return nil
}

// ── Leave ────────────────────────────────────────────────────────────────────

// LeaveMesh removes this node from the mesh, tearing down the WireGuard interface.
func (m *Manager) LeaveMesh() error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if m.state == nil || m.state.MeshID == "" {
		return fmt.Errorf("not part of any mesh")
	}

	// Tear down WireGuard interface
	exec.Command("ip", "link", "set", m.state.Interface, "down").Run()
	exec.Command("ip", "link", "del", m.state.Interface).Run()

	// Clean up config files
	os.Remove(filepath.Join(WireGuardConfigDir, m.state.Interface+".conf"))

	// Clear state
	m.state = nil
	m.peers = nil

	// Remove state files
	os.Remove(MeshStateFile)
	os.Remove(MeshPeersFile)

	return nil
}

// ── Add/Remove Peers ─────────────────────────────────────────────────────────

// AddPeer registers a new peer in the mesh and configures the WireGuard tunnel.
func (m *Manager) AddPeer(peer *Peer) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if m.state == nil {
		return fmt.Errorf("not part of any mesh")
	}

	// Check for duplicate
	for _, existing := range m.peers {
		if existing.NodeID == peer.NodeID {
			// Update existing peer
			existing.Endpoint = peer.Endpoint
			existing.PublicKey = peer.PublicKey
			existing.AllowedIPs = peer.AllowedIPs
			existing.LastSeen = time.Now().UTC()
			existing.Online = true
			if err := m.addWireGuardPeer(existing); err != nil {
				return fmt.Errorf("failed to update WireGuard peer: %w", err)
			}
			return m.savePeers()
		}
	}

	peer.LastSeen = time.Now().UTC()
	peer.Online = true
	m.peers = append(m.peers, peer)

	if err := m.addWireGuardPeer(peer); err != nil {
		return fmt.Errorf("failed to add WireGuard peer: %w", err)
	}

	return m.savePeers()
}

// RemovePeer removes a peer from the mesh.
func (m *Manager) RemovePeer(nodeID string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if m.state == nil {
		return fmt.Errorf("not part of any mesh")
	}

	var remaining []*Peer
	var removed *Peer
	for _, p := range m.peers {
		if p.NodeID == nodeID {
			removed = p
		} else {
			remaining = append(remaining, p)
		}
	}

	if removed == nil {
		return fmt.Errorf("peer %q not found", nodeID)
	}

	m.peers = remaining

	// Remove from WireGuard
	exec.Command("wg", "set", m.state.Interface,
		"peer", removed.PublicKey, "remove").Run()

	return m.savePeers()
}

// ── Latency Measurement ──────────────────────────────────────────────────────

// MeasureLatency pings all peers and updates their latency values.
func (m *Manager) MeasureLatency() {
	m.mu.Lock()
	defer m.mu.Unlock()

	for _, peer := range m.peers {
		if peer.MeshIP == "" {
			continue
		}
		// Parse mesh IP (strip /32 if present)
		ip := strings.Split(peer.MeshIP, "/")[0]
		start := time.Now()
		cmd := exec.Command("ping", "-c", "1", "-W", "2", ip)
		if err := cmd.Run(); err != nil {
			peer.Online = false
			peer.Latency = -1
			continue
		}
		peer.Latency = float64(time.Since(start).Microseconds()) / 1000.0
		peer.Online = true
		peer.LastSeen = time.Now().UTC()
	}
}

// ── WireGuard Configuration ──────────────────────────────────────────────────

// configureInterface creates and configures the WireGuard network interface.
func (m *Manager) configureInterface() error {
	iface := m.state.Interface
	meshIP := m.state.MeshIP
	listenPort := m.state.ListenPort

	// Create WireGuard interface
	if out, err := exec.Command("ip", "link", "add", iface, "type", "wireguard").CombinedOutput(); err != nil {
		return fmt.Errorf("failed to create WireGuard interface: %s", string(out))
	}

	// Write private key to temp file for wg
	keyFile := filepath.Join(MeshConfigDir, "private.key")
	os.MkdirAll(MeshConfigDir, 0700)
	if err := os.WriteFile(keyFile, []byte(m.state.PrivateKey), 0600); err != nil {
		return fmt.Errorf("failed to write private key: %w", err)
	}

	// Configure WireGuard
	if out, err := exec.Command("wg", "set", iface,
		"listen-port", fmt.Sprintf("%d", listenPort),
		"private-key", keyFile,
	).CombinedOutput(); err != nil {
		return fmt.Errorf("failed to configure WireGuard: %s", string(out))
	}

	// Assign mesh IP
	_, meshNet, _ := net.ParseCIDR(m.state.MeshCIDR)
	ones, _ := meshNet.Mask.Size()
	if out, err := exec.Command("ip", "addr", "add",
		fmt.Sprintf("%s/%d", meshIP, ones),
		"dev", iface,
	).CombinedOutput(); err != nil {
		return fmt.Errorf("failed to assign mesh IP: %s", string(out))
	}

	// Bring up interface
	if out, err := exec.Command("ip", "link", "set", iface, "up").CombinedOutput(); err != nil {
		return fmt.Errorf("failed to bring up interface: %s", string(out))
	}

	// Write WireGuard config file for wg-quick compatibility
	m.writeWireGuardConfig()

	return nil
}

// addWireGuardPeer adds or updates a peer in the WireGuard interface.
func (m *Manager) addWireGuardPeer(peer *Peer) error {
	args := []string{"set", m.state.Interface, "peer", peer.PublicKey}

	if peer.Endpoint != "" {
		args = append(args, "endpoint", peer.Endpoint)
	}

	allowedIPs := peer.AllowedIPs
	if len(allowedIPs) == 0 && peer.MeshIP != "" {
		ip := strings.Split(peer.MeshIP, "/")[0]
		allowedIPs = []string{ip + "/32"}
	}
	if len(allowedIPs) > 0 {
		args = append(args, "allowed-ips", strings.Join(allowedIPs, ","))
	}

	args = append(args, "persistent-keepalive", fmt.Sprintf("%d", KeepAliveInterval))

	if out, err := exec.Command("wg", args...).CombinedOutput(); err != nil {
		return fmt.Errorf("wg set peer failed: %s", string(out))
	}

	return nil
}

// writeWireGuardConfig generates a wg-quick compatible config file.
func (m *Manager) writeWireGuardConfig() error {
	os.MkdirAll(WireGuardConfigDir, 0700)

	_, meshNet, _ := net.ParseCIDR(m.state.MeshCIDR)
	ones, _ := meshNet.Mask.Size()

	var sb strings.Builder
	sb.WriteString("[Interface]\n")
	sb.WriteString(fmt.Sprintf("PrivateKey = %s\n", m.state.PrivateKey))
	sb.WriteString(fmt.Sprintf("ListenPort = %d\n", m.state.ListenPort))
	sb.WriteString(fmt.Sprintf("Address = %s/%d\n", m.state.MeshIP, ones))
	sb.WriteString("\n")

	for _, peer := range m.peers {
		sb.WriteString("[Peer]\n")
		sb.WriteString(fmt.Sprintf("PublicKey = %s\n", peer.PublicKey))
		if peer.Endpoint != "" {
			sb.WriteString(fmt.Sprintf("Endpoint = %s\n", peer.Endpoint))
		}
		allowedIPs := peer.AllowedIPs
		if len(allowedIPs) == 0 && peer.MeshIP != "" {
			ip := strings.Split(peer.MeshIP, "/")[0]
			allowedIPs = []string{ip + "/32"}
		}
		if len(allowedIPs) > 0 {
			sb.WriteString(fmt.Sprintf("AllowedIPs = %s\n", strings.Join(allowedIPs, ", ")))
		}
		sb.WriteString(fmt.Sprintf("PersistentKeepalive = %d\n", KeepAliveInterval))
		sb.WriteString("\n")
	}

	confPath := filepath.Join(WireGuardConfigDir, m.state.Interface+".conf")
	return os.WriteFile(confPath, []byte(sb.String()), 0600)
}

// ── Persistence ──────────────────────────────────────────────────────────────

func (m *Manager) loadState() {
	data, err := os.ReadFile(MeshStateFile)
	if err != nil {
		return
	}
	var state MeshState
	if err := json.Unmarshal(data, &state); err != nil {
		return
	}
	m.state = &state
}

func (m *Manager) saveState() error {
	os.MkdirAll(MeshConfigDir, 0700)
	data, err := json.MarshalIndent(m.state, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(MeshStateFile, data, 0600)
}

func (m *Manager) loadPeers() {
	data, err := os.ReadFile(MeshPeersFile)
	if err != nil {
		return
	}
	var peers []*Peer
	if err := json.Unmarshal(data, &peers); err != nil {
		return
	}
	m.peers = peers
}

func (m *Manager) savePeers() error {
	os.MkdirAll(MeshConfigDir, 0700)
	data, err := json.MarshalIndent(m.peers, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(MeshPeersFile, data, 0600)
}

// ── Key Generation ───────────────────────────────────────────────────────────

// generateWireGuardKeys generates a WireGuard keypair using the `wg` tool.
func generateWireGuardKeys() (privateKey, publicKey string, err error) {
	// Generate private key
	privOut, err := exec.Command("wg", "genkey").Output()
	if err != nil {
		return "", "", fmt.Errorf("wg genkey failed: %w", err)
	}
	privateKey = strings.TrimSpace(string(privOut))

	// Derive public key
	cmd := exec.Command("wg", "pubkey")
	cmd.Stdin = strings.NewReader(privateKey)
	pubOut, err := cmd.Output()
	if err != nil {
		return "", "", fmt.Errorf("wg pubkey failed: %w", err)
	}
	publicKey = strings.TrimSpace(string(pubOut))

	return privateKey, publicKey, nil
}

// generateMeshID creates a random 8-character mesh identifier.
func generateMeshID() string {
	b := make([]byte, 4)
	rand.Read(b)
	return hex.EncodeToString(b)
}

// generateNodeID creates a random 16-character node identifier.
func generateNodeID() (string, error) {
	b := make([]byte, 8)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return hex.EncodeToString(b), nil
}

// generateSecret creates a random secret of the given byte length.
func generateSecret(length int) (string, error) {
	b := make([]byte, length)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return base64.URLEncoding.EncodeToString(b), nil
}

// ── IP Allocation ────────────────────────────────────────────────────────────

// allocateFirstIP returns the first usable IP in a CIDR (x.x.x.1).
func allocateFirstIP(cidr string) (string, error) {
	ip, _, err := net.ParseCIDR(cidr)
	if err != nil {
		return "", fmt.Errorf("invalid CIDR: %w", err)
	}
	ip4 := ip.To4()
	if ip4 == nil {
		return "", fmt.Errorf("only IPv4 is supported")
	}
	// First usable: network + 1
	ip4[3] = 1
	return ip4.String(), nil
}

// allocateIPFromNodeID deterministically derives a mesh IP from a node ID,
// using a hash to distribute IPs across the CIDR space.
func allocateIPFromNodeID(cidr, nodeID string) (string, error) {
	_, ipNet, err := net.ParseCIDR(cidr)
	if err != nil {
		return "", fmt.Errorf("invalid CIDR: %w", err)
	}

	ones, bits := ipNet.Mask.Size()
	hostBits := bits - ones
	maxHosts := (1 << hostBits) - 2 // exclude network and broadcast

	// Hash node ID to get a host number
	hash := sha256.Sum256([]byte(nodeID))
	hostNum := int(hash[0])<<8 | int(hash[1])
	hostNum = (hostNum % maxHosts) + 2 // +2 to skip .0 (network) and .1 (control)

	ip := make(net.IP, 4)
	copy(ip, ipNet.IP.To4())

	// Add host number to network address
	for i := 3; i >= 0 && hostNum > 0; i-- {
		ip[i] += byte(hostNum & 0xFF)
		hostNum >>= 8
	}

	return ip.String(), nil
}

// ── Status ───────────────────────────────────────────────────────────────────

// GetWireGuardStatus retrieves the current WireGuard interface status.
func (m *Manager) GetWireGuardStatus() (string, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	if m.state == nil {
		return "", fmt.Errorf("not part of any mesh")
	}

	out, err := exec.Command("wg", "show", m.state.Interface).CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("wg show failed: %s", string(out))
	}
	return string(out), nil
}