package services

import (
	"fmt"
	"math/rand/v2"
	"sync"

	"git.saqut.com/saqut/mwse/internal/protocol"
	"git.saqut.com/saqut/mwse/internal/ws"
)

// shortCodeAlphabet is the 22-letter set the original used (J, Q, U, W absent).
// Three letters give 22³ = 10,648 unique codes.
const shortCodeAlphabet = "ABCDEFGHIKLMNOPRSTVXYZ"

// randomProbes is how many random candidates we try before falling back to a
// sequential scan. At low occupancy the first probe almost always succeeds.
const randomProbes = 256

// Announcer receives address allocation events for external monitoring.
type Announcer interface {
	Announce(kind, action, clientID string, value any)
}

type noopAnnouncer struct{}

func (noopAnnouncer) Announce(string, string, string, any) {}

// SubNet is a /24 virtual sub-network (10.A.B.0/24).
// Clients in a group or room can draw IPs from within a dedicated prefix
// so they appear to share the same network segment.
//
// Thread-safety: protected by its own mutex; IPPressure's outer mutex guards
// the subnet registry but not the per-subnet host tables.
type SubNet struct {
	Prefix string // "10.A.B" — the /24 prefix
	Owner  string // room / group id that claimed this subnet

	mu    sync.Mutex
	hosts map[byte]string // host byte [1..254] → clientID
}

// Alloc assigns a random host address within this /24 to clientID.
func (sn *SubNet) Alloc(clientID string) (string, bool) {
	sn.mu.Lock()
	defer sn.mu.Unlock()

	// Random probe first.
	for range randomProbes {
		h := byte(rand.IntN(254)) + 1 // [1..254]
		if _, busy := sn.hosts[h]; !busy {
			sn.hosts[h] = clientID
			return fmt.Sprintf("%s.%d", sn.Prefix, h), true
		}
	}
	// Sequential fallback (subnet is more than 98% full).
	for h := byte(1); h < 255; h++ {
		if _, busy := sn.hosts[h]; !busy {
			sn.hosts[h] = clientID
			return fmt.Sprintf("%s.%d", sn.Prefix, h), true
		}
	}
	return "", false // /24 exhausted (~254 hosts)
}

// Release frees the host address previously allocated to clientID.
func (sn *SubNet) Release(clientID string) {
	sn.mu.Lock()
	defer sn.mu.Unlock()
	for h, id := range sn.hosts {
		if id == clientID {
			delete(sn.hosts, h)
			return
		}
	}
}

// Whois returns the clientID that holds the given IP within this subnet.
func (sn *SubNet) Whois(ip string) (string, bool) {
	sn.mu.Lock()
	defer sn.mu.Unlock()
	var host byte
	_, err := fmt.Sscanf(ip, sn.Prefix+".%d", &host)
	if err != nil {
		return "", false
	}
	id, ok := sn.hosts[host]
	return id, ok
}

// IPPressure allocates virtual addresses to clients.
//
// Allocation strategy (#41): random probe with sequential fallback.
// At typical occupancy (< 1 % of the address space) the first random
// candidate is free, so allocation is O(1) in practice.
//
// Sub-networks (#40): callers can reserve a /24 prefix and hand out IPs
// within it. The global flat space and the subnet space are disjoint; a
// client that holds a subnet IP is also tracked in busyIP so whois/APIPAddress
// queries still work.
type IPPressure struct {
	ann Announcer

	mu         sync.Mutex
	busyNumber map[int]string    // number → clientID
	busyCode   map[string]string // shortcode → clientID
	busyIP     map[string]string // ip → clientID (both flat and subnet)

	// Sub-network registry.
	subnets  map[string]*SubNet // prefix ("10.A.B") → SubNet
	clientSN map[string]string  // clientID → subnet prefix (for disconnect cleanup)
}

// NewIPPressure builds an allocator. A nil announcer becomes a no-op.
func NewIPPressure(ann Announcer) *IPPressure {
	if ann == nil {
		ann = noopAnnouncer{}
	}
	return &IPPressure{
		ann:        ann,
		busyNumber: make(map[int]string),
		busyCode:   make(map[string]string),
		busyIP:     make(map[string]string),
		subnets:    make(map[string]*SubNet),
		clientSN:   make(map[string]string),
	}
}

