package mysql

import (
	"context"
	"net"
	"os"

	"github.com/sagernet/sing-box/adapter"
	"github.com/sagernet/sing-box/adapter/inbound"
	"github.com/sagernet/sing-box/common/listener"
	"github.com/sagernet/sing-box/common/mux"
	boxTLS "github.com/sagernet/sing-box/common/tls"
	"github.com/sagernet/sing-box/common/uot"
	C "github.com/sagernet/sing-box/constant"
	"github.com/sagernet/sing-box/log"
	"github.com/sagernet/sing-box/option"
	"github.com/sagernet/sing/common"
	E "github.com/sagernet/sing/common/exceptions"
	"github.com/sagernet/sing/common/logger"
	M "github.com/sagernet/sing/common/metadata"
	N "github.com/sagernet/sing/common/network"
	"github.com/sagernet/sing/common/task"
	"github.com/sagernet/smux"

	gmysql "github.com/go-mysql-org/go-mysql/mysql"
	"github.com/go-mysql-org/go-mysql/server"
)

func RegisterInbound(registry *inbound.Registry) {
	inbound.Register[option.MySQLInboundOptions](registry, C.TypeMySQL, NewInbound)
}

var _ adapter.TCPInjectableInbound = (*Inbound)(nil)

type Inbound struct {
	inbound.Adapter
	router           adapter.ConnectionRouterEx
	logger           logger.ContextLogger
	listener         *listener.Listener
	tlsConfig        boxTLS.ServerConfig
	identityProvider *server.InMemoryProvider
	mysqlServer      *server.Server
}

func NewInbound(ctx context.Context, router adapter.Router, logger log.ContextLogger, tag string, options option.MySQLInboundOptions) (adapter.Inbound, error) {
	inbound := &Inbound{
		Adapter:          inbound.NewAdapter(C.TypeMySQL, tag),
		router:           uot.NewRouter(router, logger),
		logger:           logger,
		identityProvider: server.NewInMemoryProvider(),
	}
	for _, user := range options.Users {
		inbound.identityProvider.AddUser(user.User, user.Password)
	}

	if options.TLS == nil || !options.TLS.Enabled {
		return nil, C.ErrTLSRequired
	}

	tlsConfig, err := boxTLS.NewServer(ctx, logger, common.PtrValueOrDefault(options.TLS))
	if err != nil {
		return nil, err
	}
	inbound.tlsConfig = tlsConfig

	// Get the standard *tls.Config from our TLS config for go-mysql server
	stdTLSConfig, err := tlsConfig.STDConfig()
	if err != nil {
		return nil, E.Cause(err, "get std tls config")
	}

	// Create a go-mysql server with our TLS config
	inbound.mysqlServer = server.NewServer(
		"9.7.0",
		gmysql.DEFAULT_COLLATION_ID,
		gmysql.AUTH_NATIVE_PASSWORD,
		nil,
		stdTLSConfig,
	)

	inbound.router, err = mux.NewRouterWithOptions(inbound.router, logger, common.PtrValueOrDefault(options.Multiplex))
	if err != nil {
		return nil, err
	}

	inbound.listener = listener.New(listener.Options{
		Context:           ctx,
		Logger:            logger,
		Network:           []string{N.NetworkTCP},
		Listen:            options.ListenOptions,
		ConnectionHandler: inbound,
	})
	return inbound, nil
}

func (h *Inbound) Start(stage adapter.StartStage) error {
	if stage != adapter.StartStateStart {
		return nil
	}
	if h.tlsConfig != nil {
		err := h.tlsConfig.Start()
		if err != nil {
			return E.Cause(err, "create TLS config")
		}
	}
	return h.listener.Start()
}

func (h *Inbound) Close() error {
	return common.Close(
		h.listener,
		h.tlsConfig,
	)
}

func (h *Inbound) NewConnectionEx(ctx context.Context, conn net.Conn, metadata adapter.InboundContext, onClose N.CloseHandlerFunc) {
	// Use go-mysql server to perform the MySQL handshake (which negotiates TLS)
	mysqlConn, err := h.mysqlServer.NewCustomizedConn(conn, h.identityProvider, &emptyHandler{})
	if err != nil {
		N.CloseOnHandshakeFailure(conn, onClose, err)
		h.logger.ErrorContext(ctx, E.Cause(err, "process connection from ", metadata.Source, ": MySQL handshake"))
		return
	}

	// After MySQL handshake, the underlying connection is TLS-encrypted.
	// Now get the underlying net.Conn (which is a *tls.Conn) and use smux on top of it.
	tlsConn := mysqlConn.Conn.Conn

	h.logger.InfoContext(ctx, "MySQL handshake completed from ", metadata.Source)

	// Handle smux session over the TLS-encrypted connection
	err = h.handleMuxSession(ctx, tlsConn, metadata.Source, onClose, mysqlConn.GetUser())
	if err != nil && !E.IsClosed(err) {
		h.logger.ErrorContext(ctx, E.Cause(err, "process mux session from ", metadata.Source))
	}
}

