package ssh

import (
	"net"

	"golang.org/x/crypto/ssh"

	"git.maze.io/maze/conduit/auth"
	"git.maze.io/maze/conduit/internal/netutil"
	"git.maze.io/maze/conduit/logger"
)

type Server struct {
	// ConnectHandler gets called before accepting new connections.
	ConnectHandler []ConnectHandler

	// AcceptHandler accepts new SSH connections, it takes care of authenticating, etc.
	AcceptHandler AcceptHandler

	// Handler per channel type.
	ChannelHandler map[string]ChannelHandler

	// PortForwardHandler
	PortForwardHandler PortForwardRequestHandler

	// FIPSMode enables FIPS 140-2 compatible ciphers, key exchanges, etc.
	FIPSMode bool // TODO(maze): implement

	// Logger for our server.
	Logger logger.Structured

	// serverConfig is our SSH server configuration.
	serverConfig *ssh.ServerConfig
}

func NewServer(keys []ssh.Signer) *Server {
	config := new(ssh.ServerConfig)
	config.SetDefaults()
	config.ServerVersion = "SSH-2.0-conduit"

	for _, key := range keys {
		config.AddHostKey(key)
	}

	return &Server{
		ChannelHandler: make(map[string]ChannelHandler),
		Logger:         logger.StandardLog,
		serverConfig:   config,
	}
}

func (s *Server) Serve(l net.Listener) error {
	for {
		c, err := l.Accept()
		if err != nil {
			return err
		}
		go s.handle(c)
	}
}

func (s *Server) handle(c net.Conn) {
	log := s.Logger.Values(logger.Values{
		"client": c.RemoteAddr().String(),
		"server": c.LocalAddr().String(),
	})
	log.Debug("New client connection")

	defer func() {
		if err := c.Close(); err != nil && !netutil.IsClosing(err) {
			log = log.Err(err)
		}
		log.Debug("Closing client connection")
	}()

	for _, h := range s.ConnectHandler {
		n, err := h.HandleConnect(c)
		if err != nil {
			log.Err(err).Warn("Error from connect handler, closing client connection")
			return
		} else if n != nil {
			log.Debugf("Replacing client connection with %T", n)
			c = n
		}
	}

	// Configure our SSH server.
	handler := s.AcceptHandler
	if handler == nil {
		log.Warn("No accept handler configured, using NO AUTHENTICATION")
		handler = auth.None{}
	}

	// We made it, now let's talk some SSH.
	sshConn, channels, requests, err := handler.HandleAccept(c, s.serverConfig)
	if err != nil {
		log.Err(err).Warn("Error establishing SSH session with client")
		return
	}
	go ssh.DiscardRequests(requests)

	log = log.Value("user", sshConn.User())
	ctx := newSSHContext(s, c, sshConn, log)
	log = log.Value("context", ctx.ID())
	log.Value("version", string(sshConn.ClientVersion())).Info("New SSH client")

	if err = ctx.handleChannels(channels); err != nil {
		if netutil.IsClosing(err) {
			log.Err(err).Debug("Client handler terminated")
		} else {
			log.Err(err).Warn("Error handling channel requests")
		}
		return
	}
}