// ---- Flat IP address (10.x.x.x, random) --------------------------------

func (p *IPPressure) lockIP(clientID string) string {
	p.mu.Lock()
	defer p.mu.Unlock()

	// Random probe across the full 10.0.0.0/8.
	for range randomProbes {
		a := byte(rand.IntN(255)) + 1
		b := byte(rand.IntN(256))
		c := byte(rand.IntN(255)) + 1
		ip := fmt.Sprintf("10.%d.%d.%d", a, b, c)
		if _, busy := p.busyIP[ip]; !busy {
			p.busyIP[ip] = clientID
			p.ann.Announce("AP_IPADDRESS", "LOCK", clientID, ip)
			return ip
		}
	}
	// Sequential fallback (extremely unlikely with random probes).
	for a := 1; a <= 255; a++ {
		for b := 0; b <= 255; b++ {
			for c := 1; c <= 255; c++ {
				ip := fmt.Sprintf("10.%d.%d.%d", a, b, c)
				if _, busy := p.busyIP[ip]; !busy {
					p.busyIP[ip] = clientID
					p.ann.Announce("AP_IPADDRESS", "LOCK", clientID, ip)
					return ip
				}
			}
		}
	}
	return "" // address space exhausted (> ~16 million clients)
}

func (p *IPPressure) releaseIP(ip string) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if clientID, ok := p.busyIP[ip]; ok {
		p.ann.Announce("AP_IPADDRESS", "RELEASE", clientID, ip)
		delete(p.busyIP, ip)
	}
}

func (p *IPPressure) whoisIP(ip string) (string, bool) {
	p.mu.Lock()
	defer p.mu.Unlock()
	id, ok := p.busyIP[ip]
	return id, ok
}

// ---- Number (random in [24, 99999]) -------------------------------------

func (p *IPPressure) lockNumber(clientID string) int {
	p.mu.Lock()
	defer p.mu.Unlock()

	const lo, hi = 24, 99999
	for range randomProbes {
		n := rand.IntN(hi-lo+1) + lo
		if _, busy := p.busyNumber[n]; !busy {
			p.busyNumber[n] = clientID
			p.ann.Announce("AP_NUMBER", "LOCK", clientID, n)
			return n
		}
	}
	// Sequential fallback.
	for n := lo; n <= hi; n++ {
		if _, busy := p.busyNumber[n]; !busy {
			p.busyNumber[n] = clientID
			p.ann.Announce("AP_NUMBER", "LOCK", clientID, n)
			return n
		}
	}
	// Beyond hi: keep counting up without bound.
	for n := hi + 1; ; n++ {
		if _, busy := p.busyNumber[n]; !busy {
			p.busyNumber[n] = clientID
			p.ann.Announce("AP_NUMBER", "LOCK", clientID, n)
			return n
		}
	}
}

func (p *IPPressure) releaseNumber(n int) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if clientID, ok := p.busyNumber[n]; ok {
		p.ann.Announce("AP_NUMBER", "RELEASE", clientID, n)
		delete(p.busyNumber, n)
	}
}

func (p *IPPressure) whoisNumber(n int) (string, bool) {
	p.mu.Lock()
	defer p.mu.Unlock()
	id, ok := p.busyNumber[n]
	return id, ok
}

// ---- Short code (random 3-letter from restricted alphabet) --------------

