package xray

import (
	"context"
	"errors"
	"reflect"
	"sort"
	"sync"
	"sync/atomic"
	_ "unsafe"

	xrayDispatcher "github.com/xtls/xray-core/app/dispatcher"
	"github.com/xtls/xray-core/common"
	"github.com/xtls/xray-core/common/buf"
	"github.com/xtls/xray-core/common/net"
	"github.com/xtls/xray-core/common/session"
	"github.com/xtls/xray-core/features/routing"
	"github.com/xtls/xray-core/transport"
	"golang.org/x/time/rate"

	"github.com/cedar2025/xboard-node/internal/nlog"
)

// Access xray's internal config creator registry so we can replace the
// default dispatcher factory with ours. This runs AFTER xray's init()
// functions because our package imports xray (dependency order guarantee).
//
//go:linkname typeCreatorRegistry github.com/xtls/xray-core/common.typeCreatorRegistry
var typeCreatorRegistry map[reflect.Type]common.ConfigCreator

var origDispatcherFactory common.ConfigCreator

// globalLimitDispatcher is set when the factory creates a LimitDispatcher.
// The Xray kernel reads it to configure limits and get connections.
var globalLimitDispatcher atomic.Pointer[LimitDispatcher]

func init() {
	configType := reflect.TypeOf((*xrayDispatcher.Config)(nil))
	origDispatcherFactory = typeCreatorRegistry[configType]
	typeCreatorRegistry[configType] = limitDispatcherFactory
}

func limitDispatcherFactory(ctx context.Context, config interface{}) (interface{}, error) {
	orig, err := origDispatcherFactory(ctx, config)
	if err != nil {
		return nil, err
	}
	inner, ok := orig.(routing.Dispatcher)
	if !ok {
		return orig, nil
	}
	ld := &LimitDispatcher{
		inner:      orig,
		innerDisp:  inner,
		limitedIPs: make(map[string]map[string]int),
	}
	globalLimitDispatcher.Store(ld)
	nlog.Core().Debug("xray: limit dispatcher installed")
	return ld, nil
}

// LimitDispatcher wraps xray's DefaultDispatcher to add per-user device
// limit enforcement, speed limiting, and per-user traffic counting.
//
// Architecture: traffic bytes are accumulated in per-user atomic counters
// (not per-connection maps). IP tracking uses sync.Map for unlimited users
// (lock-free) and RWMutex-protected map for limited users (deterministic).
// This means GetUserTraffic() is O(users), not O(connections).
type LimitDispatcher struct {
	inner     interface{}        // original DefaultDispatcher (Feature + Dispatcher)
	innerDisp routing.Dispatcher // same object, typed as Dispatcher

	// limitedUsers: users with device limit > 0, protected by mu.
	// Needs deterministic IP ordering for kick decisions.
	mu           sync.RWMutex
	limitedIPs   map[string]map[string]int // email → sourceIP → refcount
	deviceLimits map[string]int            // email → max devices
	speedLimits  map[string]int            // email → Mbps
	emailToUID   map[string]int            // email → panel user ID

	// unlimitedIPs: users without device limit — sync.Map for lock-free access.
	// Each entry is *ipCounter{count atomic.Int64, ips sync.Map}.
	unlimitedIPs sync.Map // email → *ipCounter

	speedBuckets sync.Map // email → *rate.Limiter

	// Per-user traffic counters — lock-free atomics for the data path.
	userTraffic sync.Map // email → *userTrafficCounter

	connCount atomic.Int64 // total active connections (for metrics)
}

// ipCounter tracks IPs for unlimited users without any lock.
type ipCounter struct {
	count atomic.Int64 // total TCP connections (for connCount reporting)
	ips   sync.Map     // sourceIP → *atomic.Int64 (refcount)
}

// aliveIPs returns a snapshot of distinct IPs.
func (ic *ipCounter) aliveIPs() map[string]bool {
	result := make(map[string]bool)
	ic.ips.Range(func(key, _ interface{}) bool {
		result[key.(string)] = true
		return true
	})
	return result
}

// userTrafficCounter holds per-user cumulative traffic counters.
type userTrafficCounter struct {
	upload   atomic.Int64
	download atomic.Int64
}

// ─── routing.Dispatcher ──────────────────────────────────────────────────────

