package main

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

	"dragontcp/internal/protocol"
	"dragontcp/internal/wire"
)

var active int64

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,
	chunkMax int,
	bufferBytes int,
	chunkPollWait time.Duration,
	debug *serverDebug,
) {
	defer func() {
		<-slots
		atomic.AddInt64(&active, -1)
		_ = conn.Close()
	}()

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

	for {
		_ = conn.SetDeadline(time.Now().Add(30 * time.Second))
		req, err := wire.ReadRequest(conn)
		if err != nil {
			return
		}
		if err := processWireRequest(
			conn,
			req,
			token,
			allowPrivate,
			cache,
			tcpBuffer,
			manager,
			chunkMax,
			bufferBytes,
			chunkPollWait,
			debug,
		); err != nil {
			return
		}
	}
}

func main() {
	var (
		host           = flag.String("host", "0.0.0.0", "listen host")
		port           = flag.Int("port", 53, "listen port")
		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", 256, "compatibility buffer units; 256 = about 16 MiB per active session")
		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()

	fmt.Printf("DragonTCP Go server listening on %s\n", listenAddr)
	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)
	fmt.Printf("binary_transport=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 {
		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 peer=%v active_connections=%d", conn.RemoteAddr(), atomic.LoadInt64(&active))
			}
			go handle(
				conn,
				*token,
				*allowPrivate,
				cache,
				*tcpBuffer,
				slots,
				manager,
				*chunkMax,
				bufferBytes,
				*chunkPollWait,
				debug,
			)
		default:
			if debug.enabled {
				debug.errorf("REJECT peer=%v reason=max-connections", conn.RemoteAddr())
			}
			_ = conn.Close()
		}
	}
}