func (p *IPPressure) lockCode(clientID string) string {
	p.mu.Lock()
	defer p.mu.Unlock()

	n := len(shortCodeAlphabet)
	for range randomProbes {
		code := string([]byte{
			shortCodeAlphabet[rand.IntN(n)],
			shortCodeAlphabet[rand.IntN(n)],
			shortCodeAlphabet[rand.IntN(n)],
		})
		if _, busy := p.busyCode[code]; !busy {
			p.busyCode[code] = clientID
			p.ann.Announce("AP_SHORTCODE", "LOCK", clientID, code)
			return code
		}
	}
	// Sequential fallback.
	for _, a := range shortCodeAlphabet {
		for _, b := range shortCodeAlphabet {
			for _, d := range shortCodeAlphabet {
				code := string([]rune{a, b, d})
				if _, busy := p.busyCode[code]; !busy {
					p.busyCode[code] = clientID
					p.ann.Announce("AP_SHORTCODE", "LOCK", clientID, code)
					return code
				}
			}
		}
	}
	return "" // all 10,648 codes in use
}

func (p *IPPressure) releaseCode(code string) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if clientID, ok := p.busyCode[code]; ok {
		p.ann.Announce("AP_SHORTCODE", "RELEASE", clientID, code)
		delete(p.busyCode, code)
	}
}

func (p *IPPressure) whoisCode(code string) (string, bool) {
	p.mu.Lock()
	defer p.mu.Unlock()
	id, ok := p.busyCode[code]
	return id, ok
}

// ---- Sub-network (#40) -------------------------------------------------

// allocSubNet reserves a random /24 prefix (10.A.B.0/24) for the caller.
// The prefix is held until releaseSubNet is called or the client disconnects.
func (p *IPPressure) allocSubNet(ownerID string) (*SubNet, bool) {
	p.mu.Lock()
	defer p.mu.Unlock()

	for range randomProbes {
		a := byte(rand.IntN(255)) + 1
		b := byte(rand.IntN(256))
		prefix := fmt.Sprintf("10.%d.%d", a, b)
		if _, busy := p.subnets[prefix]; !busy {
			sn := &SubNet{Prefix: prefix, Owner: ownerID, hosts: make(map[byte]string)}
			p.subnets[prefix] = sn
			p.ann.Announce("AP_SUBNET", "ALLOC", ownerID, prefix)
			return sn, true
		}
	}
	// Sequential fallback.
	for a := 1; a <= 255; a++ {
		for b := 0; b <= 255; b++ {
			prefix := fmt.Sprintf("10.%d.%d", a, b)
			if _, busy := p.subnets[prefix]; !busy {
				sn := &SubNet{Prefix: prefix, Owner: ownerID, hosts: make(map[byte]string)}
				p.subnets[prefix] = sn
				p.ann.Announce("AP_SUBNET", "ALLOC", ownerID, prefix)
				return sn, true
			}
		}
	}
	return nil, false
}

// getSubNet looks up the SubNet owned by ownerID.
func (p *IPPressure) getSubNet(ownerID string) (*SubNet, bool) {
	p.mu.Lock()
	defer p.mu.Unlock()
	for _, sn := range p.subnets {
		if sn.Owner == ownerID {
			return sn, true
		}
	}
	return nil, false
}

// releaseSubNet frees a /24 prefix and removes all host addresses it contained
// from the global busyIP table.
func (p *IPPressure) releaseSubNet(prefix string) {
	p.mu.Lock()
	sn, ok := p.subnets[prefix]
	if !ok {
		p.mu.Unlock()
		return
	}
	delete(p.subnets, prefix)
	p.mu.Unlock()

	sn.mu.Lock()
	defer sn.mu.Unlock()
	for h, cid := range sn.hosts {
		ip := fmt.Sprintf("%s.%d", prefix, h)
		p.mu.Lock()
		delete(p.busyIP, ip)
		p.mu.Unlock()
		p.ann.Announce("AP_SUBNET_IP", "RELEASE", cid, ip)
	}
	p.ann.Announce("AP_SUBNET", "RELEASE", sn.Owner, prefix)
}

// ---- Service registration ----------------------------------------------