func (d *LimitDispatcher) Dispatch(ctx context.Context, dest net.Destination) (*transport.Link, error) {
	email, sourceIP, isTCP, err := d.identifyAndCheck(ctx, dest)
	if err != nil {
		return nil, err
	}

	link, err := d.innerDisp.Dispatch(ctx, dest)
	if err != nil {
		if email != "" && isTCP {
			d.delConn(email, sourceIP)
		}
		return nil, err
	}

	if email != "" {
		d.wrapLink(ctx, link, email, sourceIP, isTCP)
	}
	return link, nil
}

func (d *LimitDispatcher) DispatchLink(ctx context.Context, dest net.Destination, link *transport.Link) error {
	email, sourceIP, isTCP, err := d.identifyAndCheck(ctx, dest)
	if err != nil {
		return err
	}

	if email != "" {
		d.wrapLink(ctx, link, email, sourceIP, isTCP)
	}
	return d.innerDisp.DispatchLink(ctx, dest, link)
}

// identifyAndCheck extracts user identity from the session context, enforces
// device limits, and returns the user's email, source IP, and TCP flag.
// Returns a non-nil error only when the connection should be rejected.
func (d *LimitDispatcher) identifyAndCheck(ctx context.Context, dest net.Destination) (email, sourceIP string, isTCP bool, err error) {
	si := session.InboundFromContext(ctx)
	if si == nil || si.User == nil || len(si.User.Email) == 0 {
		return "", "", false, nil
	}
	email = si.User.Email
	sourceIP = si.Source.Address.IP().String()
	isTCP = dest.Network == net.Network_TCP

	if d.checkDeviceLimit(email, sourceIP, isTCP) {
		nlog.Core().Info("xray: device limit exceeded", "email", email, "ip", sourceIP)
		return "", "", false, errors.New("device limit exceeded for " + email)
	}
	return email, sourceIP, isTCP, nil
}

// wrapLink instruments a transport.Link with per-user byte counting,
// rate limiting, and lifecycle tracking. Must only be called when email != "".
func (d *LimitDispatcher) wrapLink(ctx context.Context, link *transport.Link, email, sourceIP string, isTCP bool) {
	// Get or create per-user traffic counter
	v, _ := d.userTraffic.LoadOrStore(email, &userTrafficCounter{})
	utc := v.(*userTrafficCounter)

	d.connCount.Add(1)

	onClose := func() {
		if isTCP {
			d.delConn(email, sourceIP)
		}
		d.connCount.Add(-1)
	}

	limiter := d.getBucket(email)

	link.Reader = &statsCloseReader{
		Reader:  link.Reader,
		counter: &utc.upload,
		limiter: limiter,
		ctx:     ctx,
	}
	link.Writer = &statsCloseWriter{
		Writer:  link.Writer,
		counter: &utc.download,
		limiter: limiter,
		ctx:     ctx,
		onClose: onClose,
	}
}

// ─── features.Feature (delegated) ───────────────────────────────────────────

func (d *LimitDispatcher) Type() interface{} { return routing.DispatcherType() }

func (d *LimitDispatcher) Start() error {
	if s, ok := d.inner.(interface{ Start() error }); ok {
		return s.Start()
	}
	return nil
}

func (d *LimitDispatcher) Close() error {
	if c, ok := d.inner.(interface{ Close() error }); ok {
		return c.Close()
	}
	return nil
}

// ─── Limit management (called by Xray kernel) ──────────────────────────────

func (d *LimitDispatcher) UpdateLimits(emailToUID map[string]int, deviceLimits, speedLimits map[string]int) {
	d.mu.Lock()
	d.emailToUID = emailToUID
	d.deviceLimits = deviceLimits
	d.speedLimits = speedLimits
	d.mu.Unlock()

	d.speedBuckets.Range(func(key, _ interface{}) bool {
		if _, ok := speedLimits[key.(string)]; !ok {
			d.speedBuckets.Delete(key)
		}
		return true
	})
}

func (d *LimitDispatcher) ResetConns() {
	d.mu.Lock()
	d.limitedIPs = make(map[string]map[string]int)
	d.mu.Unlock()

	// Clear unlimited IPs
	d.unlimitedIPs.Range(func(key, _ interface{}) bool {
		d.unlimitedIPs.Delete(key)
		return true
	})

	d.userTraffic.Range(func(key, _ interface{}) bool {
		d.userTraffic.Delete(key)
		return true
	})
	d.connCount.Store(0)

	d.speedBuckets.Range(func(key, _ interface{}) bool {
		d.speedBuckets.Delete(key)
		return true
	})
}

