package cachecon

import (
	"fmt"
	"golang.org/x/net/proxy"
	"io"
	"log"
	"net"
)

type Config struct {
	socksAddr, dstAddr string
}

type Server struct {
	N int

	conns chan *net.Conn
}

func New(config Config) (*Server, error) {
	cache := new(Server)

	cache.N = 10
	cache.conns = make(chan *net.Conn, cache.N)

	dialer, err := proxy.SOCKS5("tcp", config.socksAddr, nil, &net.Dialer{})
	if err != nil {
		return nil, err
	}

	for i := 0; i < cache.N; i++ {
		conn, err := dialer.Dial("tcp", config.dstAddr)
		if err != nil {
			log.Printf("Dst connection error: %s", err)
			return nil, fmt.Errorf("Can't connect to destination address")
		}

		cache.conns <- &conn
	}
	return cache, nil
}

func (cache Server) Dial() (*net.Conn, error) {
	select {
	case freeConn := <-cache.conns:
		return freeConn, nil
	default:
		return nil, fmt.Errorf("No free connnections")
	}
}

func broker(src, dst net.Conn, srcClosed chan struct{}) {
	_, err := io.Copy(dst, src)

	if err != nil {
		log.Printf("Copy error: %s", err)
	}

	if err := src.Close(); err != nil {
		log.Printf("Close error: %s", err)
	}

	srcClosed <- struct{}{}
}

func proxyTcp(serverConn, clientConn *net.TCPConn) {
	serverClosed := make(chan struct{}, 1)
	clientClosed := make(chan struct{}, 1)

	go broker(clientConn, serverConn, clientClosed)
	go broker(serverConn, clientConn, serverClosed)

	var waitFor chan struct{}

	select {
	case <-clientClosed:
		serverConn.CloseRead()
		waitFor = serverClosed
	case <-serverClosed:
		serverConn.CloseRead()
		clientConn.CloseRead()
		waitFor = clientClosed
	}

	<-waitFor
}

func handleConnection(cacheCon *Server, srcConn *net.Conn) {
	dstConn, err := (*cacheCon).Dial()
	if err != nil {
		log.Panic(err)
	}

	src, okSrc := (*srcConn).(*net.TCPConn)
	dst, okDst := (*dstConn).(*net.TCPConn)

	if !okSrc || !okDst {
		log.Panic()
	}

	proxyTcp(dst, src)
}

func (cacheCon *Server) Serve(ln net.Listener) {
	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Panic(err)
			break
		}
		go handleConnection(cacheCon, &conn)
	}
}

func (cacheCon *Server) ListenAndServe(servePort string) error {
	ln, err := net.Listen("tcp", servePort)
	if err != nil {
		return err
	}
	cacheCon.Serve(ln)
	return nil
}