package tunnel

import (
	"fmt"
	"io"
	"log"
	"net"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

type Tunnel struct {
	wsURL     string
	port      int
	sessionID string
	ws        *websocket.Conn
	tcp       net.Conn
	stop      chan struct{}
	once      sync.Once
}

// New returns a tunnel that pipes traffic between the Fleet Forge server
// (WebSocket) and the local VNC server (127.0.0.1:port). It does not connect
// until Run is called.
func New(serverURL, credential, sessionID string, vncPort int) *Tunnel {
	wsURL := strings.Replace(serverURL, "https://", "wss://", 1)
	wsURL = strings.Replace(wsURL, "http://", "ws://", 1)
	wsURL = strings.TrimSuffix(wsURL, "/")

	q := url.Values{}
	q.Set("session_id", sessionID)
	q.Set("token", credential)

	full := wsURL + "/api/v1/remote/agent?" + q.Encode()

	return &Tunnel{
		wsURL:     full,
		port:      vncPort,
		sessionID: sessionID,
		stop:      make(chan struct{}),
	}
}

// Run dials the server WebSocket and the local VNC port, then pipes bytes in
// both directions until either side closes or Stop is called.
func (t *Tunnel) Run() error {
	ws, _, err := websocket.DefaultDialer.Dial(t.wsURL, nil)
	if err != nil {
		return fmt.Errorf("tunnel dial (server): %w", err)
	}
	defer ws.Close()
	t.ws = ws

	tcp, err := net.DialTimeout("tcp", fmt.Sprintf("127.0.0.1:%d", t.port), 10*time.Second)
	if err != nil {
		return fmt.Errorf("tunnel dial (local VNC): %w", err)
	}
	defer tcp.Close()
	t.tcp = tcp

	log.Printf("[tunnel] connected for session %s (vnc 127.0.0.1:%d)", t.sessionID, t.port)

	done := make(chan error, 2)

	go func() {
		for {
			_, data, err := ws.ReadMessage()
			if err != nil {
				done <- err
				return
			}
			if _, err := tcp.Write(data); err != nil {
				done <- err
				return
			}
		}
	}()

	go func() {
		buf := make([]byte, 8192)
		for {
			n, err := tcp.Read(buf)
			if n > 0 {
				if err := ws.WriteMessage(websocket.BinaryMessage, buf[:n]); err != nil {
					done <- err
					return
				}
			}
			if err != nil {
				done <- err
				return
			}
		}
	}()

	select {
	case err := <-done:
		return err
	case <-t.stop:
		return io.EOF
	}
}

// Stop closes the underlying connections, causing Run to return.
func (t *Tunnel) Stop() {
	t.once.Do(func() {
		close(t.stop)
		if t.ws != nil {
			_ = t.ws.Close()
		}
		if t.tcp != nil {
			_ = t.tcp.Close()
		}
	})
}