// GetUserTraffic returns per-user cumulative traffic and alive IPs.
// O(users), not O(connections).
func (d *LimitDispatcher) GetUserTraffic() (traffic map[int][2]int64, aliveIPs map[int]map[string]bool, connCount int) {
	d.mu.RLock()
	emailToUID := d.emailToUID
	limitedIPs := d.limitedIPs
	d.mu.RUnlock()

	aliveIPs = make(map[int]map[string]bool)

	// Collect IPs from limited users (under RLock snapshot).
	for email, ipsMap := range limitedIPs {
		uid := emailToUID[email]
		if uid == 0 {
			continue
		}
		ipSet := make(map[string]bool, len(ipsMap))
		for ip := range ipsMap {
			ipSet[ip] = true
		}
		if len(ipSet) > 0 {
			aliveIPs[uid] = ipSet
		}
	}

	// Collect IPs from unlimited users (lock-free).
	d.unlimitedIPs.Range(func(key, value interface{}) bool {
		email := key.(string)
		uid := emailToUID[email]
		if uid == 0 {
			return true
		}
		ic := value.(*ipCounter)
		if ips := ic.aliveIPs(); len(ips) > 0 {
			// Merge with limited IPs if any
			if existing, ok := aliveIPs[uid]; ok {
				for ip := range ips {
					existing[ip] = true
				}
			} else {
				aliveIPs[uid] = ips
			}
		}
		return true
	})

	traffic = make(map[int][2]int64)
	d.userTraffic.Range(func(key, value interface{}) bool {
		email := key.(string)
		utc := value.(*userTrafficCounter)
		uid := emailToUID[email]
		if uid == 0 {
			return true
		}
		up := utc.upload.Load()
		down := utc.download.Load()
		if up > 0 || down > 0 {
			traffic[uid] = [2]int64{up, down}
		}
		return true
	})

	connCount = int(d.connCount.Load())
	return
}

// ─── Internal helpers ───────────────────────────────────────────────────────

// checkDeviceLimit enforces per-user device limits.
// Fast path: unlimited users use lock-free sync.Map.
// Slow path: limited users use RWMutex with deterministic IP ordering.
func (d *LimitDispatcher) checkDeviceLimit(email, sourceIP string, isTCP bool) bool {
	d.mu.RLock()
	limit, hasLimit := d.deviceLimits[email]
	d.mu.RUnlock()

	// Fast path: no device limit — use lock-free sync.Map.
	if !hasLimit || limit <= 0 {
		if isTCP {
			v, _ := d.unlimitedIPs.LoadOrStore(email, &ipCounter{})
			ic := v.(*ipCounter)
			ic.count.Add(1)

			// Increment IP refcount atomically.
			rv, _ := ic.ips.LoadOrStore(sourceIP, &atomic.Int64{})
			rv.(*atomic.Int64).Add(1)
		}
		return false
	}

	// Slow path: user has device limit — need deterministic ordering.
	d.mu.RLock()
	ips := d.limitedIPs[email]
	if ips != nil && ips[sourceIP] > 0 {
		d.mu.RUnlock()
		if isTCP {
			d.mu.Lock()
			d.limitedIPs[email][sourceIP]++
			d.mu.Unlock()
		}
		return false
	}

	if ips != nil && len(ips) < limit {
		d.mu.RUnlock()
		if isTCP {
			d.mu.Lock()
			if d.limitedIPs[email] == nil {
				d.limitedIPs[email] = make(map[string]int)
			}
			d.limitedIPs[email][sourceIP]++
			d.mu.Unlock()
		}
		return false
	}
	d.mu.RUnlock()

	// Over limit — need write lock for deterministic check.
	d.mu.Lock()
	defer d.mu.Unlock()

	// Re-check under write lock.
	ips = d.limitedIPs[email]
	if ips == nil {
		ips = make(map[string]int)
		d.limitedIPs[email] = ips
	}

	if ips[sourceIP] > 0 {
		if isTCP {
			ips[sourceIP]++
		}
		return false
	}

	if len(ips) < limit {
		if isTCP {
			ips[sourceIP]++
		}
		return false
	}

	// Over limit — deterministic: allow lowest IPs lexicographically.
	ipList := make([]string, 0, len(ips)+1)
	for ip := range ips {
		ipList = append(ipList, ip)
	}
	ipList = append(ipList, sourceIP)
	sort.Strings(ipList)

	for i := 0; i < limit && i < len(ipList); i++ {
		if ipList[i] == sourceIP {
			if isTCP {
				ips[sourceIP]++
			}
			return false
		}
	}
	return true
}

