package proxy

import (
	"net"
	"strings"
	"time"
	"io"
)

type (
	TCPProxy interface {
		LocalPort() int
	}

	tcpProxy struct {
		listener      *net.TCPListener
		remoteAddress *net.TCPAddr
	}
)

func NewTCPProxy(local, remote string) (TCPProxy, error) {
	if strings.Index(local, ":") >= 0 {
		host, _, err := net.SplitHostPort(local)
		if err != nil {
			return nil, err
		}
		local = host
	}

	localResolved, err := net.ResolveTCPAddr("tcp", net.JoinHostPort(local, "0"))
	if err != nil {
		return nil, err
	}

	remoteResolved, err := net.ResolveTCPAddr("tcp", remote)
	if err != nil {
		return nil, err
	}

	listener, err := net.ListenTCP("tcp", localResolved)
	if err != nil {
		return nil, err
	}

	p := &tcpProxy{
		listener:      listener,
		remoteAddress: remoteResolved,
	}

	go p.listen()

	return p, nil
}

func (p *tcpProxy) listen() {
	defer p.listener.Close()

	for {
		conn, err := p.listener.AcceptTCP()

		if err != nil {
			if ne, ok := err.(net.Error); ok && ne.Temporary() {
				time.Sleep(10 * time.Millisecond)
				continue
			}
			return
		}

		go p.handleConnection(conn)
	}
}

func (p *tcpProxy) handleConnection(conn *net.TCPConn) {
	defer conn.Close()

	rc, err := net.DialTCP("tcp", nil, p.remoteAddress)
	if err != nil {
		return
	}

	defer rc.Close()

	go p.pipe(conn, rc)
	p.pipe(rc, conn)
}

func (p *tcpProxy) pipe(dst io.Writer, src io.Reader) {
	for {
		if _, err := io.Copy(dst, src); err != nil {
			break
		}
	}
}

func (p *tcpProxy) LocalPort() int {
	return p.listener.Addr().(*net.TCPAddr).Port
}