package shadowsocks_2022

import (
	"context"
	"crypto/cipher"
	"encoding/binary"
	"strconv"
	"time"

	"github.com/xtls/xray-core/common"
	"github.com/xtls/xray-core/common/buf"
	"github.com/xtls/xray-core/common/errors"
	"github.com/xtls/xray-core/common/log"
	"github.com/xtls/xray-core/common/net"
	"github.com/xtls/xray-core/common/protocol"
	"github.com/xtls/xray-core/common/session"
	"github.com/xtls/xray-core/common/uuid"
	"github.com/xtls/xray-core/core"
	"github.com/xtls/xray-core/features/policy"
	"github.com/xtls/xray-core/features/routing"
	"github.com/xtls/xray-core/transport/internet/stat"
)

func init() {
	common.Must(common.RegisterConfig((*RelayServerConfig)(nil), func(ctx context.Context, config interface{}) (interface{}, error) {
		return NewRelayServer(ctx, config.(*RelayServerConfig))
	}))
}

type relayDest struct {
	destination net.Destination
	email       string
	level       uint32
	blockCipher cipher.Block
}

type RelayInbound struct {
	networks      []net.Network
	method        *CipherMethod
	relayPSK      []byte
	relayBlock    cipher.Block
	destinations  map[[AESBlockSize]byte]*relayDest
	udpSessions   *UDPSessionManager
	policyManager policy.Manager
}

func NewRelayServer(ctx context.Context, config *RelayServerConfig) (*RelayInbound, error) {
	networks := config.Network
	if len(networks) == 0 {
		networks = []net.Network{
			net.Network_TCP,
			net.Network_UDP,
		}
	}

	method, err := GetCipherMethod(config.Method)
	if err != nil {
		return nil, err
	}
	if method.IsChaCha {
		return nil, errors.New("shadowsocks 2022 relay: only aes methods are supported")
	}

	relayPSK, err := ParseKey(config.Key, method.KeySaltLength)
	if err != nil {
		return nil, err
	}

	relayBlock, err := method.NewBlock(relayPSK)
	if err != nil {
		return nil, err
	}

	v := core.MustFromContext(ctx)
	i := &RelayInbound{
		networks:      networks,
		method:        method,
		relayPSK:      relayPSK,
		relayBlock:    relayBlock,
		destinations:  make(map[[AESBlockSize]byte]*relayDest),
		udpSessions:   NewUDPSessionManager(500 * time.Second),
		policyManager: v.GetFeature(policy.ManagerType()).(policy.Manager),
	}

	for idx, d := range config.Destinations {
		if d.Email == "" {
			u := uuid.New()
			d.Email = "unnamed-destination-" + strconv.Itoa(idx) + "-" + u.String()
		}
		destKey, err := ParseKey(d.Key, method.KeySaltLength)
		if err != nil {
			return nil, err
		}

		destBlock, err := method.NewBlock(destKey)
		if err != nil {
			return nil, err
		}

		hash := DeriveUserPSKHash(destKey)

		i.destinations[hash] = &relayDest{
			destination: net.TCPDestination(d.Address.AsAddress(), net.Port(d.Port)),
			email:       d.Email,
			level:       uint32(d.Level),
			blockCipher: destBlock,
		}
	}

	return i, nil
}

func (i *RelayInbound) Network() []net.Network {
	return i.networks
}

func (i *RelayInbound) Process(ctx context.Context, network net.Network, connection stat.Connection, dispatcher routing.Dispatcher) error {
	inbound := session.InboundFromContext(ctx)
	inbound.Name = "shadowsocks-2022-relay"
	inbound.CanSpliceCopy = 3

	if network == net.Network_TCP {
		return i.processTCP(ctx, connection, dispatcher)
	}
	return i.processUDP(ctx, connection, dispatcher)
}

