package dialer

import (
	"context"
	"errors"
	"net"
	"net/netip"
	"slices"
	"syscall"
	"time"

	"github.com/sagernet/sing-box/adapter"
	"github.com/sagernet/sing-box/common/listener"
	C "github.com/sagernet/sing-box/constant"
	"github.com/sagernet/sing-box/option"
	"github.com/sagernet/sing-box/service/powerreport"
	"github.com/sagernet/sing/common"
	"github.com/sagernet/sing/common/bufio"
	"github.com/sagernet/sing/common/control"
	E "github.com/sagernet/sing/common/exceptions"
	M "github.com/sagernet/sing/common/metadata"
	N "github.com/sagernet/sing/common/network"
	"github.com/sagernet/sing/service"

	"github.com/database64128/tfo-go/v2"
)

var (
	_ ParallelInterfaceDialer = (*DefaultDialer)(nil)
	_ UDPListener             = (*DefaultDialer)(nil)
)

type DefaultDialer struct {
	dialer4                tfo.Dialer
	dialer6                tfo.Dialer
	udpDialer4             net.Dialer
	udpDialer6             net.Dialer
	udpListener            net.ListenConfig
	udpAddr4               string
	udpAddr6               string
	netns                  string
	autoDetectBindFunc     control.Func
	connectionManager      adapter.ConnectionManager
	networkManager         adapter.NetworkManager
	powerManager           *powerreport.Manager
	outboundManager        adapter.OutboundManager
	dnsTransportManager    adapter.DNSTransportManager
	networkStrategy        *C.NetworkStrategy
	defaultNetworkStrategy bool
	networkType            []C.InterfaceType
	fallbackNetworkType    []C.InterfaceType
	networkFallbackDelay   time.Duration
	networkLastFallback    common.TypedValue[time.Time]
}

