// Copyright 2017 fatedier, fatedier@gmail.com
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package proxy

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"reflect"
	"strconv"
	"sync"
	"time"

	libio "github.com/fatedier/golib/io"
	"golang.org/x/time/rate"

	"github.com/fatedier/frp/pkg/config/types"
	v1 "github.com/fatedier/frp/pkg/config/v1"
	"github.com/fatedier/frp/pkg/msg"
	plugin "github.com/fatedier/frp/pkg/plugin/server"
	"github.com/fatedier/frp/pkg/proto/wire"
	"github.com/fatedier/frp/pkg/util/limit"
	netpkg "github.com/fatedier/frp/pkg/util/net"
	"github.com/fatedier/frp/pkg/util/xlog"
	"github.com/fatedier/frp/server/controller"
	"github.com/fatedier/frp/server/metrics"
)

var proxyFactoryRegistry = map[reflect.Type]func(*BaseProxy) Proxy{}

func RegisterProxyFactory(proxyConfType reflect.Type, factory func(*BaseProxy) Proxy) {
	proxyFactoryRegistry[proxyConfType] = factory
}

type WorkConn struct {
	conn *msg.Conn
}

func NewWorkConn(conn *msg.Conn) *WorkConn {
	return &WorkConn{conn: conn}
}

func (c *WorkConn) Start(m *msg.StartWorkConn) (net.Conn, error) {
	if err := c.conn.WriteMsg(m); err != nil {
		return nil, err
	}
	return c.conn, nil
}

func (c *WorkConn) Close() error {
	return c.conn.Close()
}

type GetWorkConnFn func() (*WorkConn, error)

type Proxy interface {
	Context() context.Context
	Run() (remoteAddr string, err error)
	GetName() string
	GetConfigurer() v1.ProxyConfigurer
	GetWorkConnFromPool(src, dst net.Addr) (workConn net.Conn, err error)
	GetUsedPortsNum() int
	GetResourceController() *controller.ResourceController
	GetUserInfo() plugin.UserInfo
	GetLimiter() *rate.Limiter
	GetLoginMsg() *msg.Login
	Close()
}

type BaseProxy struct {
	name          string
	rc            *controller.ResourceController
	listeners     []net.Listener
	usedPortsNum  int
	poolCount     int
	getWorkConnFn GetWorkConnFn
	serverCfg     *v1.ServerConfig
	encryptionKey []byte
	limiter       *rate.Limiter
	userInfo      plugin.UserInfo
	loginMsg      *msg.Login
	configurer    v1.ProxyConfigurer
	wireProtocol  string

	mu  sync.RWMutex
	xl  *xlog.Logger
	ctx context.Context
}

func (pxy *BaseProxy) GetName() string {
	return pxy.name
}

func (pxy *BaseProxy) Context() context.Context {
	return pxy.ctx
}

func (pxy *BaseProxy) GetUsedPortsNum() int {
	return pxy.usedPortsNum
}

func (pxy *BaseProxy) GetResourceController() *controller.ResourceController {
	return pxy.rc
}

func (pxy *BaseProxy) GetUserInfo() plugin.UserInfo {
	return pxy.userInfo
}

func (pxy *BaseProxy) GetLoginMsg() *msg.Login {
	return pxy.loginMsg
}

func (pxy *BaseProxy) GetLimiter() *rate.Limiter {
	return pxy.limiter
}

func (pxy *BaseProxy) GetConfigurer() v1.ProxyConfigurer {
	return pxy.configurer
}

func (pxy *BaseProxy) Close() {
	xl := xlog.FromContextSafe(pxy.ctx)
	xl.Infof("proxy closing")
	for _, l := range pxy.listeners {
		l.Close()
	}
}

// GetWorkConnFromPool try to get a new work connections from pool
// for quickly response, we immediately send the StartWorkConn message to frpc after take out one from pool
func (pxy *BaseProxy) GetWorkConnFromPool(src, dst net.Addr) (workConn net.Conn, err error) {
	xl := xlog.FromContextSafe(pxy.ctx)
	// try all connections from the pool
	for i := 0; i < pxy.poolCount+1; i++ {
		var pxyWorkConn *WorkConn
		if pxyWorkConn, err = pxy.getWorkConnFn(); err != nil {
			xl.Warnf("failed to get work connection: %v", err)
			return
		}
		xl.Debugf("get a new work connection: [%s]", pxyWorkConn.conn.RemoteAddr().String())
		xl.Spawn().AppendPrefix(pxy.GetName())

		var (
			srcAddr    string
			dstAddr    string
			srcPortStr string
			dstPortStr string
			srcPort    uint64
			dstPort    uint64
		)

		if src != nil {
			srcAddr, srcPortStr, _ = net.SplitHostPort(src.String())
			srcPort, _ = strconv.ParseUint(srcPortStr, 10, 16)
		}
		if dst != nil {
			dstAddr, dstPortStr, _ = net.SplitHostPort(dst.String())
			dstPort, _ = strconv.ParseUint(dstPortStr, 10, 16)
		}
		workConn, err = pxyWorkConn.Start(&msg.StartWorkConn{
			ProxyName: pxy.GetName(),
			SrcAddr:   srcAddr,
			SrcPort:   uint16(srcPort),
			DstAddr:   dstAddr,
			DstPort:   uint16(dstPort),
			Error:     "",
		})
		if err != nil {
			xl.Warnf("failed to send message to work connection from pool: %v, times: %d", err, i)
			pxyWorkConn.Close()
			workConn = nil
		} else {
			workConn = netpkg.NewContextConn(pxy.ctx, workConn)
			break
		}
	}

	if err != nil {
		xl.Errorf("try to get work connection failed in the end")
		return
	}
	return
}

