package socks

import (
	"math"
	"math/rand"
	"net"
	"sync"

	"github.com/v2ray/v2ray-core/common/collect"
	"github.com/v2ray/v2ray-core/common/log"
	v2net "github.com/v2ray/v2ray-core/common/net"
	"github.com/v2ray/v2ray-core/proxy/socks/protocol"
)

const (
	bufferSize = 2 * 1024
)

type portMap struct {
	access       sync.Mutex
	data         map[uint16]*net.UDPAddr
	removedPorts *collect.TimedQueue
}

func newPortMap() *portMap {
	m := &portMap{
		access:       sync.Mutex{},
		data:         make(map[uint16]*net.UDPAddr),
		removedPorts: collect.NewTimedQueue(1),
	}
	go m.removePorts(m.removedPorts.RemovedEntries())
	return m
}

func (m *portMap) assignAddressToken(addr *net.UDPAddr) uint16 {
	for {
		token := uint16(rand.Intn(math.MaxUint16))
		if _, used := m.data[token]; !used {
			m.access.Lock()
			if _, used = m.data[token]; !used {
				m.data[token] = addr
				m.access.Unlock()
				return token
			}
			m.access.Unlock()
		}
	}
}

func (m *portMap) removePorts(removedPorts <-chan interface{}) {
	for {
		rawToken := <-removedPorts
		m.access.Lock()
		delete(m.data, rawToken.(uint16))
		m.access.Unlock()
	}
}

func (m *portMap) popPort(token uint16) *net.UDPAddr {
	m.access.Lock()
	defer m.access.Unlock()
	addr, exists := m.data[token]
	if !exists {
		return nil
	}
	delete(m.data, token)
	return addr
}

var (
	ports *portMap

	udpConn *net.UDPConn
)

func (server *SocksServer) ListenUDP(port uint16) error {
	ports = newPortMap()

	addr := &net.UDPAddr{
		IP:   net.IP{0, 0, 0, 0},
		Port: int(port),
		Zone: "",
	}
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		log.Error("Socks failed to listen UDP on port %d: %v", port, err)
		return err
	}

	go server.AcceptPackets(conn)
	udpConn = conn
	return nil
}

func (server *SocksServer) AcceptPackets(conn *net.UDPConn) error {
	for {
		buffer := make([]byte, 0, bufferSize)
		nBytes, addr, err := conn.ReadFromUDP(buffer)
		if err != nil {
			log.Error("Socks failed to read UDP packets: %v", err)
			return err
		}
		request, err := protocol.ReadUDPRequest(buffer[:nBytes])
		if err != nil {
			log.Error("Socks failed to parse UDP request: %v", err)
			return err
		}
		if request.Fragment != 0 {
			// TODO handle fragments
			continue
		}

		token := ports.assignAddressToken(addr)

		udpPacket := v2net.NewUDPPacket(request.Destination(), request.Data, token)
		server.vPoint.DispatchToOutbound(udpPacket)
	}
}

func (server *SocksServer) Dispatch(packet v2net.Packet) {
	if udpPacket, ok := packet.(*v2net.UDPPacket); ok {
		token := udpPacket.Token()
		addr := ports.popPort(token)
		if udpConn != nil {
			udpConn.WriteToUDP(udpPacket.Chunk(), addr)
		}
	}
	// We don't expect TCP Packets here
}