package main

import (
	"flag"
	"log"
	"net"
	"strings"
	"time"

	kcp "github.com/xtaci/kcp-go/v5"

	"omnisocketgo/cmd/internal/latencylog"
	"omnisocketgo/cmd/internal/server"
	"omnisocketgo/cmd/internal/transport"
)

const (
	kcpServerModeHub   = "hub"
	kcpServerModeRelay = "relay"
)

func main() {
	mode := flag.String("mode", kcpServerModeHub, "kcpserver mode: hub or relay")
	listenAddr := flag.String("listen", ":9002", "listen address; KCP listener in hub mode, UDP relay listener in relay mode")
	bindDevice := flag.String("bind-device", "", "optional Linux network device used when listening")
	logPath := flag.String("latency-log", "", "optional JSONL file path for latency timestamp logs")
	kcpTimestampDebugLogPath := flag.String("kcp-ts-debug-log", "", "optional JSONL file path for KCP packet kernel timestamp debug records")
	kcpSessionStatsLogPath := flag.String("kcp-session-stats-log", "", "optional JSONL file path for KCP session stats records")
	kcpSessionStatsInterval := flag.String("kcp-session-stats-interval", transport.DefaultKCPSessionStatsInterval.String(), "sampling interval for KCP session stats, for example 100ms")
	relayListenAddr := flag.String("relay-listen", "", "deprecated alias for -listen in relay mode")
	relayRemoteAddr := flag.String("relay-remote", "", "fixed remote UDP address used in relay mode")
	relayPeerAddr := flag.String("relay-peer", "", "deprecated alias for -relay-remote")
	flag.Parse()

	var relayRemoteFlagSet bool
	var relayPeerFlagSet bool
	var relayListenFlagSet bool
	flag.Visit(func(f *flag.Flag) {
		switch f.Name {
		case "relay-listen":
			relayListenFlagSet = true
		case "relay-remote":
			relayRemoteFlagSet = true
		case "relay-peer":
			relayPeerFlagSet = true
		}
	})

	switch {
	case relayRemoteFlagSet && relayPeerFlagSet && *relayRemoteAddr != *relayPeerAddr:
		log.Fatal("flags -relay-remote and -relay-peer must match when both are set")
	case *relayRemoteAddr == "" && *relayPeerAddr != "":
		*relayRemoteAddr = *relayPeerAddr
	}
	if relayPeerFlagSet {
		log.Printf("warning: flag -relay-peer is deprecated; use -relay-remote instead")
	}
	if relayListenFlagSet {
		if *relayListenAddr == "" {
			log.Fatal("flag -relay-listen must not be empty when set")
		}
		if *mode != kcpServerModeRelay {
			log.Fatal("flag -relay-listen may only be used in relay mode")
		}
		if *listenAddr != ":9002" && *listenAddr != *relayListenAddr {
			log.Fatal("flags -listen and -relay-listen must match when both are set in relay mode")
		}
		*listenAddr = *relayListenAddr
		log.Printf("warning: flag -relay-listen is deprecated; use -listen with -mode=relay instead")
	}

	statsInterval, err := transport.ParseKCPSessionStatsInterval(*kcpSessionStatsInterval)
	if err != nil {
		log.Fatalf("parse -kcp-session-stats-interval=%q: %v", *kcpSessionStatsInterval, err)
	}

	switch *mode {
	case kcpServerModeHub:
		if *relayRemoteAddr != "" {
			log.Fatal("flag -relay-remote may only be used in relay mode")
		}
		runHubServer(*listenAddr, *bindDevice, *logPath, *kcpTimestampDebugLogPath, *kcpSessionStatsLogPath, statsInterval)
	case kcpServerModeRelay:
		if *bindDevice != "" {
			log.Fatal("flag -bind-device is not supported in relay mode")
		}
		if *relayRemoteAddr == "" {
			log.Fatal("flag -relay-remote is required in relay mode")
		}
		runUDPRelayServer(*listenAddr, *relayRemoteAddr)
	default:
		log.Fatalf("unsupported -mode=%q; want %q or %q", *mode, kcpServerModeHub, kcpServerModeRelay)
	}
}

func runHubServer(listenAddr, bindDevice, logPath, packetDebugLogPath, sessionStatsLogPath string, statsInterval time.Duration) {
	hubOptions := make([]server.KCPOption, 0, 2)
	if logPath != "" {
		logger, err := latencylog.NewJSONLLogger(logPath)
		if err != nil {
			log.Fatalf("create latency logger %s: %v", logPath, err)
		}
		defer logger.Close()
		hubOptions = append(hubOptions, server.WithKCPLogger(logger))
	}

	var packetLogger transport.KCPPacketDebugLogger
	if packetDebugLogPath != "" {
		logger, err := transport.NewJSONLKCPPacketDebugLogger(packetDebugLogPath)
		if err != nil {
			log.Fatalf("create kcp packet debug logger %s: %v", packetDebugLogPath, err)
		}
		defer logger.Close()
		packetLogger = logger
	}
	if sessionStatsLogPath != "" {
		logger, err := transport.NewJSONLKCPSessionStatsLogger(sessionStatsLogPath)
		if err != nil {
			log.Fatalf("create kcp session stats logger %s: %v", sessionStatsLogPath, err)
		}
		defer logger.Close()
		hubOptions = append(hubOptions, server.WithKCPSessionStatsLogger(logger, statsInterval))
	}

	listener, packetConn, err := transport.ListenKCPSessions(listenAddr, bindDevice, packetLogger, latencylog.NodeRoleServer, "hub")
	if err != nil {
		log.Fatalf("listen kcp on %s: %v", listenAddr, err)
	}
	defer packetConn.Close()
	defer listener.Close()

	hub := server.NewKCPHub(hubOptions...)

	log.Printf("kcp hub listening on %s", listener.Addr())

	for {
		session, err := listener.AcceptKCP()
		if err != nil {
			if strings.Contains(err.Error(), "closed") {
				return
			}
			log.Printf("accept kcp session: %v", err)
			continue
		}

		go func(sess *kcp.UDPSession) {
			if serveErr := hub.ServeSession(sess); serveErr != nil {
				log.Printf("kcp session closed: %v", serveErr)
			}
		}(session)
	}
}

func runUDPRelayServer(listenAddr, remoteAddr string) {
	conn, err := net.ListenPacket("udp", listenAddr)
	if err != nil {
		log.Fatalf("listen udp relay on %s: %v", listenAddr, err)
	}
	defer conn.Close()

	remote, err := net.ResolveUDPAddr("udp", remoteAddr)
	if err != nil {
		log.Fatalf("resolve relay remote %s: %v", remoteAddr, err)
	}

	relay, err := server.NewUDPRelay(conn, remote)
	if err != nil {
		_ = conn.Close()
		log.Fatalf("create udp relay: %v", err)
	}

	log.Printf("udp relay listening on %s and forwarding to %s", conn.LocalAddr(), remote)
	if err := relay.Serve(); err != nil {
		log.Fatalf("udp relay stopped: %v", err)
	}
}