// Command remotedesk-hysteria2-agent exposes a RemoteDesk datagram endpoint
// through the official Hysteria2 QUIC server. No TCP proxy or SOCKS5 mode is
// enabled: every authenticated Hysteria2 session is delivered to the local
// Agent UDP endpoint.
package main

import (
	"context"
	"crypto/tls"
	"errors"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/apernet/hysteria/core/v2/client"
	"github.com/apernet/hysteria/core/v2/server"
)

type passwordAuth struct{ password string }

func (a passwordAuth) Authenticate(_ net.Addr, auth string, _ uint64) (bool, string) {
	return auth != "" && auth == a.password, "remotedesk"
}

// localDatagramOutbound maps a Hysteria2 UDP request to a local UDP Agent.
// The public side remains QUIC/UDP; this socket is local-only and never a
// fallback transport for the remote connection.
type localDatagramOutbound struct{ target *net.UDPAddr }

func (o localDatagramOutbound) TCP(string) (net.Conn, error) {
	return nil, fmt.Errorf("TCP proxy is disabled for RemoteDesk Hysteria2")
}
func (o localDatagramOutbound) CheckUDP(string) error { return nil }
func (o localDatagramOutbound) UDP(_ string) (server.UDPConn, error) {
	conn, err := net.DialUDP("udp", nil, o.target)
	if err != nil {
		return nil, err
	}
	return &udpConn{conn: conn}, nil
}

type udpConn struct{ conn *net.UDPConn }

type datagramKind uint8

const (
	datagramOpaque datagramKind = iota
	datagramVideo
)

// classifyDatagram keeps the relay payload opaque. Only the RemoteDesk video
// framing has a local size contract; control and audio packets are forwarded
// without parsing so they can evolve independently.
func classifyDatagram(data []byte) (datagramKind, error) {
	if len(data) >= len(videoDatagramMagic) && string(data[:len(videoDatagramMagic)]) == videoDatagramMagic {
		if err := validateVideoDatagramSize(data); err != nil {
			return datagramVideo, err
		}
		return datagramVideo, nil
	}
	if len(data) >= len(audioDatagramMagic) && string(data[:len(audioDatagramMagic)]) == audioDatagramMagic {
		if err := validateAudioDatagramSize(data); err != nil {
			return datagramOpaque, err
		}
		if _, err := decodeAudioDatagram(data); err != nil {
			return datagramOpaque, err
		}
	}
	return datagramOpaque, nil
}

func prepareDatagram(data []byte) ([]byte, datagramKind, error) {
	if len(data) == 0 {
		return nil, datagramOpaque, fmt.Errorf("empty RemoteDesk datagram")
	}
	kind, err := classifyDatagram(data)
	if err != nil {
		return nil, kind, err
	}
	// Always copy a UDP read buffer before handing it to Hysteria2.
	return append([]byte(nil), data...), kind, nil
}

func (c *udpConn) ReadFrom(b []byte) (int, string, error) {
	n, addr, err := c.conn.ReadFromUDP(b)
	if addr == nil {
		return n, "", err
	}
	return n, addr.String(), err
}
func (c *udpConn) WriteTo(b []byte, _ string) (int, error) { return c.conn.Write(b) }
func (c *udpConn) Close() error                            { return c.conn.Close() }