func (i *RelayInbound) processTCP(ctx context.Context, conn net.Conn, dispatcher routing.Dispatcher) error {
	defer conn.Close()

	sessionPolicy := i.policyManager.ForLevel(0)
	if err := conn.SetReadDeadline(time.Now().Add(sessionPolicy.Timeouts.Handshake)); err != nil {
		return errors.New("unable to set read deadline").Base(err)
	}

	// Read initial handshake in a single read call per SIP022 §3.1.3 & §3.1.4
	needed := i.method.KeySaltLength + AESBlockSize
	requestHeader := buf.New()
	n, err := requestHeader.ReadFrom(conn)
	if err != nil {
		requestHeader.Release()
		ResetTCPConn(conn)
		return err
	}
	if int(n) < needed {
		requestHeader.Release()
		ResetTCPConn(conn)
		return ErrInvalidRequest
	}

	headerSlice := requestHeader.Bytes()
	salt := headerSlice[:i.method.KeySaltLength]
	eih := headerSlice[i.method.KeySaltLength:needed]

	decryptedHash, err := DecryptEIH(i.method, i.relayPSK, salt, eih)
	if err != nil {
		requestHeader.Release()
		ResetTCPConn(conn)
		return err
	}

	targetDest, ok := i.destinations[decryptedHash]
	if !ok {
		requestHeader.Release()
		ResetTCPConn(conn)
		return ErrInvalidRequest
	}
	conn.SetReadDeadline(time.Time{})

	inbound := session.InboundFromContext(ctx)
	inbound.User = &protocol.MemoryUser{
		Email: targetDest.email,
		Level: targetDest.level,
	}

	ctx = log.ContextWithAccessMessage(ctx, &log.AccessMessage{
		From:   conn.RemoteAddr(),
		To:     targetDest.destination,
		Status: log.AccessAccepted,
		Email:  targetDest.email,
	})

	errors.LogInfo(ctx, "relaying connection to ", targetDest.destination)

	link, err := dispatcher.Dispatch(ctx, targetDest.destination)
	if err != nil {
		requestHeader.Release()
		return err
	}

	// Unwrap outer EIH: send client salt and remaining handshake bytes to next hop
	// in a single write call, satisfying downstream server's single-read handshake expectation (SIP022 §3.1.3).
	var saltCopy [32]byte
	copy(saltCopy[:i.method.KeySaltLength], salt)
	copy(requestHeader.Bytes()[AESBlockSize:AESBlockSize+i.method.KeySaltLength], saltCopy[:i.method.KeySaltLength])
	requestHeader.Advance(AESBlockSize)

	if err := link.Writer.WriteMultiBuffer(buf.MultiBuffer{requestHeader}); err != nil {
		return err
	}

	return TransportTCP(ctx, i.policyManager.ForLevel(targetDest.level), buf.NewReader(conn), buf.NewWriter(conn), link)
}

func (i *RelayInbound) processUDP(ctx context.Context, conn stat.Connection, dispatcher routing.Dispatcher) error {
	reader := buf.NewPacketReader(conn)
	for {
		mb, err := reader.ReadMultiBuffer()
		if err != nil {
			buf.ReleaseMulti(mb)
			return err
		}

		for _, b := range mb {
			data := b.Bytes()
			if len(data) < 2*AESBlockSize {
				b.Release()
				continue
			}

			var packetHeader [AESBlockSize]byte
			i.relayBlock.Decrypt(packetHeader[:], data[:AESBlockSize])

			eiHeader := DecryptUDPEIH(i.relayBlock, packetHeader[:], data[AESBlockSize:2*AESBlockSize])

			targetDest, ok := i.destinations[eiHeader]
			if !ok {
				b.Release()
				continue
			}

			// Extract sessionID from raw packetHeader for session-level link caching before re-encrypting
			sessionID := binary.BigEndian.Uint64(packetHeader[:8])

			// Re-encrypt packetHeader with next hop block cipher
			targetDest.blockCipher.Encrypt(packetHeader[:], packetHeader[:])

			// Strip outer EIH: replace second block with re-encrypted packetHeader and advance
			copy(data[AESBlockSize:2*AESBlockSize], packetHeader[:])
			b.Advance(int32(AESBlockSize))

			dest := targetDest.destination
			dest.Network = net.Network_UDP

			sessionItem := i.udpSessions.GetOrCreate(sessionID)
			if sessionItem.User == nil {
				sessionItem.Lock()
				if sessionItem.User == nil {
					sessionItem.User = &protocol.MemoryUser{
						Email: targetDest.email,
						Level: targetDest.level,
					}
				}
				sessionItem.Unlock()
			}
			link, err := sessionItem.EnsureLink(ctx, conn, dest, dispatcher, i.policyManager, nil)
			if err != nil {
				b.Release()
				continue
			}

			_ = link.Writer.WriteMultiBuffer(buf.MultiBuffer{b})
		}
	}
}