// startVisitorListener sets up a VisitorManager listener for visitor-based proxies (STCP, SUDP).
func (pxy *BaseProxy) startVisitorListener(secretKey string, allowUsers []string, proxyType string) error {
	// if allowUsers is empty, only allow same user from proxy
	if len(allowUsers) == 0 {
		allowUsers = []string{pxy.GetUserInfo().User}
	}
	listener, err := pxy.rc.VisitorManager.Listen(pxy.GetName(), secretKey, allowUsers)
	if err != nil {
		return err
	}
	pxy.listeners = append(pxy.listeners, listener)
	pxy.xl.Infof("%s proxy custom listen success", proxyType)
	pxy.startCommonTCPListenersHandler()
	return nil
}

// buildDomains constructs a list of domains from custom domains and subdomain configuration.
func (pxy *BaseProxy) buildDomains(customDomains []string, subDomain string) []string {
	domains := make([]string, 0, len(customDomains)+1)
	for _, d := range customDomains {
		if d != "" {
			domains = append(domains, d)
		}
	}
	if subDomain != "" {
		domains = append(domains, subDomain+"."+pxy.serverCfg.SubDomainHost)
	}
	return domains
}

// startCommonTCPListenersHandler start a goroutine handler for each listener.
func (pxy *BaseProxy) startCommonTCPListenersHandler() {
	xl := xlog.FromContextSafe(pxy.ctx)
	for _, listener := range pxy.listeners {
		go func(l net.Listener) {
			var tempDelay time.Duration // how long to sleep on accept failure

			for {
				// block
				// if listener is closed, err returned
				c, err := l.Accept()
				if err != nil {
					if err, ok := err.(interface{ Temporary() bool }); ok && err.Temporary() {
						if tempDelay == 0 {
							tempDelay = 5 * time.Millisecond
						} else {
							tempDelay *= 2
						}
						if maxTime := 1 * time.Second; tempDelay > maxTime {
							tempDelay = maxTime
						}
						xl.Infof("met temporary error: %s, sleep for %s ...", err, tempDelay)
						time.Sleep(tempDelay)
						continue
					}

					xl.Warnf("listener is closed: %s", err)
					return
				}
				xl.Infof("get a user connection [%s]", c.RemoteAddr().String())
				go pxy.handleUserTCPConnection(c)
			}
		}(listener)
	}
}