func main() {
	mode := flag.String("mode", "server", "transport mode: server or client")
	listen := flag.String("listen", env("REMOTEDESK_HYSTERIA2_LISTEN", "0.0.0.0:39502"), "Hysteria2 UDP listen address")
	serverAddr := flag.String("server", env("REMOTEDESK_HYSTERIA2_SERVER", "127.0.0.1:39502"), "Hysteria2 server address (client mode)")
	password := flag.String("password", os.Getenv("REMOTEDESK_HYSTERIA2_PASSWORD"), "Hysteria2 password")
	certFile := flag.String("cert", os.Getenv("REMOTEDESK_HYSTERIA2_CERT"), "TLS certificate PEM")
	keyFile := flag.String("key", os.Getenv("REMOTEDESK_HYSTERIA2_KEY"), "TLS private key PEM")
	agent := flag.String("agent-udp", env("REMOTEDESK_AGENT_UDP", "127.0.0.1:39502"), "local Agent UDP endpoint")
	serverName := flag.String("server-name", env("REMOTEDESK_HYSTERIA2_SERVER_NAME", "localhost"), "TLS server name (client mode)")
	insecure := flag.Bool("insecure", false, "disable TLS certificate verification (client mode only)")
	ringPath := flag.String("encoded-ring", os.Getenv("REMOTEDESK_ENCODED_RING"), "local Rust encoded-access-unit ring path (optional)")
	ringStreamID := flag.Uint("ring-stream-id", 0, "RemoteDesk video stream ID for the encoded ring")
	ringGeneration := flag.Uint64("ring-generation", 0, "RemoteDesk encoded ring generation")
	agentControl := flag.String("windows-agent-control", os.Getenv("REMOTEDESK_WINDOWS_AGENT_CONTROL"), "loopback Windows Agent TCP control endpoint (optional)")
	agentPipe := flag.String("windows-agent-pipe", os.Getenv("REMOTEDESK_WINDOWS_AGENT_PIPE"), "authenticated local Windows Agent Named Pipe (recommended)")
	sessionID := flag.String("session-id", os.Getenv("REMOTEDESK_SESSION_ID"), "encoded desktop session ID")
	ringFPS := flag.Uint("ring-fps", 30, "encoded desktop frame rate (1-120)")
	allowSoftwareFallback := flag.Bool("allow-software-fallback", true, "allow Media Foundation software H.264 fallback")
	flag.Parse()
	if *mode == "client" {
		if *agentControl != "" && *agentPipe != "" {
			log.Fatal("--windows-agent-control and --windows-agent-pipe are mutually exclusive")
		}
		controlEndpoint := *agentPipe
		if controlEndpoint == "" {
			controlEndpoint = *agentControl
		}
		if uint64(*ringStreamID) > uint64(^uint32(0)) {
			log.Fatal("--ring-stream-id exceeds uint32")
		}
		if *ringFPS == 0 || *ringFPS > 120 {
			log.Fatal("--ring-fps must be between 1 and 120")
		}
		runClient(*listen, *serverAddr, *password, *agent, *serverName, *insecure, *ringPath, uint32(*ringStreamID), *ringGeneration, controlEndpoint, *sessionID, uint8(*ringFPS), *allowSoftwareFallback)
		return
	}
	if *mode != "server" {
		log.Fatalf("invalid --mode %q (must be server or client)", *mode)
	}
	if *password == "" || *certFile == "" || *keyFile == "" {
		log.Fatal("password, cert and key are required; refusing unauthenticated Hysteria2")
	}
	target, err := net.ResolveUDPAddr("udp", *agent)
	if err != nil {
		log.Fatalf("invalid --agent-udp: %v", err)
	}
	cert, err := tls.LoadX509KeyPair(*certFile, *keyFile)
	if err != nil {
		log.Fatalf("load TLS certificate: %v", err)
	}
	packetConn, err := net.ListenPacket("udp", *listen)
	if err != nil {
		log.Fatalf("listen Hysteria2 UDP: %v", err)
	}
	srv, err := server.NewServer(&server.Config{
		TLSConfig:     server.TLSConfig{Certificates: []tls.Certificate{cert}},
		Conn:          packetConn,
		Authenticator: passwordAuth{password: *password},
		Outbound:      localDatagramOutbound{target: target},
		DisableUDP:    false,
	})
	if err != nil {
		log.Fatalf("create Hysteria2 server: %v", err)
	}
	log.Printf("RemoteDesk Hysteria2 server listening on %s (UDP/QUIC only)", *listen)
	log.Printf("authenticated datagrams delivered to local Agent UDP %s", target)
	stop := make(chan os.Signal, 1)
	signal.Notify(stop, os.Interrupt, syscall.SIGTERM)
	done := make(chan error, 1)
	go func() { done <- srv.Serve() }()
	select {
	case sig := <-stop:
		log.Printf("shutting down on %s", sig)
		_ = srv.Close()
	case err := <-done:
		if err != nil && !strings.Contains(err.Error(), "closed") {
			log.Fatal(err)
		}
	}
	// Allow the QUIC listener to finish releasing its packet socket.
	time.Sleep(50 * time.Millisecond)
}

