package main

import (
	"context"
	"fmt"
	"log"
	"net"
	"os"
	"io"
	"time"

	"tailscale.com/tsnet"
)

var globalTSServer *tsnet.Server

func startTailscale(dataDir, nodeID, headscaleURL, authKey string) (string, error) {
	// Configure tsnet to use our Headscale server
	os.Setenv("TS_AUTHKEY", authKey)
	os.Setenv("TS_CONTROL_URL", headscaleURL)

	s := &tsnet.Server{
		Hostname: nodeID,
		Dir:      dataDir,
		Logf:     log.Printf,
	}

	if err := s.Start(); err != nil {
		return "", fmt.Errorf("tailscale start: %w", err)
	}

	globalTSServer = s

	// Wait for Tailscale to come up and retrieve IP
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	lc, err := s.LocalClient()
	if err != nil {
		return "", fmt.Errorf("tailscale local client: %w", err)
	}

	var tailscaleIP string
	for {
		status, err := lc.Status(ctx)
		if err != nil {
			return "", fmt.Errorf("tailscale status: %w", err)
		}
		if status.Self != nil && len(status.Self.TailscaleIPs) > 0 {
			tailscaleIP = status.Self.TailscaleIPs[0].String()
			break
		}
		select {
		case <-ctx.Done():
			return "", fmt.Errorf("tailscale IP timeout")
		case <-time.After(1 * time.Second):
		}
	}

	log.Printf("Tailscale started with IP: %s", tailscaleIP)
	return tailscaleIP, nil
}

func startTailscaleProxy(port int) (net.Listener, error) {
	if globalTSServer == nil {
		return nil, fmt.Errorf("tailscale server not started")
	}
	ln, err := globalTSServer.Listen("tcp", fmt.Sprintf(":%d", port))
	if err != nil {
		return nil, fmt.Errorf("tailscale listen on port %d: %w", port, err)
	}
	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				log.Printf("tailscale proxy accept error on port %d: %v", port, err)
				return
			}
			go handleProxyConn(conn, port)
		}
	}()
	log.Printf("Tailscale proxy started on port %d", port)
	return ln, nil
}

func handleProxyConn(src net.Conn, port int) {
	defer src.Close()
	dst, err := net.Dial("tcp", fmt.Sprintf("localhost:%d", port))
	if err != nil {
		log.Printf("tailscale proxy dial localhost:%d error: %v", port, err)
		return
	}
	defer dst.Close()

	done := make(chan struct{}, 2)
	go func() {
		_, _ = io.Copy(dst, src)
		done <- struct{}{}
	}()
	go func() {
		_, _ = io.Copy(src, dst)
		done <- struct{}{}
	}()
	<-done
}