func NewDefault(ctx context.Context, options option.DialerOptions) (*DefaultDialer, error) {
	connectionManager := service.FromContext[adapter.ConnectionManager](ctx)
	networkManager := service.FromContext[adapter.NetworkManager](ctx)
	platformInterface := service.FromContext[adapter.PlatformInterface](ctx)

	var (
		dialer                 net.Dialer
		listener               net.ListenConfig
		interfaceFinder        control.InterfaceFinder
		networkStrategy        *C.NetworkStrategy
		defaultNetworkStrategy bool
		networkType            []C.InterfaceType
		fallbackNetworkType    []C.InterfaceType
		networkFallbackDelay   time.Duration
		autoDetectBindFunc     control.Func
	)
	if networkManager != nil {
		interfaceFinder = networkManager.InterfaceFinder()
	} else {
		interfaceFinder = control.NewDefaultInterfaceFinder()
	}
	if options.BindInterface != "" {
		if !(C.IsLinux || C.IsDarwin || C.IsWindows) {
			return nil, E.New("`bind_interface` is only supported on Linux, macOS and Windows")
		}
		bindFunc := control.BindToInterface(interfaceFinder, options.BindInterface, -1)
		dialer.Control = control.Append(dialer.Control, bindFunc)
		listener.Control = control.Append(listener.Control, bindFunc)
	}
	if options.RoutingMark > 0 {
		if !C.IsLinux {
			return nil, E.New("`routing_mark` is only supported on Linux")
		}
		dialer.Control = control.Append(dialer.Control, setMarkWrapper(networkManager, uint32(options.RoutingMark), false))
		listener.Control = control.Append(listener.Control, setMarkWrapper(networkManager, uint32(options.RoutingMark), false))
	}
	disableDefaultBind := options.BindInterface != "" || options.Inet4BindAddress != nil || options.Inet6BindAddress != nil
	if disableDefaultBind || options.TCPFastOpen {
		if options.NetworkStrategy != nil || len(options.NetworkType) > 0 && options.FallbackNetworkType == nil && options.FallbackDelay == 0 {
			return nil, E.New("`network_strategy` is conflict with `bind_interface`, `inet4_bind_address`, `inet6_bind_address` and `tcp_fast_open`")
		}
	}

	if networkManager != nil {
		defaultOptions := networkManager.DefaultOptions()
		if defaultOptions.BindInterface != "" && !disableDefaultBind {
			bindFunc := control.BindToInterface(networkManager.InterfaceFinder(), defaultOptions.BindInterface, -1)
			dialer.Control = control.Append(dialer.Control, bindFunc)
			listener.Control = control.Append(listener.Control, bindFunc)
		} else if networkManager.AutoDetectInterface() && !disableDefaultBind {
			if platformInterface != nil && platformInterface.UsePlatformNetworkInterfaces() {
				networkStrategy = (*C.NetworkStrategy)(options.NetworkStrategy)
				networkType = common.Map(options.NetworkType, option.InterfaceType.Build)
				fallbackNetworkType = common.Map(options.FallbackNetworkType, option.InterfaceType.Build)
				if networkStrategy == nil && len(networkType) == 0 && len(fallbackNetworkType) == 0 {
					networkStrategy = defaultOptions.NetworkStrategy
					networkType = defaultOptions.NetworkType
					fallbackNetworkType = defaultOptions.FallbackNetworkType
				}
				networkFallbackDelay = time.Duration(options.FallbackDelay)
				if networkFallbackDelay == 0 && defaultOptions.FallbackDelay != 0 {
					networkFallbackDelay = defaultOptions.FallbackDelay
				}
				if networkStrategy == nil {
					networkStrategy = common.Ptr(C.NetworkStrategyDefault)
					defaultNetworkStrategy = true
				}
				bindFunc := networkManager.ProtectFunc()
				dialer.Control = control.Append(dialer.Control, bindFunc)
				listener.Control = control.Append(listener.Control, bindFunc)
			} else {
				bindFunc := networkManager.AutoDetectInterfaceFunc()
				dialer.Control = control.Append(dialer.Control, bindFunc)
				autoDetectBindFunc = bindFunc
			}
		}
		if options.RoutingMark == 0 && defaultOptions.RoutingMark != 0 {
			dialer.Control = control.Append(dialer.Control, setMarkWrapper(networkManager, defaultOptions.RoutingMark, true))
			listener.Control = control.Append(listener.Control, setMarkWrapper(networkManager, defaultOptions.RoutingMark, true))
		}
	}
	if networkManager != nil {
		markFunc := networkManager.AutoRedirectOutputMarkFunc()
		dialer.Control = control.Append(dialer.Control, markFunc)
		listener.Control = control.Append(listener.Control, markFunc)
	}
	if options.ReuseAddr {
		listener.Control = control.Append(listener.Control, control.ReuseAddr())
	}
	if options.ProtectPath != "" {
		dialer.Control = control.Append(dialer.Control, control.ProtectPath(options.ProtectPath))
		listener.Control = control.Append(listener.Control, control.ProtectPath(options.ProtectPath))
	}
	if options.BindAddressNoPort {
		if !C.IsLinux {
			return nil, E.New("`bind_address_no_port` is only supported on Linux")
		}
		dialer.Control = control.Append(dialer.Control, control.BindAddressNoPort())
	}
	if options.ConnectTimeout != 0 {
		dialer.Timeout = time.Duration(options.ConnectTimeout)
	} else {
		dialer.Timeout = C.TCPConnectTimeout
	}
	if options.DisableTCPKeepAlive {
		dialer.KeepAlive = -1
		dialer.KeepAliveConfig.Enable = false
	} else if options.TCPKeepAliveSystemDefaults {
		dialer.KeepAliveConfig = net.KeepAliveConfig{
			Enable:   true,
			Idle:     -1,
			Interval: -1,
			Count:    -1,
		}
	} else {
		keepIdle := time.Duration(options.TCPKeepAlive)
		if keepIdle == 0 {
			keepIdle = C.TCPKeepAliveInitial
		}
		keepInterval := time.Duration(options.TCPKeepAliveInterval)
		if keepInterval == 0 {
			keepInterval = C.TCPKeepAliveInterval
		}
		dialer.KeepAliveConfig = net.KeepAliveConfig{
			Enable:   true,
			Idle:     keepIdle,
			Interval: keepInterval,
		}
	}
	var udpFragment bool
	if options.UDPFragment != nil {
		udpFragment = *options.UDPFragment
	} else {
		udpFragment = options.UDPFragmentDefault
	}
	if !udpFragment {
		dialer.Control = control.Append(dialer.Control, control.DisableUDPFragment())
		listener.Control = control.Append(listener.Control, control.DisableUDPFragment())
	}
	var (
		dialer4    = dialer
		udpDialer4 = dialer
		udpAddr4   string
	)
	if options.Inet4BindAddress != nil {
		bindAddr := options.Inet4BindAddress.Build(netip.IPv4Unspecified())
		dialer4.LocalAddr = &net.TCPAddr{IP: bindAddr.AsSlice()}
		udpDialer4.LocalAddr = &net.UDPAddr{IP: bindAddr.AsSlice(), Port: int(options.UDPBindPort)}
		udpAddr4 = M.SocksaddrFrom(bindAddr, options.UDPBindPort).String()
	} else if options.UDPBindPort != 0 {
		udpDialer4.LocalAddr = &net.UDPAddr{IP: net.IPv4zero, Port: int(options.UDPBindPort)}
		udpAddr4 = M.SocksaddrFrom(netip.IPv4Unspecified(), options.UDPBindPort).String()
	}
	var (
		dialer6    = dialer
		udpDialer6 = dialer
		udpAddr6   string
	)
	if options.Inet6BindAddress != nil {
		bindAddr := options.Inet6BindAddress.Build(netip.IPv6Unspecified())
		dialer6.LocalAddr = &net.TCPAddr{IP: bindAddr.AsSlice()}
		udpDialer6.LocalAddr = &net.UDPAddr{IP: bindAddr.AsSlice(), Port: int(options.UDPBindPort)}
		udpAddr6 = M.SocksaddrFrom(bindAddr, options.UDPBindPort).String()
	} else if options.UDPBindPort != 0 {
		udpDialer6.LocalAddr = &net.UDPAddr{IP: net.IPv6unspecified, Port: int(options.UDPBindPort)}
		udpAddr6 = M.SocksaddrFrom(netip.IPv6Unspecified(), options.UDPBindPort).String()
	}
	if options.TCPMultiPath {
		dialer4.SetMultipathTCP(true)
	}
	tcpDialer4 := tfo.Dialer{Dialer: dialer4, DisableTFO: !options.TCPFastOpen}
	tcpDialer6 := tfo.Dialer{Dialer: dialer6, DisableTFO: !options.TCPFastOpen}
	return &DefaultDialer{
		dialer4:                tcpDialer4,
		dialer6:                tcpDialer6,
		udpDialer4:             udpDialer4,
		udpDialer6:             udpDialer6,
		udpListener:            listener,
		udpAddr4:               udpAddr4,
		udpAddr6:               udpAddr6,
		netns:                  options.NetNs,
		autoDetectBindFunc:     autoDetectBindFunc,
		connectionManager:      connectionManager,
		networkManager:         networkManager,
		powerManager:           service.FromContext[*powerreport.Manager](ctx),
		outboundManager:        service.FromContext[adapter.OutboundManager](ctx),
		dnsTransportManager:    service.FromContext[adapter.DNSTransportManager](ctx),
		networkStrategy:        networkStrategy,
		defaultNetworkStrategy: defaultNetworkStrategy,
		networkType:            networkType,
		fallbackNetworkType:    fallbackNetworkType,
		networkFallbackDelay:   networkFallbackDelay,
	}, nil
}

