package main

import (
	"fmt"
	"io"
	"log"
	"net"

	"github.com/gorilla/websocket"
)

func handleConnection(tcpConn net.Conn, wsURL string) {
	defer tcpConn.Close()

	conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
	wsConn := NewWebSocketConn(conn)
	if err != nil {
		log.Println("WebSocket connection error:", err)
		return
	}
	defer wsConn.Close()

	waitChannel := make(chan bool)

	go func() {
		defer func() {
			waitChannel <- true
		}()
		_, _ = io.Copy(tcpConn, wsConn)
		// TODO print error
	}()

	go func() {
		defer func() {
			waitChannel <- true
		}()
		for {
			buffer := make([]byte, 1024)
			n, err := tcpConn.Read(buffer)
			if err != nil {
				if err != io.EOF {
					log.Println("TCP read error:", err)
				}
				return
			}
			err = conn.WriteMessage(websocket.BinaryMessage, buffer[:n])
			if err != nil {
				log.Println("WebSocket write error:", err)
				return
			}
		}
	}()

	<-waitChannel
}

func main() {
	tcpPort := "7000"
	wsURL := "ws://localhost:8000/ws" // Replace with your WebSocket server URL

	listener, err := net.Listen("tcp", ":"+tcpPort)
	if err != nil {
		log.Fatal(err)
	}
	defer listener.Close()

	log.Printf("TCP server listening on port %s", tcpPort)
	log.Printf("Forwarding connections to WebSocket server at %s", wsURL)

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Println(err)
			continue
		}
		go handleConnection(conn, wsURL)
	}
}

type WebSocketConn struct {
	conn *websocket.Conn
	buf  []byte
}

func NewWebSocketConn(conn *websocket.Conn) *WebSocketConn {
	return &WebSocketConn{conn: conn}
}

func (websocketConn *WebSocketConn) Read(p []byte) (n int, err error) {
	if len(websocketConn.buf) == 0 {
		_, message, err := websocketConn.conn.ReadMessage()
		fmt.Println("Got message ", message)
		if err != nil {
			return 0, err
		}
		websocketConn.buf = message
	}

	n = copy(p, websocketConn.buf)
	websocketConn.buf = websocketConn.buf[n:]

	log.Println("Read ", n, " bytes")
	return n, err
}

func (websocketConn *WebSocketConn) Write(p []byte) (n int, err error) {
	err = websocketConn.conn.WriteMessage(websocket.BinaryMessage, p)
	if err != nil {
		n = len(p)
	}
	return n, err
}

func (websocketConn *WebSocketConn) Close() error {
	return websocketConn.conn.Close()
}