package main

import (
	"context"
	"crypto/subtle"
	"flag"
	"fmt"
	"net"
	"net/netip"
	"os"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"dragontcp/internal/protocol"
)

var active int64

// idleDeadline avoids a SetDeadline system call for every small protocol
// record. It refreshes halfway through the idle window, preserving idle-client
// cleanup while making persistent high-throughput lanes substantially cheaper.
type idleDeadline struct {
	conn    net.Conn
	timeout time.Duration
	next    time.Time
}

func newIdleDeadline(conn net.Conn, timeout time.Duration) *idleDeadline {
	return &idleDeadline{conn: conn, timeout: timeout}
}

func (d *idleDeadline) refresh() error {
	now := time.Now()
	if !d.next.IsZero() && now.Before(d.next.Add(-d.timeout/2)) {
		return nil
	}
	d.next = now.Add(d.timeout)
	return d.conn.SetDeadline(d.next)
}

type dnsEntry struct {
	ips     []netip.Addr
	expires time.Time
}

type dnsCache struct {
	mu      sync.RWMutex
	entries map[string]dnsEntry
	ttl     time.Duration
	max     int
}

func newDNSCache(ttl time.Duration, max int) *dnsCache {
	return &dnsCache{
		entries: make(map[string]dnsEntry),
		ttl:     ttl,
		max:     max,
	}
}

func (c *dnsCache) resolve(ctx context.Context, host string) ([]netip.Addr, error) {
	if ip, err := netip.ParseAddr(host); err == nil {
		return []netip.Addr{ip}, nil
	}

	now := time.Now()
	c.mu.RLock()
	entry, ok := c.entries[host]
	c.mu.RUnlock()
	if ok && now.Before(entry.expires) {
		return entry.ips, nil
	}

	ips, err := net.DefaultResolver.LookupNetIP(ctx, "ip", host)
	if err != nil {
		return nil, err
	}

	c.mu.Lock()
	if len(c.entries) >= c.max {
		// Simple bounded reset keeps the hot cache cheap and prevents growth.
		c.entries = make(map[string]dnsEntry, c.max)
	}
	c.entries[host] = dnsEntry{ips: ips, expires: now.Add(c.ttl)}
	c.mu.Unlock()

	return ips, nil
}

func tokenEqual(a, b string) bool {
	if len(a) != len(b) {
		return false
	}
	return subtle.ConstantTimeCompare([]byte(a), []byte(b)) == 1
}

var blockedSpecial = []netip.Prefix{
	netip.MustParsePrefix("0.0.0.0/8"),
	netip.MustParsePrefix("100.64.0.0/10"),
	netip.MustParsePrefix("192.0.0.0/24"),
	netip.MustParsePrefix("192.0.2.0/24"),
	netip.MustParsePrefix("198.18.0.0/15"),
	netip.MustParsePrefix("198.51.100.0/24"),
	netip.MustParsePrefix("203.0.113.0/24"),
	netip.MustParsePrefix("240.0.0.0/4"),
	netip.MustParsePrefix("2001:db8::/32"),
}

func addressAllowed(addr netip.Addr, allowPrivate bool) bool {
	if addr.IsUnspecified() || addr.IsMulticast() {
		return false
	}

	if allowPrivate {
		return true
	}

	if !addr.IsGlobalUnicast() ||
		addr.IsPrivate() ||
		addr.IsLoopback() ||
		addr.IsLinkLocalUnicast() {
		return false
	}

	for _, prefix := range blockedSpecial {
		if prefix.Contains(addr) {
			return false
		}
	}

	return true
}

func dialTarget(ctx context.Context, host string, port int, allowPrivate bool, cache *dnsCache, tcpBuffer int) (net.Conn, error) {
	ips, err := cache.resolve(ctx, host)
	if err != nil {
		return nil, err
	}

	var lastErr error
	var blocked []string

	d := net.Dialer{
		Timeout:   10 * time.Second,
		KeepAlive: 30 * time.Second,
	}

	for _, ip := range ips {
		if !addressAllowed(ip, allowPrivate) {
			blocked = append(blocked, ip.String())
			continue
		}

		addr := net.JoinHostPort(ip.String(), strconv.Itoa(port))
		conn, err := d.DialContext(ctx, "tcp", addr)
		if err == nil {
			protocol.TuneTCP(conn)
			protocol.TuneTCPBuffer(conn, tcpBuffer)
			return conn, nil
		}
		lastErr = err
	}

	if lastErr != nil {
		return nil, lastErr
	}
	if len(blocked) > 0 {
		return nil, fmt.Errorf("target resolves only to blocked addresses: %s", strings.Join(blocked, ","))
	}
	return nil, fmt.Errorf("no usable target address")
}

