package relay

import (
	"io"
	"log"
	"net"
	"net/http"
	"sync"
	"time"
)

// Server bridges WS connections on RelayPath to the upstream VNC TCP server.
type Server struct {
	vncAddr string
}

// New returns a relay Server targeting the given VNC TCP address.
func New(vncAddr string) *Server {
	return &Server{vncAddr: vncAddr}
}

// ServeHTTP upgrades to WebSocket and bridges to the VNC server.
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	ws, err := handshake(w, r)
	if err != nil {
		// handshake failed before hijack; respond with an error.
		http.Error(w, "websocket handshake failed: "+err.Error(), http.StatusBadRequest)
		return
	}
	defer ws.close()

	vnc, err := net.DialTimeout("tcp", s.vncAddr, 10*time.Second)
	if err != nil {
		log.Printf("relay: dial vnc %s failed: %v", s.vncAddr, err)
		ws.writeClose()
		return
	}
	defer vnc.Close()

	bridge(ws, vnc)
}

// bridge pumps bytes between the WebSocket and the TCP connection until either side closes.
func bridge(ws *wsConn, vnc net.Conn) {
	var wg sync.WaitGroup
	wg.Add(2)

	// TCP -> WS
	go func() {
		defer wg.Done()
		buf := make([]byte, 4096)
		for {
			n, err := vnc.Read(buf)
			if n > 0 {
				if werr := ws.writeBinary(buf[:n]); werr != nil {
					return
				}
			}
			if err != nil {
				if err != io.EOF {
					log.Printf("relay: vnc read: %v", err)
				}
				ws.writeClose()
				return
			}
		}
	}()

	// WS -> TCP
	go func() {
		defer wg.Done()
		for {
			opcode, payload, err := ws.readFrame()
			if err != nil {
				if err != io.EOF && !isClosedConnErr(err) {
					log.Printf("relay: ws read: %v", err)
				}
				_ = vnc.Close()
				return
			}
			switch opcode {
			case opBinary, opText, opContinuation:
				if len(payload) > 0 {
					if _, err := vnc.Write(payload); err != nil {
						return
					}
				}
			case opPing:
				_ = ws.writePong(payload)
			case opPong:
				// ignore
			case opClose:
				ws.writeClose()
				_ = vnc.Close()
				return
			}
		}
	}()

	wg.Wait()
}

func isClosedConnErr(err error) bool {
	if err == nil {
		return false
	}
	s := err.Error()
	return contains(s, "use of closed network connection") ||
		contains(s, "connection reset") ||
		contains(s, "EOF")
}

func contains(s, sub string) bool {
	return len(s) >= len(sub) && (s == sub || indexOf(s, sub) >= 0)
}

func indexOf(s, sub string) int {
	for i := 0; i+len(sub) <= len(s); i++ {
		if s[i:i+len(sub)] == sub {
			return i
		}
	}
	return -1
}