package relay

import (
	"net"
	"sync"
	"time"
)

// Binding is the identity an allocation was created under. Src is the
// typed source key (typedSrc) pion will present on later requests.
type Binding struct {
	ID          Identity
	Expiry      time.Time
	Kid         string
	Incarnation string
	Src         string
}

// typedSrc keys the owner binding on transport AND address, because pion
// tags every allocation's tuple as UDP even for TCP clients and a UDP and a
// TCP client can share an address string.
func typedSrc(a net.Addr) string {
	switch a.(type) {
	case *net.TCPAddr:
		return "tcp:" + a.String()
	default:
		return "udp:" + a.String()
	}
}

type portSet struct {
	mu sync.Mutex
	m  map[int]bool
}

func (p *portSet) add(port int)    { p.mu.Lock(); p.m[port] = true; p.mu.Unlock() }
func (p *portSet) remove(port int) { p.mu.Lock(); delete(p.m, port); p.mu.Unlock() }
func (p *portSet) has(port int) bool {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.m[port]
}

// relayConn is the PacketConn the relay hands pion for one allocation.
// Closing it is how the relay ends the allocation: pion's read loop deletes
// the allocation when ReadFrom errors (allocation.go:301 in pion/turn
// v4.1.4). WriteTo is the client→peer leg and is where forwarded bytes are
// counted, the per-allocation and box buckets applied, and the EIP port
// gate enforced. ReadFrom is not overridden: peer→relay traffic is counted
// by the server only after pion forwards it (see Task 6), so this wrapper
// counts nothing on read.
type relayConn struct {
	net.PacketConn
	b          *Binding
	eip        net.IP
	ports      *portSet
	appBucket  *Bucket
	boxBucket  *Bucket
	onIn       func(Identity, int)
	onRejected func(Identity, int)
	onClose    func(*relayConn)
	once       sync.Once
}

func newRelayConn(inner net.PacketConn, b *Binding, eip net.IP, ports *portSet, appBucket, boxBucket *Bucket,
	onIn func(Identity, int), onRejected func(Identity, int), onClose func(*relayConn)) *relayConn {
	return &relayConn{PacketConn: inner, b: b, eip: eip, ports: ports, appBucket: appBucket, boxBucket: boxBucket,
		onIn: onIn, onRejected: onRejected, onClose: onClose}
}

func (c *relayConn) WriteTo(p []byte, addr net.Addr) (int, error) {
	if ua, ok := addr.(*net.UDPAddr); ok && c.eip != nil && ua.IP.Equal(c.eip) && !c.ports.has(ua.Port) {
		if c.onRejected != nil {
			c.onRejected(c.b.ID, len(p))
		}
		return len(p), nil // dropped: only live relay ports are reachable on our own EIP
	}
	// Check both buckets before spending either: if we spent on appBucket
	// and boxBucket then refused, the packet is dropped but appBucket is
	// left debited, throttling the app harder than its rate whenever the
	// box is under pressure.
	if c.appBucket != nil && !c.appBucket.Peek(len(p)) {
		return len(p), nil // dropped by the app bucket
	}
	if c.boxBucket != nil && !c.boxBucket.Peek(len(p)) {
		return len(p), nil // dropped by the box bucket
	}
	if c.appBucket != nil {
		c.appBucket.Spend(len(p))
	}
	if c.boxBucket != nil {
		c.boxBucket.Spend(len(p))
	}
	n, err := c.PacketConn.WriteTo(p, addr)
	if n > 0 && c.onIn != nil {
		c.onIn(c.b.ID, n)
	}
	return n, err
}

func (c *relayConn) Close() error {
	var err error
	c.once.Do(func() {
		err = c.PacketConn.Close()
		if c.onClose != nil {
			c.onClose(c)
		}
	})
	return err
}