func setMarkWrapper(networkManager adapter.NetworkManager, mark uint32, isDefault bool) control.Func {
	if networkManager == nil {
		return control.RoutingMark(mark)
	}
	return func(network, address string, conn syscall.RawConn) error {
		if networkManager.AutoRedirectOutputMark() != 0 {
			if isDefault {
				return E.New("`route.default_mark` is conflict with `tun.auto_redirect`")
			} else {
				return E.New("`routing_mark` is conflict with `tun.auto_redirect`")
			}
		}
		return control.RoutingMark(mark)(network, address, conn)
	}
}

func (d *DefaultDialer) DialContext(ctx context.Context, network string, address M.Socksaddr) (net.Conn, error) {
	if !address.IsValid() {
		return nil, E.New("invalid address")
	} else if address.IsDomain() {
		return nil, E.New("domain not resolved")
	}
	if d.networkStrategy == nil {
		conn, err := listener.ListenNetworkNamespace[net.Conn](ctx, d.netns, func() (net.Conn, error) {
			switch N.NetworkName(network) {
			case N.NetworkUDP:
				if !address.IsIPv6() {
					return d.udpDialer4.DialContext(ctx, network, address.String())
				} else {
					return d.udpDialer6.DialContext(ctx, network, address.String())
				}
			}
			if !address.IsIPv6() {
				return DialSlowContext(&d.dialer4, ctx, network, address)
			} else {
				return DialSlowContext(&d.dialer6, ctx, network, address)
			}
		})
		return d.trackConn(ctx, address, conn, err)
	} else {
		return d.DialParallelInterface(ctx, network, address, d.networkStrategy, d.networkType, d.fallbackNetworkType, d.networkFallbackDelay)
	}
}

