package unified

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"
)

// IPAM manages IP address allocation for networks
type IPAM struct {
	stateDir string
	pools    map[string]*Pool
	mu       sync.RWMutex
}

// Pool represents an IP address pool for a network
type Pool struct {
	// Network name
	Name string `json:"name"`

	// Subnet
	Subnet *net.IPNet `json:"subnet"`

	// Gateway address
	Gateway net.IP `json:"gateway"`

	// Pool start (first allocatable address)
	Start net.IP `json:"start"`

	// Pool end (last allocatable address)
	End net.IP `json:"end"`

	// Static reservations (workloadID -> IP)
	Reservations map[string]net.IP `json:"reservations"`

	// Active leases
	Leases map[string]*Lease `json:"leases"`

	// Free IPs (bitmap for fast allocation)
	allocated map[uint32]bool
}

// NewIPAM creates a new IPAM instance
func NewIPAM(stateDir string) (*IPAM, error) {
	if err := os.MkdirAll(stateDir, 0755); err != nil {
		return nil, fmt.Errorf("create IPAM state dir: %w", err)
	}

	ipam := &IPAM{
		stateDir: stateDir,
		pools:    make(map[string]*Pool),
	}

	// Load existing state
	if err := ipam.loadState(); err != nil {
		// Non-fatal, might be first run
		_ = err
	}

	return ipam, nil
}

// AddPool adds a new IP pool for a network
func (i *IPAM) AddPool(name string, subnet *net.IPNet, gateway net.IP, reservations map[string]net.IP) error {
	i.mu.Lock()
	defer i.mu.Unlock()

	// Calculate pool range
	start := nextIP(subnet.IP)
	if gateway != nil && gateway.Equal(start) {
		start = nextIP(start)
	}

	// Broadcast address is last in subnet
	end := lastIP(subnet)

	pool := &Pool{
		Name:         name,
		Subnet:       subnet,
		Gateway:      gateway,
		Start:        start,
		End:          end,
		Reservations: reservations,
		Leases:       make(map[string]*Lease),
		allocated:    make(map[uint32]bool),
	}

	// Mark gateway as allocated
	if gateway != nil {
		pool.allocated[ipToUint32(gateway)] = true
	}

	// Mark reservations as allocated
	for _, ip := range reservations {
		pool.allocated[ipToUint32(ip)] = true
	}

	i.pools[name] = pool
	return i.saveState()
}

// Allocate allocates an IP address for a workload
func (i *IPAM) Allocate(network, workloadID string, mac net.HardwareAddr) (*Lease, error) {
	i.mu.Lock()
	defer i.mu.Unlock()

	pool, ok := i.pools[network]
	if !ok {
		return nil, fmt.Errorf("network %s not found", network)
	}

	// Check if workload already has a lease
	if lease, ok := pool.Leases[workloadID]; ok {
		return lease, nil
	}

	// Check for static reservation
	if ip, ok := pool.Reservations[workloadID]; ok {
		lease := &Lease{
			IP:         ip,
			MAC:        mac,
			WorkloadID: workloadID,
			Start:      time.Now(),
			Expires:    time.Now().Add(365 * 24 * time.Hour), // Long lease for static
			Static:     true,
		}
		pool.Leases[workloadID] = lease
		pool.allocated[ipToUint32(ip)] = true
		_ = i.saveState()
		return lease, nil
	}

	// Find free IP in pool
	ip, err := pool.findFreeIP()
	if err != nil {
		return nil, err
	}

	lease := &Lease{
		IP:         ip,
		MAC:        mac,
		WorkloadID: workloadID,
		Start:      time.Now(),
		Expires:    time.Now().Add(24 * time.Hour), // Default 24h lease
		Static:     false,
	}

	pool.Leases[workloadID] = lease
	pool.allocated[ipToUint32(ip)] = true
	_ = i.saveState()

	return lease, nil
}

// Release releases an IP address allocation
func (i *IPAM) Release(network, workloadID string) error {
	i.mu.Lock()
	defer i.mu.Unlock()

	pool, ok := i.pools[network]
	if !ok {
		return nil // Network doesn't exist, nothing to release
	}

	lease, ok := pool.Leases[workloadID]
	if !ok {
		return nil // No lease, nothing to release
	}

	// Don't release static reservations from allocated map
	if !lease.Static {
		delete(pool.allocated, ipToUint32(lease.IP))
	}

	delete(pool.Leases, workloadID)
	return i.saveState()
}

