package browser

import (
	"api/internal/database"
	"context"
	"fmt"
	"net"

	"github.com/coder/websocket"
)

type VNCProxy struct {
	session *database.BrowserSession
}

func NewVNCProxy(session *database.BrowserSession) *VNCProxy {
	return &VNCProxy{
		session: session,
	}
}

func (p *VNCProxy) ServeWebSocket(ctx context.Context, ws *websocket.Conn) error {
	// Connect to VNC server
	vncConn, err := net.Dial("tcp", fmt.Sprintf("localhost:%d", p.session.VNCPort))
	if err != nil {
		return fmt.Errorf("failed to connect to VNC server: %w", err)
	}
	defer vncConn.Close()

	// Create channels for error handling
	errChan := make(chan error, 2)

	// Forward WebSocket -> VNC
	go func() {
		for {
			messageType, data, err := ws.Read(ctx)
			if err != nil {
				errChan <- err
				return
			}
			if messageType == websocket.MessageBinary {
				_, err = vncConn.Write(data)
				if err != nil {
					errChan <- err
					return
				}
			}
		}
	}()

	// Forward VNC -> WebSocket
	go func() {
		buffer := make([]byte, 4096)
		for {
			n, err := vncConn.Read(buffer)
			if err != nil {
				errChan <- err
				return
			}
			err = ws.Write(ctx, websocket.MessageBinary, buffer[:n])
			if err != nil {
				errChan <- err
				return
			}
		}
	}()

	// Wait for any errors
	return <-errChan
}