// HandleUserTCPConnection is used for incoming user TCP connections.
func (pxy *BaseProxy) handleUserTCPConnection(userConn net.Conn) {
	xl := xlog.FromContextSafe(pxy.Context())
	defer userConn.Close()

	cfg := pxy.configurer.GetBaseConfig()
	// server plugin hook
	rc := pxy.GetResourceController()
	content := &plugin.NewUserConnContent{
		User:       pxy.GetUserInfo(),
		ProxyName:  pxy.GetName(),
		ProxyType:  cfg.Type,
		RemoteAddr: userConn.RemoteAddr().String(),
	}
	_, err := rc.PluginManager.NewUserConn(content)
	if err != nil {
		xl.Warnf("the user conn [%s] was rejected, err:%v", content.RemoteAddr, err)
		return
	}

	// try all connections from the pool
	workConn, err := pxy.GetWorkConnFromPool(userConn.RemoteAddr(), userConn.LocalAddr())
	if err != nil {
		return
	}
	defer workConn.Close()

	var local io.ReadWriteCloser = workConn
	xl.Tracef("handler user tcp connection, use_encryption: %t, use_compression: %t",
		cfg.Transport.UseEncryption, cfg.Transport.UseCompression)
	if cfg.Transport.UseEncryption {
		local, err = libio.WithEncryption(local, pxy.encryptionKey)
		if err != nil {
			xl.Errorf("create encryption stream error: %v", err)
			return
		}
	}
	if cfg.Transport.UseCompression {
		var recycleFn func()
		local, recycleFn = libio.WithCompressionFromPool(local)
		defer recycleFn()
	}

	if pxy.GetLimiter() != nil {
		local = libio.WrapReadWriteCloser(limit.NewReader(local, pxy.GetLimiter()), limit.NewWriter(local, pxy.GetLimiter()), func() error {
			return local.Close()
		})
	}

	xl.Debugf("join connections, workConn(l[%s] r[%s]) userConn(l[%s] r[%s])", workConn.LocalAddr().String(),
		workConn.RemoteAddr().String(), userConn.LocalAddr().String(), userConn.RemoteAddr().String())

	name := pxy.GetName()
	proxyType := cfg.Type
	var watcher *idleWatcher
	touch := func() {
		if watcher != nil {
			watcher.touch()
		}
	}
	local = wrapCountingReadWriteCloser(local, nil, func(bytes int64) {
		touch()
		metrics.Server.AddTrafficIn(name, proxyType, bytes)
	})
	userConn = netpkg.WrapReadWriteCloserToConn(
		wrapCountingReadWriteCloser(userConn, nil, func(bytes int64) {
			touch()
			metrics.Server.AddTrafficOut(name, proxyType, bytes)
		}),
		userConn,
	)
	if idleTimeout := time.Duration(pxy.serverCfg.Transport.ProxyIdleTimeout) * time.Second; idleTimeout > 0 {
		wrappedLocal, wrappedUserConn := local, userConn
		watcher = startIdleWatcher(idleTimeout, func() {
			xl.Infof("close user connection [%s] of proxy [%s]: no traffic in either direction for %s",
				wrappedUserConn.RemoteAddr().String(), name, idleTimeout)
			_ = wrappedUserConn.Close()
			_ = wrappedLocal.Close()
		})
		defer watcher.stop()
	}
	metrics.Server.OpenConnection(name, proxyType)
	// Traffic is counted incrementally via the counting wrappers above, so the
	// byte totals returned by joinUserConnection are intentionally discarded here.
	_, _, _ = pxy.joinUserConnection(local, userConn, proxyType, xl)
	metrics.Server.CloseConnection(name, proxyType)
	xl.Debugf("join connections closed")
}

func (pxy *BaseProxy) joinUserConnection(local io.ReadWriteCloser, userConn net.Conn, proxyType string, xl *xlog.Logger) (int64, int64, []error) {
	visitorWireProtocol := wireProtocolFromConn(userConn)
	if proxyType == string(v1.ProxyTypeSUDP) && isMixedWireProtocol(pxy.wireProtocol, visitorWireProtocol) {
		xl.Infof("bridge mixed SUDP payload codecs, proxy wireProtocol [%s], visitor wireProtocol [%s]",
			normalizeWireProtocol(pxy.wireProtocol), normalizeWireProtocol(visitorWireProtocol))
		return joinSUDPMessageBridge(local, userConn, pxy.wireProtocol, visitorWireProtocol, xl)
	}
	return libio.Join(local, userConn)
}

type wireProtocolGetter interface {
	WireProtocol() string
}

func wireProtocolFromConn(conn net.Conn) string {
	if getter, ok := conn.(wireProtocolGetter); ok {
		return getter.WireProtocol()
	}
	return ""
}

func isMixedWireProtocol(left, right string) bool {
	return normalizeWireProtocol(left) != normalizeWireProtocol(right)
}

func normalizeWireProtocol(wireProtocol string) string {
	if wireProtocol == wire.ProtocolV2 {
		return wire.ProtocolV2
	}
	return wire.ProtocolV1
}

func joinSUDPMessageBridge(
	proxyConn io.ReadWriteCloser,
	visitorConn io.ReadWriteCloser,
	proxyWireProtocol string,
	visitorWireProtocol string,
	xl *xlog.Logger,
) (inCount int64, outCount int64, errs []error) {
	// The mixed bridge decodes and re-encodes messages, so raw framed byte counts
	// are not available. Count UDP payload bytes and ignore heartbeat traffic.
	proxyRW := msg.NewReadWriter(proxyConn, proxyWireProtocol)
	visitorRW := msg.NewReadWriter(visitorConn, visitorWireProtocol)

	var (
		once       sync.Once
		wait       sync.WaitGroup
		recordErrs = make([]error, 2)
	)
	closeBoth := func() {
		_ = proxyConn.Close()
		_ = visitorConn.Close()
	}

	wait.Add(2)
	go func() {
		defer wait.Done()
		defer once.Do(closeBoth)
		recordErrs[0] = bridgeSUDPProxyToVisitor(proxyRW, visitorRW, &outCount, xl)
	}()
	go func() {
		defer wait.Done()
		defer once.Do(closeBoth)
		recordErrs[1] = bridgeSUDPVisitorToProxy(visitorRW, proxyRW, &inCount, xl)
	}()
	wait.Wait()

	for _, err := range recordErrs {
		if err != nil {
			errs = append(errs, err)
		}
	}
	return
}

