package shadow

import (
	"fmt"
	"io"
	"net"
	"time"
)

type Server struct {
	ListenAddress string
	TargetAddress string
	FakeAddress   string
}

func NewServer(listenAddress string, targetAddress string, fakeAddress string) *Server {
	server := &Server{
		ListenAddress: listenAddress,
		TargetAddress: targetAddress,
		FakeAddress:   fakeAddress,
	}
	return server
}

func (s *Server) Start() {
	listen, err := net.Listen("tcp", s.ListenAddress)
	if err != nil {
		fmt.Printf("[Server] Start server error: %v\n", err)
		return
	}
	defer listen.Close()
	fmt.Printf("[Server] Listening at:%v\n", s.ListenAddress)
	for {
		conn, err := listen.Accept()
		if err != nil {
			fmt.Printf("[Server] Accept error: %v\n", err)
			continue
		}
		go handler(conn, s.TargetAddress, s.FakeAddress)
	}
}

func handler(conn net.Conn, targetAddress string, fakeAddress string) {

	//Process fake TLS handshake
	fmt.Println("[Server] Perform handshake")
	fakeConn, err := net.Dial("tcp", fakeAddress)
	if err != nil {
		fmt.Printf("[Server] Dial fake error : %v\n", err)
		return
	}
	waitCh := make(chan int, 1)

	go processHandshake(conn, fakeConn, waitCh)
	go processHandshake(fakeConn, conn, waitCh)

	<-waitCh

	//Clean up previous buffered data
	conn.SetDeadline(time.Now())
	conn.SetDeadline(time.Time{})

	realConnection, err := net.Dial("tcp", targetAddress)
	if err != nil {
		fmt.Printf("[Server] Dial target error : %v\n", err)
		return
	}
	if err != nil {
		return
	}

	p := &PackAppData{Conn: conn}
	go io.Copy(realConnection, p)
	go io.Copy(p, realConnection)
}

func processHandshake(src net.Conn, dst net.Conn, waitCh chan int) {
	buf := make([]byte, 32*1024)
	for {
		nr, er := src.Read(buf)
		if nr > 0 {

			header := ParseAndVerifyTLSHeader(buf[0:nr])
			nw, ew := dst.Write(buf[0:nr])
			if header != nil && header.Type == ChangeCipherSpec {
				//fmt.Println(header.toString())
				fmt.Println("[Server] handshake complete")
				if header.ChangeCipherSpecNext == AppData {
					dst.Close()
					waitCh <- 1
				} else {
					src.Close()
					waitCh <- 1
					return
				}
				break
			}
			if nw < 0 || nr < nw {
				nw = 0
				if ew == nil {
					//fmt.Printf("ERR1 %v \n", ew)
				}
			}
			if ew != nil {
				//fmt.Printf("ERR2 %v \n", ew)
				break
			}
			if nr != nw {
				//fmt.Printf("ERR3 %v \n", "shortwrite")
				break
			}
		}
		if er != nil {
			if er != io.EOF {
				//fmt.Printf("ERR4 %v \n", er)
			}
			break
		}
	}
	waitCh <- 1
}