func handle(
	conn net.Conn,
	token string,
	allowPrivate bool,
	cache *dnsCache,
	tcpBuffer int,
	slots chan struct{},
	manager *streamManager,
	bhttpManager *bhttpSessionManager,
	xorManager *chunkManager,
	chunkMax int,
	bufferBytes int,
	xorBufferBytes int,
	chunkPollWait time.Duration,
	debug *serverDebug,
) {
	defer func() {
		<-slots
		atomic.AddInt64(&active, -1)
		_ = conn.Close()
	}()

	protocol.TuneTCP(conn)
	protocol.TuneTCPBuffer(conn, tcpBuffer)

	// One listener serves both wires and every startup-selected header profile.
	// sniffWire partitions the full first-byte space so B and X remain
	// unambiguous even when their legacy mode/UP bytes are masked.
	_ = conn.SetDeadline(time.Now().Add(30 * time.Second))
	conn, isXOR, headerMask, err := sniffWire(conn)
	if err != nil {
		return
	}
	if isXOR {
		if debug != nil && debug.enabled {
			debug.logf("WIRE peer=%v mode=xor header_mask=%02x", conn.RemoteAddr(), headerMask)
		}
		handleXOR(conn, headerMask, token, allowPrivate, cache, tcpBuffer, xorManager,
			chunkMax, xorBufferBytes, chunkPollWait, debug)
		return
	}
	clearPayload := false
	if profiled, ok := conn.(interface{ ClearPayload() bool }); ok {
		clearPayload = profiled.ClearPayload()
	}
	if debug != nil && debug.enabled {
		debug.logf("WIRE peer=%v mode=binary header_mask=%02x clear_payload=%t", conn.RemoteAddr(), headerMask, clearPayload)
	}

	handleBinary(conn, headerMask, clearPayload, token, allowPrivate, cache, tcpBuffer, manager,
		bhttpManager, chunkMax, bufferBytes, chunkPollWait, debug)
}

func acceptLoop(
	ln net.Listener,
	token string,
	allowPrivate bool,
	cache *dnsCache,
	tcpBuffer int,
	slots chan struct{},
	manager *streamManager,
	bhttpManager *bhttpSessionManager,
	xorManager *chunkManager,
	chunkMax int,
	bufferBytes int,
	xorBufferBytes int,
	chunkPollWait time.Duration,
	debug *serverDebug,
) {
	for {
		conn, err := ln.Accept()
		if err != nil {
			fmt.Fprintln(os.Stderr, "accept:", err)
			continue
		}

		select {
		case slots <- struct{}{}:
			atomic.AddInt64(&active, 1)
			if debug.enabled {
				debug.logf("ACCEPT local=%v peer=%v active_connections=%d", conn.LocalAddr(), conn.RemoteAddr(), atomic.LoadInt64(&active))
			}
			go handle(conn, token, allowPrivate, cache, tcpBuffer, slots, manager,
				bhttpManager, xorManager, chunkMax, bufferBytes, xorBufferBytes,
				chunkPollWait, debug)
		default:
			if debug.enabled {
				debug.errorf("REJECT peer=%v reason=max-connections", conn.RemoteAddr())
			}
			_ = conn.Close()
		}
	}
}