// registerIPPressure wires the alloc/realloc/release/whois handlers and the
// disconnect cleanup. The allocator instance is returned for tests.
func registerIPPressure(hub *ws.Hub, ann Announcer) *IPPressure {
	p := NewIPPressure(ann)

	// --- Flat IP address ---
	hub.Register("alloc/APIPAddress", func(c *ws.Client, m protocol.Message) any {
		if ip := c.APIP(); ip != "" {
			return map[string]any{"status": "success", "ip": ip}
		}
		ip := p.lockIP(c.ID)
		if ip == "" {
			return map[string]any{"status": "fail"}
		}
		c.SetAPIP(ip)
		return map[string]any{"status": "success", "ip": ip}
	})
	hub.Register("realloc/APIPAddress", func(c *ws.Client, m protocol.Message) any {
		old := c.APIP()
		if old == "" {
			return map[string]any{"status": "fail"}
		}
		ip := p.lockIP(c.ID)
		if ip == "" {
			return map[string]any{"status": "fail"}
		}
		p.releaseIP(old)
		c.SetAPIP(ip)
		return map[string]any{"status": "success", "ip": ip}
	})
	hub.Register("release/APIPAddress", func(c *ws.Client, m protocol.Message) any {
		p.releaseIP(c.APIP())
		c.SetAPIP("")
		return success()
	})
	hub.Register("whois/APIPAddress", func(c *ws.Client, m protocol.Message) any {
		if id, ok := p.whoisIP(m.Str("whois")); ok {
			return map[string]any{"status": "success", "socket": id}
		}
		return map[string]any{"status": "fail"}
	})

	// --- Number ---
	hub.Register("alloc/APNumber", func(c *ws.Client, m protocol.Message) any {
		if n := c.APNumber(); n != 0 {
			return map[string]any{"status": "success", "number": n}
		}
		n := p.lockNumber(c.ID)
		c.SetAPNumber(n)
		return map[string]any{"status": "success", "number": n}
	})
	hub.Register("realloc/APNumber", func(c *ws.Client, m protocol.Message) any {
		old := c.APNumber()
		if old == 0 {
			return map[string]any{"status": "fail"}
		}
		n := p.lockNumber(c.ID)
		p.releaseNumber(old)
		c.SetAPNumber(n)
		return map[string]any{"status": "success", "number": n}
	})
	hub.Register("release/APNumber", func(c *ws.Client, m protocol.Message) any {
		p.releaseNumber(c.APNumber())
		c.SetAPNumber(0)
		return success()
	})
	hub.Register("whois/APNumber", func(c *ws.Client, m protocol.Message) any {
		if id, ok := p.whoisNumber(m.Int("whois")); ok {
			return map[string]any{"status": "success", "socket": id}
		}
		return map[string]any{"status": "fail"}
	})

	// --- Short code ---
	hub.Register("alloc/APShortCode", func(c *ws.Client, m protocol.Message) any {
		if code := c.APShortCode(); code != "" {
			return map[string]any{"status": "success", "code": code}
		}
		code := p.lockCode(c.ID)
		if code == "" {
			return map[string]any{"status": "fail"}
		}
		c.SetAPShortCode(code)
		return map[string]any{"status": "success", "code": code}
	})
	hub.Register("realloc/APShortCode", func(c *ws.Client, m protocol.Message) any {
		old := c.APShortCode()
		if old == "" {
			return map[string]any{"status": "fail"}
		}
		code := p.lockCode(c.ID)
		if code == "" {
			return map[string]any{"status": "fail"}
		}
		p.releaseCode(old)
		c.SetAPShortCode(code)
		return map[string]any{"status": "success", "code": code}
	})
	hub.Register("release/APShortCode", func(c *ws.Client, m protocol.Message) any {
		p.releaseCode(c.APShortCode())
		c.SetAPShortCode("")
		return success()
	})
	hub.Register("whois/APShortCode", func(c *ws.Client, m protocol.Message) any {
		if id, ok := p.whoisCode(m.Str("whois")); ok {
			return map[string]any{"status": "success", "socket": id}
		}
		return map[string]any{"status": "fail"}
	})

	// --- Sub-network (#40) ---

	// alloc/APSubNet: claim a /24 prefix for a group/room. The caller becomes
	// the subnet owner; its members can then call alloc/APSubNetIP.
	hub.Register("alloc/APSubNet", func(c *ws.Client, m protocol.Message) any {
		// Idempotent: return existing prefix if the caller already owns one.
		if sn, ok := p.getSubNet(c.ID); ok {
			return map[string]any{"status": "success", "prefix": sn.Prefix}
		}
		sn, ok := p.allocSubNet(c.ID)
		if !ok {
			return map[string]any{"status": "fail"}
		}
		// Track on the client so disconnect cleanup works.
		p.mu.Lock()
		p.clientSN[c.ID] = sn.Prefix
		p.mu.Unlock()
		return map[string]any{"status": "success", "prefix": sn.Prefix}
	})

	// release/APSubNet: free the prefix. All host IPs are released.
	hub.Register("release/APSubNet", func(c *ws.Client, m protocol.Message) any {
		p.mu.Lock()
		prefix, had := p.clientSN[c.ID]
		if had {
			delete(p.clientSN, c.ID)
		}
		p.mu.Unlock()
		if had {
			p.releaseSubNet(prefix)
		}
		return success()
	})

	// alloc/APSubNetIP: allocate a host IP within the subnet identified by
	// the "prefix" field. The caller need not be the owner.
	hub.Register("alloc/APSubNetIP", func(c *ws.Client, m protocol.Message) any {
		prefix := m.Str("prefix")
		p.mu.Lock()
		sn, ok := p.subnets[prefix]
		p.mu.Unlock()
		if !ok {
			return map[string]any{"status": "fail", "message": "subnet not found"}
		}
		ip, ok := sn.Alloc(c.ID)
		if !ok {
			return map[string]any{"status": "fail", "message": "subnet exhausted"}
		}
		// Register in the global table so flat whois works.
		p.mu.Lock()
		p.busyIP[ip] = c.ID
		p.mu.Unlock()
		// If the client later disconnects, release via the subnet too.
		p.mu.Lock()
		p.clientSN[c.ID] = prefix
		p.mu.Unlock()
		c.SetAPIP(ip)
		return map[string]any{"status": "success", "ip": ip}
	})

	// release/APSubNetIP: free the current subnet IP of the caller.
	hub.Register("release/APSubNetIP", func(c *ws.Client, m protocol.Message) any {
		ip := c.APIP()
		if ip == "" {
			return success()
		}
		p.mu.Lock()
		prefix := p.clientSN[c.ID]
		sn, hasSN := p.subnets[prefix]
		delete(p.clientSN, c.ID)
		delete(p.busyIP, ip)
		p.mu.Unlock()
		if hasSN {
			sn.Release(c.ID)
		}
		c.SetAPIP("")
		return success()
	})

	// whois/APSubNetIP: look up a host IP within a given subnet.
	hub.Register("whois/APSubNetIP", func(c *ws.Client, m protocol.Message) any {
		prefix := m.Str("prefix")
		ip     := m.Str("whois")
		p.mu.Lock()
		sn, ok := p.subnets[prefix]
		p.mu.Unlock()
		if !ok {
			return map[string]any{"status": "fail"}
		}
		if id, ok := sn.Whois(ip); ok {
			return map[string]any{"status": "success", "socket": id}
		}
		return map[string]any{"status": "fail"}
	})

	// Release every address when a client disconnects.
	hub.OnDisconnect(func(c *ws.Client) {
		if ip := c.APIP(); ip != "" {
			p.releaseIP(ip)
			// Also release from subnet if applicable.
			p.mu.Lock()
			if prefix, ok := p.clientSN[c.ID]; ok {
				if sn, snOK := p.subnets[prefix]; snOK {
					sn.Release(c.ID)
				}
				delete(p.clientSN, c.ID)
			}
			p.mu.Unlock()
		}
		if c.APNumber() != 0 {
			p.releaseNumber(c.APNumber())
		}
		if code := c.APShortCode(); code != "" {
			p.releaseCode(code)
		}
		// Release owned subnet prefix on disconnect.
		p.mu.Lock()
		prefix, hadSN := p.clientSN[c.ID]
		if hadSN {
			delete(p.clientSN, c.ID)
		}
		p.mu.Unlock()
		if hadSN {
			p.releaseSubNet(prefix)
		}
	})

	return p
}