func (h *Inbound) handleMuxSession(ctx context.Context, conn net.Conn, source M.Socksaddr, onClose N.CloseHandlerFunc, user string) error {
	session, err := smux.Server(conn, smuxConfig())
	if err != nil {
		if onClose != nil {
			onClose(err)
		}
		return err
	}
	var group task.Group
	group.Append0(func(_ context.Context) error {
		for {
			stream, sErr := session.AcceptStream()
			if sErr != nil {
				return sErr
			}
			go h.handleMuxStream(ctx, stream, source, user)
		}
	})
	group.Cleanup(func() {
		session.Close()
		if onClose != nil {
			onClose(os.ErrClosed)
		}
	})
	return group.Run(ctx)
}

func (h *Inbound) handleMuxStream(ctx context.Context, conn net.Conn, source M.Socksaddr, user string) {
	err := h.handleMuxStream0(ctx, conn, source, user)
	if err != nil {
		h.logger.ErrorContext(ctx, E.Cause(err, "process mux stream"))
	}
}

func (h *Inbound) handleMuxStream0(ctx context.Context, conn net.Conn, source M.Socksaddr, user string) error {
	// Read destination from the stream header:
	// 1 byte command (0x01=TCP, 0x03=UDP)
	// then socks address (using SocksaddrSerializer)
	var cmdBuf [1]byte
	_, err := conn.Read(cmdBuf[:])
	if err != nil {
		return E.Cause(err, "read command")
	}
	command := cmdBuf[0]

	destination, err := M.SocksaddrSerializer.ReadAddrPort(conn)
	if err != nil {
		return E.Cause(err, "read destination")
	}

	var metadata adapter.InboundContext
	metadata.Inbound = h.Tag()
	metadata.InboundType = h.Type()
	metadata.Source = source
	metadata.User = user

	switch command {
	case commandTCP:
		metadata.Destination = destination
		h.logger.InfoContext(ctx, "inbound connection to ", metadata.Destination)
		h.router.RouteConnectionEx(ctx, conn, metadata, nil)
	case commandUDP:
		metadata.Destination = destination
		h.logger.InfoContext(ctx, "inbound UoT packet connection to ", metadata.Destination)
		h.router.RouteConnectionEx(ctx, conn, metadata, nil)
	default:
		return E.New("unknown command ", command)
	}
	return nil
}

func smuxConfig() *smux.Config {
	config := smux.DefaultConfig()
	config.KeepAliveDisabled = true
	return config
}

const (
	commandTCP byte = 0x01
	commandUDP byte = 0x03
)

// emptyHandler implements go-mysql server.Handler with no-op operations.
// It is used because we only need the MySQL handshake for TLS negotiation,
// not actual MySQL query handling.
type emptyHandler struct{}

func (h *emptyHandler) UseDB(dbName string) error {
	return nil
}

func (h *emptyHandler) HandleQuery(query string) (*gmysql.Result, error) {
	return nil, gmysql.NewError(gmysql.ER_UNKNOWN_ERROR, "not supported")
}

func (h *emptyHandler) HandleFieldList(table string, fieldWildcard string) ([]*gmysql.Field, error) {
	return nil, gmysql.NewError(gmysql.ER_UNKNOWN_ERROR, "not supported")
}

func (h *emptyHandler) HandleStmtPrepare(query string) (int, int, interface{}, error) {
	return 0, 0, nil, gmysql.NewError(gmysql.ER_UNKNOWN_ERROR, "not supported")
}

func (h *emptyHandler) HandleStmtExecute(context interface{}, query string, args []interface{}) (*gmysql.Result, error) {
	return nil, gmysql.NewError(gmysql.ER_UNKNOWN_ERROR, "not supported")
}

func (h *emptyHandler) HandleStmtClose(context interface{}) error {
	return nil
}

func (h *emptyHandler) HandleOtherCommand(cmd byte, data []byte) error {
	return gmysql.NewError(gmysql.ER_UNKNOWN_ERROR, "not supported")
}

// compile-time check
var _ server.Handler = (*emptyHandler)(nil)