package main

import (
	"bytes"
	"flag"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"time"
)

func makeRead(r io.Reader) (chan []byte, chan error) {
	read := make(chan []byte)
	cerr := make(chan error)
	go func() {
		for {
			b := make([]byte, 1024)
			n, err := r.Read(b)
			if err == io.EOF {
				cerr <- err
				return
			}
			if err != nil {
				continue
			}
			if n > 0 {
				read <- b[0:n]
			}
		}
	}()
	return read, cerr
}

func client(server, listen, connect string) error {
	var (
		read  chan []byte
		cerr  chan error
		write io.Writer
	)
	switch listen {
	case "-":
		read, cerr = makeRead(os.Stdin)
		write = os.Stdout
	default:
		l, err := net.Listen("tcp", listen)
		if err != nil {
			return err
		}
		conn, err := l.Accept()
		if err != nil {
			return err
		}
		read, cerr = makeRead(conn)
		write = conn
	}
	buf := bytes.NewBuffer([]byte(connect))
	resp, err := http.Post("http://"+server+"/new", "text/plain", buf)
	if err != nil {
		return err
	}
	bkey, _ := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	key := string(bkey)

	tick := time.NewTicker(time.Millisecond * 250)
	buf.Reset()
	cl := &http.Client{Transport: &http.Transport{MaxIdleConnsPerHost: 1}}
	req, _ := http.NewRequest("POST", "http://"+server+"/", nil)
	req.Header.Set("Content-Type", "application/octet-stream")
	req.Header.Set("Key", key)
	for {
		select {
		case <-tick.C:
			b := bytes.NewBuffer(nil)
			buf.WriteTo(b)
			req.Body = ioutil.NopCloser(b)
			resp, err := cl.Do(req)
			if err != nil {
				log.Println(err.Error())
				continue
			}
			_, err = io.Copy(write, resp.Body)
			if err != nil {
				log.Println(err.Error())
			}
			resp.Body.Close()
		case b := <-read:
			buf.Write(b)
		case <-cerr:
			req, _ := http.NewRequest("GET", "http://"+server+"/close", nil)
			req.Header.Set("Key", key)
			cl.Do(req)
			return nil
		}
	}
	return nil
}

func main() {
	var (
		server  string
		listen  string
		connect string
	)
	flag.StringVar(&server, "server", "", "server addr")
	flag.StringVar(&listen, "listen", ":8080", "client listen address (use '-' for stdin/stdout)")
	flag.StringVar(&connect, "connect", "", "remote connect addr")
	if err := client(server, listen, connect); err != nil {
		panic(err)
	}
}