// delConn decrements the IP refcount when a connection closes.
func (d *LimitDispatcher) delConn(email, sourceIP string) {
	// Check if this is an unlimited user first (lock-free).
	if v, ok := d.unlimitedIPs.Load(email); ok {
		ic := v.(*ipCounter)
		if rv, ok := ic.ips.Load(sourceIP); ok {
			counter := rv.(*atomic.Int64)
			if counter.Add(-1) <= 0 {
				ic.ips.Delete(sourceIP)
			}
		}
		return
	}

	// Limited user — use write lock.
	d.mu.Lock()
	defer d.mu.Unlock()
	if ips, ok := d.limitedIPs[email]; ok {
		ips[sourceIP]--
		if ips[sourceIP] <= 0 {
			delete(ips, sourceIP)
		}
		if len(ips) == 0 {
			delete(d.limitedIPs, email)
		}
	}
}

func (d *LimitDispatcher) getBucket(email string) *rate.Limiter {
	d.mu.RLock()
	mbps, ok := d.speedLimits[email]
	d.mu.RUnlock()
	if !ok || mbps <= 0 {
		return nil
	}
	bytesPerSec := int(mbps) * 1_000_000 / 8
	burst := bytesPerSec
	if burst < 64*1024 {
		burst = 64 * 1024
	}
	v, loaded := d.speedBuckets.LoadOrStore(email, rate.NewLimiter(rate.Limit(bytesPerSec), burst))
	lim := v.(*rate.Limiter)
	if loaded {
		// Update existing limiter in case speed limit changed.
		lim.SetLimit(rate.Limit(bytesPerSec))
		lim.SetBurst(burst)
	}
	return lim
}

// ─── I/O wrappers ───────────────────────────────────────────────────────────

const xrayRateLimitThreshold = 64 * 1024

type statsCloseReader struct {
	buf.Reader
	counter       *atomic.Int64
	limiter       *rate.Limiter
	ctx           context.Context
	pendingTokens atomic.Int64
}

func (r *statsCloseReader) ReadMultiBuffer() (buf.MultiBuffer, error) {
	mb, err := r.Reader.ReadMultiBuffer()
	if n := int64(mb.Len()); n > 0 {
		r.counter.Add(n)
		if r.limiter != nil {
			if pending := r.pendingTokens.Add(n); pending >= xrayRateLimitThreshold {
				tokens := int(r.pendingTokens.Swap(0))
				if burst := r.limiter.Burst(); tokens > burst {
					tokens = burst
				}
				_ = r.limiter.WaitN(r.ctx, tokens)
			}
		}
	}
	return mb, err
}

func (r *statsCloseReader) Close() error { return common.Close(r.Reader) }
func (r *statsCloseReader) Interrupt()   { common.Interrupt(r.Reader) }

type statsCloseWriter struct {
	buf.Writer
	counter       *atomic.Int64
	limiter       *rate.Limiter
	ctx           context.Context
	onClose       func()
	closed        atomic.Bool
	pendingTokens atomic.Int64
}

func (w *statsCloseWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
	n := int64(mb.Len())
	if w.limiter != nil {
		if pending := w.pendingTokens.Add(n); pending >= xrayRateLimitThreshold {
			tokens := int(w.pendingTokens.Swap(0))
			if burst := w.limiter.Burst(); tokens > burst {
				tokens = burst
			}
			_ = w.limiter.WaitN(w.ctx, tokens)
		}
	}
	w.counter.Add(n)
	return w.Writer.WriteMultiBuffer(mb)
}

func (w *statsCloseWriter) Close() error {
	if w.closed.CompareAndSwap(false, true) {
		w.onClose()
	}
	return common.Close(w.Writer)
}

func (w *statsCloseWriter) Interrupt() {
	if w.closed.CompareAndSwap(false, true) {
		w.onClose()
	}
	common.Interrupt(w.Writer)
}