package main

import (
	"io"
	"log"
	"net"

	"github.com/gorilla/websocket"
)

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

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

	go func() {
		for {
			_, message, err := wsConn.ReadMessage()
			if err != nil {
				log.Println("WebSocket read error:", err)
				return
			}
			_, err = conn.Write(message)
			if err != nil {
				log.Println("TCP write error:", err)
				return
			}
		}
	}()

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

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)
	}
}