package socks5

import (
	"bytes"
	"errors"
	"net"
)

var (
	// ErrNoConnectedAddress is the no connected address error
	ErrNoConnectedAddress = errors.New("no connected address")
)

type UDPConn struct {
	net.PacketConn

	TCPConn       net.Conn
	ServerAddr    net.Addr
	ConnectedAddr net.Addr
}

func (s *UDPConn) ReadFrom(b []byte) (int, net.Addr, error) {
	_, _, err := s.PacketConn.ReadFrom(b)
	if err != nil {
		return 0, nil, err
	}

	datagram, err := NewDatagramFrom(bytes.NewReader(b))
	if err != nil {
		return 0, nil, err
	}

	n := copy(b, datagram.Data)
	return n, AddrFromSocks(datagram.ATYP, datagram.DstAddr, datagram.DstPort), nil
}

func (s *UDPConn) Read(b []byte) (int, error) {
	n, _, err := s.ReadFrom(b)
	return n, err
}

func (s *UDPConn) WriteTo(b []byte, addr net.Addr) (int, error) {
	socksAddr, err := AddrFromString(addr.String())
	if err != nil {
		return 0, err
	}

	atyp, host, port := socksAddr.Socks()
	datagram := NewDatagram(atyp, host, port, b)
	return s.PacketConn.WriteTo(datagram.Bytes(), s.ServerAddr)
}

func (s *UDPConn) Write(b []byte) (int, error) {
	if s.ConnectedAddr == nil {
		return 0, ErrNoConnectedAddress
	}

	return s.WriteTo(b, s.ConnectedAddr)
}

func (s *UDPConn) Close() error {
	s.TCPConn.Close()
	return s.PacketConn.Close()
}

func (s *UDPConn) RemoteAddr() net.Addr {
	return s.ConnectedAddr
}