func (d *DefaultDialer) DialParallelInterface(ctx context.Context, network string, address M.Socksaddr, strategy *C.NetworkStrategy, interfaceType []C.InterfaceType, fallbackInterfaceType []C.InterfaceType, fallbackDelay time.Duration) (net.Conn, error) {
	if strategy == nil {
		strategy = d.networkStrategy
	}
	if strategy == nil {
		return d.DialContext(ctx, network, address)
	}
	if len(interfaceType) == 0 {
		interfaceType = d.networkType
	}
	if len(fallbackInterfaceType) == 0 {
		fallbackInterfaceType = d.fallbackNetworkType
	}
	if fallbackDelay == 0 {
		fallbackDelay = d.networkFallbackDelay
	}
	var dialer net.Dialer
	if N.NetworkName(network) == N.NetworkTCP {
		dialer = d.dialer4.Dialer
	} else {
		dialer = d.udpDialer4
	}
	fastFallback := time.Since(d.networkLastFallback.Load()) < C.TCPTimeout
	var (
		conn      net.Conn
		isPrimary bool
		err       error
	)
	if !fastFallback {
		conn, isPrimary, err = d.dialParallelInterface(ctx, dialer, network, address.String(), *strategy, interfaceType, fallbackInterfaceType, fallbackDelay)
	} else {
		conn, isPrimary, err = d.dialParallelInterfaceFastFallback(ctx, dialer, network, address.String(), *strategy, interfaceType, fallbackInterfaceType, fallbackDelay, d.networkLastFallback.Store)
	}
	if err != nil {
		// bind interface failed on legacy xiaomi systems
		if d.defaultNetworkStrategy && errors.Is(err, syscall.EPERM) {
			d.networkStrategy = nil
			return d.DialContext(ctx, network, address)
		} else {
			return nil, err
		}
	}
	if !fastFallback && !isPrimary {
		d.networkLastFallback.Store(time.Now())
	}
	return d.trackConn(ctx, address, conn, nil)
}

