package main

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

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

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

func main() {
	listenAddr := flag.String("listen", ":9002", "KCP server listen address")
	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", "", "optional raw UDP relay listen address")
	relayPeerAddr := flag.String("relay-peer", "", "optional fixed raw UDP relay peer address")
	relayLearnPeer := flag.Bool("relay-learn-peer", false, "learn the relay peer address from the first inbound relay packet")
	flag.Parse()

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

	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 *kcpTimestampDebugLogPath != "" {
		logger, err := transport.NewJSONLKCPPacketDebugLogger(*kcpTimestampDebugLogPath)
		if err != nil {
			log.Fatalf("create kcp packet debug logger %s: %v", *kcpTimestampDebugLogPath, err)
		}
		defer logger.Close()
		packetLogger = logger
	}
	if *kcpSessionStatsLogPath != "" {
		logger, err := transport.NewJSONLKCPSessionStatsLogger(*kcpSessionStatsLogPath)
		if err != nil {
			log.Fatalf("create kcp session stats logger %s: %v", *kcpSessionStatsLogPath, 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...)

	if *relayPeerAddr != "" && *relayListenAddr == "" {
		log.Fatal("flag -relay-listen is required when -relay-peer is set")
	}
	if *relayLearnPeer && *relayListenAddr == "" {
		log.Fatal("flag -relay-listen is required when -relay-learn-peer is set")
	}
	if *relayListenAddr != "" {
		relayConn, err := net.ListenPacket("udp", *relayListenAddr)
		if err != nil {
			log.Fatalf("listen relay udp on %s: %v", *relayListenAddr, err)
		}
		defer relayConn.Close()

		var relayPeer net.Addr
		if *relayPeerAddr != "" {
			relayPeer, err = net.ResolveUDPAddr("udp", *relayPeerAddr)
			if err != nil {
				log.Fatalf("resolve relay peer %s: %v", *relayPeerAddr, err)
			}
		}

		hub.SetRelaySocket(relayConn, relayPeer, *relayLearnPeer)
		go func() {
			if serveErr := hub.ServeRelay(); serveErr != nil {
				log.Printf("kcp relay loop ended: %v", serveErr)
			}
		}()
		log.Printf("kcp relay listening on %s", relayConn.LocalAddr())
		if relayPeer != nil {
			log.Printf("kcp relay peer configured as %s", relayPeer)
		}
	}

	log.Printf("kcp server 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)
	}
}