// go-nc project main.go
package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"syscall"
)

// Handles TC connection and perform synchorinization:
// TCP -> Stdout and Stdin -> TCP
func tcp_con_handle(con net.Conn) {
	chan_to_stdout := stream_copy(con, os.Stdout)
	chan_to_remote := stream_copy(os.Stdin, con)
	select {
	case <-chan_to_stdout:
		log.Println("Remote connection is closed")
	case <-chan_to_remote:
		log.Println("Local program is terminated")
	}
}

// Performs copy operation between streams: os and tcp streams
func stream_copy(src io.Reader, dst io.Writer) <-chan int {
	buf := make([]byte, 1024)
	sync_channel := make(chan int)
	go func() {
		defer func() {
			if con, ok := dst.(net.Conn); ok {
				con.Close()
				log.Printf("Connection from %v is closed\n", con.RemoteAddr())
			}
			sync_channel <- 0 // Notify that processing is finished
		}()
		for {
			var nBytes int
			var err error
			nBytes, err = src.Read(buf)
			if err != nil {
				if err != io.EOF {
					log.Printf("Read error: %s\n", err)
				}
				break
			}
			_, err = dst.Write(buf[0:nBytes])
			if err != nil {
				log.Fatalf("Write error: %s\n", err)
			}
		}
	}()
	return sync_channel
}

//Accept data from UPD connection and copy it to the stream
func accept_from_udp_to_stream(src net.Conn, dst io.Writer) <-chan net.Addr {
	buf := make([]byte, 1024)
	sync_channel := make(chan net.Addr)
	con, ok := src.(*net.UDPConn)
	if !ok {
		log.Printf("Input must be UDP connection")
		return sync_channel
	}
	go func() {
		var remoteAddr net.Addr
		for {
			var nBytes int
			var err error
			var addr net.Addr
			nBytes, addr, err = con.ReadFromUDP(buf)
			if err != nil {
				if err != io.EOF {
					log.Printf("Read error: %s\n", err)
				}
				break
			}
			if remoteAddr == nil && remoteAddr != addr {
				remoteAddr = addr
				sync_channel <- remoteAddr
			}
			_, err = dst.Write(buf[0:nBytes])
			if err != nil {
				log.Fatalf("Write error: %s\n", err)
			}
		}
	}()
	log.Println("Exit write_from_udp_to_stream")
	return sync_channel
}

// Put input date from the stream to UDP connection
func put_from_stream_to_udp(src io.Reader, dst net.Conn, remoteAddr net.Addr) <-chan net.Addr {
	buf := make([]byte, 1024)
	sync_channel := make(chan net.Addr)
	go func() {
		for {
			var nBytes int
			var err error
			nBytes, err = src.Read(buf)
			if err != nil {
				if err != io.EOF {
					log.Printf("Read error: %s\n", err)
				}
				break
			}
			log.Println("Write to the remote address:", remoteAddr)
			if con, ok := dst.(*net.UDPConn); ok && remoteAddr != nil {
				_, err = con.WriteTo(buf[0:nBytes], remoteAddr)
			}
			if err != nil {
				log.Fatalf("Write error: %s\n", err)
			}
		}
	}()
	return sync_channel
}

// Handle UDP connection
func udp_con_handle(con net.Conn) {
	in_channel := accept_from_udp_to_stream(con, os.Stdout)
	log.Println("Waiting for remote connection")
	remoteAddr := <-in_channel
	log.Println("Connected from", remoteAddr)
	out_channel := put_from_stream_to_udp(os.Stdin, con, remoteAddr)
	select {
	case <-in_channel:
		log.Println("Remote connection is closed")
	case <-out_channel:
		log.Println("Local program is terminated")
	}
}

func main() {
	var destinationPort string
	var resizeCommand string
	var isUdp bool

	sigs := make(chan os.Signal, 1)

	flag.BoolVar(&isUdp, "u", false, "Use UDP instead of the default option of TCP.")
	flag.StringVar(&resizeCommand, "r", "", "Use UDP instead of the default option of TCP.")
	flag.Parse()
	if flag.NFlag() == 0 && flag.NArg() == 0 {
		fmt.Println("go-nc [-u] [-r command] [port]")
		flag.Usage()
		os.Exit(1)
	}

	if isUdp {
		log.Println("Protocol:", "udp")
	} else {
		log.Println("Protocol:", "tcp")
	}

	if flag.NArg() < 1 {
		fmt.Println("when you use -l option [port] is mandatory argument")
		os.Exit(1)
	}
	if _, err := strconv.Atoi(flag.Arg(0)); err != nil {
		log.Println("Destination port shall be not empty and have integer value")
		os.Exit(1)
	}
	destinationPort = fmt.Sprintf(":%v", flag.Arg(0))

	log.Println("Port:", destinationPort)

	if len(resizeCommand) != 0 {
		log.Println("Resize command:", resizeCommand)
		signal.Notify(sigs, syscall.SIGWINCH)
		go func() {
			for true{
				<-sigs
				log.Println("Window Resized")
				cmd := exec.Command(resizeCommand)
				cmd.Stdin = os.Stdin
				_, err := cmd.Output()

				log.Printf("err: %#v\n", err)
			}
		}()
	}

	if !isUdp {

		listener, err := net.Listen("tcp", destinationPort)
		if err != nil {
			log.Fatalln(err)
		}
		log.Println("Listening on", destinationPort)
		con, err := listener.Accept()
		if err != nil {
			log.Fatalln(err)
		}
		log.Println("Connect from", con.RemoteAddr())
		tcp_con_handle(con)

	} else {
		addr, err := net.ResolveUDPAddr("udp", destinationPort)
		if err != nil {
			log.Fatalln(err)
		}
		con, err := net.ListenUDP("udp", addr)
		if err != nil {
			log.Fatalln(err)
		}
		log.Println("Has been resolved UDP address:", addr)
		log.Println("Listening on", destinationPort)
		udp_con_handle(con)
	}
}