func runClient(localListen, serverAddr, password, remoteAgent, serverName string, insecure bool, ringPath string, ringStreamID uint32, ringGeneration uint64, agentControl, sessionID string, ringFPS uint8, allowSoftwareFallback bool) {
	if password == "" {
		log.Fatal("--password is required in client mode")
	}
	serverUDP, err := net.ResolveUDPAddr("udp", serverAddr)
	if err != nil {
		log.Fatalf("invalid --server: %v", err)
	}
	localUDP, err := net.ResolveUDPAddr("udp", localListen)
	if err != nil {
		log.Fatalf("invalid --listen: %v", err)
	}
	localConn, err := net.ListenUDP("udp", localUDP)
	if err != nil {
		log.Fatalf("listen local UDP: %v", err)
	}
	defer localConn.Close()
	hy, _, err := client.NewClient(&client.Config{
		ServerAddr: serverUDP,
		Auth:       password,
		TLSConfig: client.TLSConfig{
			ServerName:         serverName,
			InsecureSkipVerify: insecure,
		},
	})
	if err != nil {
		log.Fatalf("connect Hysteria2 server: %v", err)
	}
	defer hy.Close()
	stream, err := hy.UDP()
	if err != nil {
		log.Fatalf("open Hysteria2 UDP session: %v", err)
	}
	defer stream.Close()
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	var sendMu sync.Mutex
	log.Printf("RemoteDesk Hysteria2 client listening on %s (UDP only)", localConn.LocalAddr())
	log.Printf("authenticated Hysteria2 server %s; remote Agent %s", serverAddr, remoteAgent)
	var agentControlConnection net.Conn
	var audioDescriptor *audioRingDescriptor
	if ringPath == "" && agentControl != "" {
		connection, descriptor, err := openWindowsEncodedSession(agentControl, sessionID, ringFPS, allowSoftwareFallback)
		if err != nil {
			log.Fatalf("start Windows encoded desktop: %v", err)
		}
		agentControlConnection = connection
		ringPath = descriptor.RingPath
		ringGeneration = descriptor.Generation
		if ringStreamID == 0 {
			ringStreamID = 1
		}
		audioDescriptor = descriptor.AudioRing
		log.Printf("Windows encoded desktop session %s generation %d opened", sessionID, ringGeneration)
	}
	if agentControlConnection != nil {
		defer agentControlConnection.Close()
	}
	if ringPath != "" {
		if ringStreamID == 0 || ringGeneration == 0 {
			log.Fatal("--ring-stream-id and --ring-generation are required with --encoded-ring")
		}
		ring, err := openEncodedRing(ringPath)
		if err != nil {
			log.Fatalf("open encoded media ring: %v", err)
		}
		go func() {
			defer ring.Close()
			err := pumpEncodedRing(ctx, ring, ringStreamID, ringGeneration, func(data []byte) error {
				sendMu.Lock()
				defer sendMu.Unlock()
				return stream.Send(data, remoteAgent)
			})
			if err != nil && !errors.Is(err, context.Canceled) {
				log.Printf("encoded media ring stopped: %v", err)
			}
			cancel()
			if agentControlConnection != nil {
				_ = agentControlConnection.Close()
			}
			_ = localConn.Close()
		}()
	}
	if audioDescriptor != nil {
		audioRing, err := openAudioRing(audioDescriptor.RingPath)
		if err != nil {
			log.Printf("audio ring unavailable, continuing without audio: %v", err)
		} else {
			go func() {
				defer audioRing.Close()
				err := pumpAudioRing(ctx, audioRing, ringStreamID, audioDescriptor.Generation, func(data []byte) error {
					sendMu.Lock()
					defer sendMu.Unlock()
					return stream.Send(data, remoteAgent)
				})
				if err != nil && !errors.Is(err, context.Canceled) {
					log.Printf("audio media ring stopped: %v", err)
				}
			}()
		}
	}
	var peerMu sync.RWMutex
	var peer *net.UDPAddr
	go func() {
		for {
			data, _, err := stream.Receive()
			if err != nil {
				cancel()
				_ = localConn.Close()
				return
			}
			prepared, _, err := prepareDatagram(data)
			if err != nil {
				log.Printf("drop invalid datagram from Hysteria2: %v", err)
				continue
			}
			peerMu.RLock()
			currentPeer := peer
			if currentPeer != nil {
				_, _ = localConn.WriteToUDP(prepared, currentPeer)
			}
			peerMu.RUnlock()
		}
	}()
	buf := make([]byte, 64*1024)
	for {
		n, sourcePeer, err := localConn.ReadFromUDP(buf)
		if err != nil {
			return
		}
		peerMu.Lock()
		peer = sourcePeer
		peerMu.Unlock()
		prepared, _, err := prepareDatagram(buf[:n])
		if err != nil {
			log.Printf("drop invalid local datagram from %s: %v", sourcePeer, err)
			continue
		}
		// The Hysteria2 destination is metadata only. The server's outbound
		// implementation pins delivery to its configured local Agent UDP port.
		sendMu.Lock()
		err = stream.Send(prepared, remoteAgent)
		sendMu.Unlock()
		if err != nil {
			log.Printf("send Hysteria2 datagram from %s: %v", sourcePeer, err)
			return
		}
	}
}

func env(name, fallback string) string {
	if value := os.Getenv(name); value != "" {
		return value
	}
	return fallback
}