func main() {
	var (
		host           = flag.String("host", "0.0.0.0", "listen host")
		port           = flag.Int("port", 53, "listen port")
		portAlt        = flag.Int("port-alt", 80, "second simultaneous listen port; 0 disables")
		token          = flag.String("token", "", "optional shared token")
		maxConnections = flag.Int("max-connections", 20000, "max simultaneous tunnels")
		allowPrivate   = flag.Bool("allow-private", false, "allow private/loopback targets")
		dnsCacheTTL    = flag.Duration("dns-cache-ttl", 30*time.Second, "server DNS cache TTL")
		dnsCacheSize   = flag.Int("dns-cache-size", 4096, "maximum cached DNS hostnames")
		tcpBuffer      = flag.Int("tcp-buffer", 0, "optional TCP read/write buffer bytes; 0 keeps OS autotuning")
		chunkMax       = flag.Int("chunk-max", 1048576, "maximum adaptive chunk payload bytes (32 bytes to 1 MiB)")
		chunkBuffered  = flag.Int("chunk-buffered", 32, "per-session download buffer in 64 KiB units; 32 = about 2 MiB")
		chunkPollWait  = flag.Duration("chunk-poll-wait", 200*time.Millisecond, "server long-poll wait for chunk data")
		sessionTimeout = flag.Duration("chunk-session-timeout", 2*time.Minute, "idle chunk session timeout")
		debugEnabled   = flag.Bool("debug", false, "log session/connect/errors and periodic statistics")
		debugChunks    = flag.Bool("debug-chunks", false, "log every chunk protocol record; very verbose")
		debugStats     = flag.Duration("debug-stats-interval", 5*time.Second, "periodic debug statistics interval; 0 disables")
	)
	flag.Parse()

	if *chunkMax < 32 || *chunkMax > protocol.MaxChunkPayload {
		fmt.Fprintf(os.Stderr, "--chunk-max must be between 32 and %d\n", protocol.MaxChunkPayload)
		os.Exit(2)
	}
	if *chunkBuffered < 8 {
		fmt.Fprintln(os.Stderr, "--chunk-buffered must be at least 8")
		os.Exit(2)
	}

	listenAddr := net.JoinHostPort(*host, strconv.Itoa(*port))
	ln, err := net.Listen("tcp", listenAddr)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	defer ln.Close()
	listeners := []net.Listener{ln}
	fmt.Printf("DragonTCP Go server listening on %s\n", listenAddr)
	if *portAlt < 0 || *portAlt > 65535 {
		fmt.Fprintln(os.Stderr, "--port-alt must be between 0 and 65535")
		os.Exit(2)
	}
	if *portAlt != 0 && *portAlt != *port {
		altAddr := net.JoinHostPort(*host, strconv.Itoa(*portAlt))
		alt, altErr := net.Listen("tcp", altAddr)
		if altErr != nil {
			fmt.Fprintf(os.Stderr, "warning: secondary listener %s unavailable: %v\n", altAddr, altErr)
		} else {
			defer alt.Close()
			listeners = append(listeners, alt)
			fmt.Printf("DragonTCP Go server listening on %s\n", altAddr)
		}
	}
	fmt.Printf("max_connections=%d tcp_buffer=%d\n", *maxConnections, *tcpBuffer)

	slots := make(chan struct{}, *maxConnections)
	cache := newDNSCache(*dnsCacheTTL, *dnsCacheSize)
	debug := newServerDebug(*debugEnabled, *debugChunks, *debugStats)
	bufferBytes := *chunkBuffered * 65536
	if bufferBytes < 1024*1024 {
		bufferBytes = 1024 * 1024
	}
	if bufferBytes > 64*1024*1024 {
		bufferBytes = 64 * 1024 * 1024
	}
	manager := newStreamManager(*sessionTimeout, debug)
	bhttpManager := newBHTTPSessionManager(*sessionTimeout, *maxConnections)
	xorManager := newChunkManager(*sessionTimeout, debug)
	fmt.Printf("binary_transport=true bp_compat=true chunk_max=%d buffer_bytes=%d poll_wait=%s\n", *chunkMax, bufferBytes, chunkPollWait.String())
	if debug.enabled {
		fmt.Printf("debug=true debug_chunks=%t stats_interval=%s\n", debug.chunks, debug.statsEvery)
	}

	for _, listener := range listeners {
		go acceptLoop(listener, *token, *allowPrivate, cache, *tcpBuffer, slots,
			manager, bhttpManager, xorManager, *chunkMax, bufferBytes,
			bufferBytes, *chunkPollWait, debug)
	}
	select {}
}