func (d *DefaultDialer) ListenPacket(ctx context.Context, destination M.Socksaddr) (net.PacketConn, error) {
	if d.networkStrategy == nil {
		packetConn, err := listener.ListenNetworkNamespace[net.PacketConn](ctx, d.netns, func() (net.PacketConn, error) {
			listenConfig := d.udpListener
			if d.autoDetectBindFunc != nil {
				listenConfig.Control = control.Append(listenConfig.Control, func(network, address string, conn syscall.RawConn) error {
					if destination.Addr.IsValid() {
						return d.autoDetectBindFunc(network, destination.String(), conn)
					}
					return d.autoDetectBindFunc(network, address, conn)
				})
			}
			if destination.IsIPv6() {
				return listenConfig.ListenPacket(ctx, N.NetworkUDP, d.udpAddr6)
			} else if destination.IsIPv4() && !destination.Addr.IsUnspecified() {
				return listenConfig.ListenPacket(ctx, N.NetworkUDP+"4", d.udpAddr4)
			} else {
				return listenConfig.ListenPacket(ctx, N.NetworkUDP, d.udpAddr4)
			}
		})
		return d.trackPacketConn(ctx, destination, packetConn, err)
	} else {
		return d.ListenSerialInterfacePacket(ctx, destination, d.networkStrategy, d.networkType, d.fallbackNetworkType, d.networkFallbackDelay)
	}
}

func (d *DefaultDialer) DialerForICMPDestination(destination netip.Addr) net.Dialer {
	if !destination.Is6() {
		return d.dialer4.Dialer
	} else {
		return d.dialer6.Dialer
	}
}

func (d *DefaultDialer) ListenSerialInterfacePacket(ctx context.Context, destination M.Socksaddr, strategy *C.NetworkStrategy, interfaceType []C.InterfaceType, fallbackInterfaceType []C.InterfaceType, fallbackDelay time.Duration) (net.PacketConn, error) {
	if strategy == nil {
		strategy = d.networkStrategy
	}
	if strategy == nil {
		return d.ListenPacket(ctx, destination)
	}
	if len(interfaceType) == 0 {
		interfaceType = d.networkType
	}
	if len(fallbackInterfaceType) == 0 {
		fallbackInterfaceType = d.fallbackNetworkType
	}
	if fallbackDelay == 0 {
		fallbackDelay = d.networkFallbackDelay
	}
	network := N.NetworkUDP
	if destination.IsIPv4() && !destination.Addr.IsUnspecified() {
		network += "4"
	}
	packetConn, err := d.listenSerialInterfacePacket(ctx, d.udpListener, network, "", *strategy, interfaceType, fallbackInterfaceType, fallbackDelay)
	if err != nil {
		// bind interface failed on legacy xiaomi systems
		if d.defaultNetworkStrategy && errors.Is(err, syscall.EPERM) {
			d.networkStrategy = nil
			return d.ListenPacket(ctx, destination)
		} else {
			return nil, err
		}
	}
	return d.trackPacketConn(ctx, destination, packetConn, nil)
}

func (d *DefaultDialer) UDPListenerControl() (control.Func, bool) {
	egressEnabled := d.autoDetectBindFunc != nil && d.netns == ""
	listenerControl := d.udpListener.Control
	if egressEnabled && d.networkManager.AutoRedirectOutputMark() == 0 {
		listenerControl = control.Append(listenerControl, d.autoDetectBindFunc)
	}
	return listenerControl, egressEnabled
}

func (d *DefaultDialer) trackConn(ctx context.Context, destination M.Socksaddr, conn net.Conn, err error) (net.Conn, error) {
	if err != nil {
		return conn, err
	}
	if d.connectionManager != nil {
		conn = d.connectionManager.TrackConn(conn)
	}
	if d.powerManager != nil {
		recorder := d.powerManager.Recorder()
		if recorder != nil {
			recorder.CountConnectionOpened()
			attribution := d.dialAttribution(ctx, destination)
			conn = bufio.NewCounterConn(conn, []N.CountFunc{func(n int64) {
				recorder.Touch(powerreport.DirectionInbound, int(n), attribution)
			}}, []N.CountFunc{func(n int64) {
				recorder.Touch(powerreport.DirectionOutbound, int(n), attribution)
			}})
		}
	}
	return conn, nil
}