func bridgeSUDPProxyToVisitor(from msg.ReadWriter, to msg.ReadWriter, count *int64, xl *xlog.Logger) error {
	for {
		rawMsg, err := from.ReadMsg()
		if err != nil {
			return normalizeSUDPBridgeError(err)
		}

		switch m := rawMsg.(type) {
		case *msg.UDPPacket:
			if err := to.WriteMsg(m); err != nil {
				return normalizeSUDPBridgeError(err)
			}
			*count += int64(len(m.Content))
		case *msg.Ping:
			traceSUDPBridge(xl, "bridge SUDP ping from proxy to visitor")
			if err := to.WriteMsg(m); err != nil {
				return normalizeSUDPBridgeError(err)
			}
		default:
			return fmt.Errorf("unexpected SUDP proxy message %T", rawMsg)
		}
	}
}

func bridgeSUDPVisitorToProxy(from msg.ReadWriter, to msg.ReadWriter, count *int64, xl *xlog.Logger) error {
	for {
		rawMsg, err := from.ReadMsg()
		if err != nil {
			return normalizeSUDPBridgeError(err)
		}

		switch m := rawMsg.(type) {
		case *msg.UDPPacket:
			if err := to.WriteMsg(m); err != nil {
				return normalizeSUDPBridgeError(err)
			}
			*count += int64(len(m.Content))
		case *msg.Ping:
			traceSUDPBridge(xl, "drop SUDP ping from visitor to proxy")
			continue
		default:
			return fmt.Errorf("unexpected SUDP visitor message %T", rawMsg)
		}
	}
}

func normalizeSUDPBridgeError(err error) error {
	if err == nil || errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) {
		return nil
	}
	return err
}

func traceSUDPBridge(xl *xlog.Logger, format string, args ...any) {
	if xl != nil {
		xl.Tracef(format, args...)
	}
}

type Options struct {
	UserInfo           plugin.UserInfo
	LoginMsg           *msg.Login
	PoolCount          int
	ResourceController *controller.ResourceController
	GetWorkConnFn      GetWorkConnFn
	Configurer         v1.ProxyConfigurer
	ServerCfg          *v1.ServerConfig
	EncryptionKey      []byte
	WireProtocol       string
}

func NewProxy(ctx context.Context, options *Options) (pxy Proxy, err error) {
	configurer := options.Configurer
	xl := xlog.FromContextSafe(ctx).Spawn().AppendPrefix(configurer.GetBaseConfig().Name)

	var limiter *rate.Limiter
	limitBytes := configurer.GetBaseConfig().Transport.BandwidthLimit.Bytes()
	if limitBytes > 0 && configurer.GetBaseConfig().Transport.BandwidthLimitMode == types.BandwidthLimitModeServer {
		limiter = rate.NewLimiter(rate.Limit(float64(limitBytes)), int(limitBytes))
	}

	basePxy := BaseProxy{
		name:          configurer.GetBaseConfig().Name,
		rc:            options.ResourceController,
		listeners:     make([]net.Listener, 0),
		poolCount:     options.PoolCount,
		getWorkConnFn: options.GetWorkConnFn,
		serverCfg:     options.ServerCfg,
		encryptionKey: options.EncryptionKey,
		limiter:       limiter,
		xl:            xl,
		ctx:           xlog.NewContext(ctx, xl),
		userInfo:      options.UserInfo,
		loginMsg:      options.LoginMsg,
		configurer:    configurer,
		wireProtocol:  options.WireProtocol,
	}

	factory := proxyFactoryRegistry[reflect.TypeOf(configurer)]
	if factory == nil {
		return pxy, fmt.Errorf("proxy type not support")
	}
	pxy = factory(&basePxy)
	if pxy == nil {
		return nil, fmt.Errorf("proxy not created")
	}
	return pxy, nil
}

type Manager struct {
	// proxies indexed by proxy name
	pxys map[string]Proxy

	mu sync.RWMutex
}

func NewManager() *Manager {
	return &Manager{
		pxys: make(map[string]Proxy),
	}
}

func (pm *Manager) Add(name string, pxy Proxy) error {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	if _, ok := pm.pxys[name]; ok {
		return fmt.Errorf("proxy name [%s] is already in use", name)
	}

	pm.pxys[name] = pxy
	return nil
}

func (pm *Manager) Exist(name string) bool {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	_, ok := pm.pxys[name]
	return ok
}

func (pm *Manager) Del(name string) {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	delete(pm.pxys, name)
}

func (pm *Manager) GetByName(name string) (pxy Proxy, ok bool) {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	pxy, ok = pm.pxys[name]
	return
}