package sshrpc

import (
	"fmt"
	"net/rpc"

	"dev.justinjudd.org/justin/easyssh"

	"golang.org/x/crypto/ssh"
)

// DefaultRPCChannel is the Channel Name that will be used to carry the RPC traffic, can be changed in Serverr and Client
const DefaultRPCChannel = "RPCChannel"

// RPCSubsystem is the subsystem that will be used to trigger RPC endpoint creation
const RPCSubsystem = "RPCSubsystem"

// CallbackFunc to be called when reverse RPC client is created
type CallbackFunc func(RPCClient *rpc.Client, conn ssh.Conn)

// Server represents an SSH Server that spins up RPC servers when requested.
type Server struct {
	*rpc.Server
	Config      *ssh.ServerConfig
	ChannelName string

	CallbackFunc
}

// NewServer returns a new Server to handle incoming SSH and RPC requests.
func NewServer() *Server {
	c := &ssh.ServerConfig{
		PasswordCallback: func(c ssh.ConnMetadata, pass []byte) (*ssh.Permissions, error) {
			if c.User() == "sshrpc" && string(pass) == "sshrpc" {
				return nil, nil
			}
			return nil, fmt.Errorf("password rejected for %q", c.User())
		},
	}
	return &Server{Server: rpc.NewServer(), Config: c, ChannelName: DefaultRPCChannel}

}

// StartServer starts the server listening for requests
func (s *Server) StartServer(address string) {

	server := easyssh.Server{Addr: address, Config: s.Config}
	handler := easyssh.NewSSHConnHandler()
	channelsMux := easyssh.NewChannelsMux()
	channelsMux.HandleChannel(s.ChannelName, s)
	handler.MultipleChannelsHandler = channelsMux
	server.Handler = handler
	server.ListenAndServe()

}

// HandleChannel implements easyssh.HandleChannelFunc allowing Server to Listen for and respond to all requests for a egistered SSH channel
func (s *Server) HandleChannel(newChannel ssh.NewChannel, channel ssh.Channel, reqs <-chan *ssh.Request, sshConn ssh.Conn) {
	go func(in <-chan *ssh.Request) {
		for req := range in {
			ok := false
			switch req.Type {

			case easyssh.SubsystemRequest:
				ok = true
				logger.Printf("subsystem '%s'", req.Payload[4:])
				switch string(req.Payload[4:]) {
				//RPCSubsystem Request made indicates client desires RPC Server access
				case RPCSubsystem:
					go s.ServeConn(channel)
					logger.Printf("Started SSH RPC")
					// triggers reverse RPC connection as well
					clientChannel, err := openRPCClientChannel(sshConn, s.ChannelName+"-reverse")
					if err != nil {
						logger.Printf("Failed to create client channel: " + err.Error())
						continue
					}
					rpcClient := rpc.NewClient(clientChannel)
					if s.CallbackFunc != nil {
						s.CallbackFunc(rpcClient, sshConn)
					}
					logger.Printf("Started SSH RPC client")
				default:
					logger.Printf("Unknown subsystem: %s", req.Payload)
				}

			}
			if !ok {
				logger.Printf("declining %s request...", req.Type)
			}
			req.Reply(ok, nil)
		}
	}(reqs)
}