func (d *DefaultDialer) trackPacketConn(ctx context.Context, destination M.Socksaddr, conn net.PacketConn, err error) (net.PacketConn, error) {
	if err != nil {
		return conn, err
	}
	if d.connectionManager != nil {
		conn = d.connectionManager.TrackPacketConn(conn)
	}
	if d.powerManager != nil {
		recorder := d.powerManager.Recorder()
		if recorder != nil {
			recorder.CountConnectionOpened()
			attribution := d.dialAttribution(ctx, destination)
			conn = bufio.NewNetPacketConn(bufio.NewCounterPacketConn(bufio.NewPacketConn(conn), []N.CountFunc{func(n int64) {
				recorder.Touch(powerreport.DirectionInbound, int(n), attribution)
			}}, []N.CountFunc{func(n int64) {
				recorder.Touch(powerreport.DirectionOutbound, int(n), attribution)
			}}))
		}
	}
	return conn, nil
}

func (d *DefaultDialer) dialAttribution(ctx context.Context, destination M.Socksaddr) *powerreport.Attribution {
	attribution := &powerreport.Attribution{}
	dnsTransportTag, hasDNSTransport := adapter.DNSTransportTagFromContext(ctx)
	if hasDNSTransport {
		attribution.DNS = dnsTransportTag
		if d.dnsTransportManager != nil {
			transport, loaded := d.dnsTransportManager.Transport(dnsTransportTag)
			if loaded {
				attribution.DNSType = transport.Type()
			}
		}
	}
	metadata := adapter.ContextFrom(ctx)
	if metadata == nil {
		attribution.Destination = destination.String()
		return attribution
	}
	attribution.Inbound = metadata.Inbound
	attribution.InboundType = metadata.InboundType
	attribution.Network = metadata.Network
	if metadata.Source.IsValid() {
		attribution.Source = metadata.Source.String()
	}
	attribution.Domain = metadata.Domain
	attribution.Protocol = metadata.Protocol
	attribution.User = metadata.User
	if metadata.ProcessInfo != nil {
		attribution.Process = &powerreport.ProcessAttribution{
			ProcessID:    metadata.ProcessInfo.ProcessID,
			UserID:       metadata.ProcessInfo.UserId,
			UserName:     metadata.ProcessInfo.UserName,
			ProcessPath:  metadata.ProcessInfo.ProcessPath,
			PackageNames: metadata.ProcessInfo.AndroidPackageNames,
		}
	}
	attribution.Rule = metadata.RouteRule
	attribution.Outbound = metadata.Outbound
	if d.outboundManager != nil {
		if metadata.Outbound != "" {
			outbound, loaded := d.outboundManager.Outbound(metadata.Outbound)
			if loaded {
				attribution.OutboundType = outbound.Type()
			}
		}
		if metadata.RouteOutbound != "" {
			attribution.Chain = d.outboundChain(metadata.RouteOutbound)
		}
	}
	if metadata.Destination.IsValid() {
		attribution.Destination = metadata.Destination.String()
		if metadata.Destination != destination {
			attribution.Server = destination.String()
		}
	} else {
		attribution.Destination = destination.String()
	}
	return attribution
}

func (d *DefaultDialer) outboundChain(head string) []string {
	var chain []string
	next := head
	for {
		detour, loaded := d.outboundManager.Outbound(next)
		if !loaded {
			break
		}
		chain = append(chain, next)
		outboundGroup, isGroup := detour.(adapter.OutboundGroup)
		if !isGroup {
			break
		}
		next = outboundGroup.Now()
	}
	slices.Reverse(chain)
	return chain
}