// GetLease returns the current lease for a workload
func (i *IPAM) GetLease(network, workloadID string) (*Lease, error) {
	i.mu.RLock()
	defer i.mu.RUnlock()

	pool, ok := i.pools[network]
	if !ok {
		return nil, fmt.Errorf("network %s not found", network)
	}

	lease, ok := pool.Leases[workloadID]
	if !ok {
		return nil, fmt.Errorf("no lease for %s", workloadID)
	}

	return lease, nil
}

// ListLeases returns all active leases for a network
func (i *IPAM) ListLeases(network string) ([]*Lease, error) {
	i.mu.RLock()
	defer i.mu.RUnlock()

	pool, ok := i.pools[network]
	if !ok {
		return nil, fmt.Errorf("network %s not found", network)
	}

	result := make([]*Lease, 0, len(pool.Leases))
	for _, lease := range pool.Leases {
		result = append(result, lease)
	}

	return result, nil
}

// Reserve creates a static IP reservation
func (i *IPAM) Reserve(network, workloadID string, ip net.IP) error {
	i.mu.Lock()
	defer i.mu.Unlock()

	pool, ok := i.pools[network]
	if !ok {
		return fmt.Errorf("network %s not found", network)
	}

	// Check if IP is in subnet
	if !pool.Subnet.Contains(ip) {
		return fmt.Errorf("IP %s not in subnet %s", ip, pool.Subnet)
	}

	// Check if already allocated
	if pool.allocated[ipToUint32(ip)] {
		return fmt.Errorf("IP %s already allocated", ip)
	}

	if pool.Reservations == nil {
		pool.Reservations = make(map[string]net.IP)
	}
	pool.Reservations[workloadID] = ip
	pool.allocated[ipToUint32(ip)] = true

	return i.saveState()
}

// Unreserve removes a static IP reservation
func (i *IPAM) Unreserve(network, workloadID string) error {
	i.mu.Lock()
	defer i.mu.Unlock()

	pool, ok := i.pools[network]
	if !ok {
		return nil
	}

	if ip, ok := pool.Reservations[workloadID]; ok {
		delete(pool.allocated, ipToUint32(ip))
		delete(pool.Reservations, workloadID)
		return i.saveState()
	}

	return nil
}

// findFreeIP finds the next available IP in the pool
func (p *Pool) findFreeIP() (net.IP, error) {
	startUint := ipToUint32(p.Start)
	endUint := ipToUint32(p.End)

	for ip := startUint; ip <= endUint; ip++ {
		if !p.allocated[ip] {
			return uint32ToIP(ip), nil
		}
	}

	return nil, fmt.Errorf("no free IPs in pool %s", p.Name)
}

// saveState persists IPAM state to disk
func (i *IPAM) saveState() error {
	data, err := json.MarshalIndent(i.pools, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(filepath.Join(i.stateDir, "pools.json"), data, 0644)
}

// loadState loads IPAM state from disk
func (i *IPAM) loadState() error {
	data, err := os.ReadFile(filepath.Join(i.stateDir, "pools.json"))
	if err != nil {
		return err
	}

	if err := json.Unmarshal(data, &i.pools); err != nil {
		return err
	}

	// Rebuild allocated maps
	for _, pool := range i.pools {
		pool.allocated = make(map[uint32]bool)
		if pool.Gateway != nil {
			pool.allocated[ipToUint32(pool.Gateway)] = true
		}
		for _, ip := range pool.Reservations {
			pool.allocated[ipToUint32(ip)] = true
		}
		for _, lease := range pool.Leases {
			pool.allocated[ipToUint32(lease.IP)] = true
		}
	}

	return nil
}

// Helper functions for IP math

func ipToUint32(ip net.IP) uint32 {
	ip = ip.To4()
	if ip == nil {
		return 0
	}
	return binary.BigEndian.Uint32(ip)
}

func uint32ToIP(n uint32) net.IP {
	ip := make(net.IP, 4)
	binary.BigEndian.PutUint32(ip, n)
	return ip
}

func nextIP(ip net.IP) net.IP {
	return uint32ToIP(ipToUint32(ip) + 1)
}

func lastIP(subnet *net.IPNet) net.IP {
	// Get the broadcast address (last IP in subnet)
	ip := subnet.IP.To4()
	mask := subnet.Mask
	broadcast := make(net.IP, 4)
	for i := range ip {
		broadcast[i] = ip[i] | ^mask[i]
	}
	// Return one before broadcast (last usable)
	return uint32ToIP(ipToUint32(broadcast) - 1)
}