require 'socket'

module DRuby

    # A DRuby session sends messages back and forth between server and client
    # The message format is as follows:
    # +-----+-----+-----+-----+-----+---  +  ---+-----+
    # | msg. size | msg. type | marshalled msg/args   |
    # +-----+-----+-----+-----+-----+---  +  ---+-----+
    class Session
        REQUEST         = 0x1001
        REQUEST_BLOCK   = 0x1002
        RETVAL          = 0x2001
        EXCEPTION       = 0x2002
        YIELD           = 0x2003

        def initialize(io)
            @io = io
        end

        def send_message(type, obj)
            data = Marshal.dump(obj)
            header = [data.length, type]
            @io.write(header.pack("vv"))
            @io.write(data)
        end

        def get_message()
            header = @io.read(4)
            size, type = header.unpack("vv")
            data = @io.read(size)
            obj = Marshal.load(data)
            return [ type, obj ]
        end

        def finished()
            return @io.eof
        end
    end

    # The DRuby server class.  Like its drb equivalent, this class spawns off
    # a new thread which processes requests, allowing the server to do other
    # things while it is doing processing for a distributed object.  This
    # means, though, that all objects used with DRuby must be thread-safe.
    class Server
        attr_reader :host, :port, :obj, :thread

        def initialize(host, port, obj)
            @host = host
            @port = port
            @obj = obj
            @thread = Thread.new do
                server = TCPServer.new(@host, @port)
                while(socket = server.accept)
                    socket.setsockopt(Socket::SOL_TCP, Socket::TCP_NODELAY, 1)
                    Thread.new(socket) do |socket|
                        begin
                            session_loop(Session.new(socket))
                        rescue Exception
                            # Let the user know we got an exception and exit
                            # the loop
                            puts $!
                            puts $!.backtrace
                        end
                    end
                end
            end
        end

        # Main server loop.  Wait for a REQUEST message, process it, and send
        # back a YIELD, EXCEPTION, or RETVAL message.
        def session_loop(session)
            while not session.finished()
                type, message = session.get_message()
                begin
                    case type
                        when Session::REQUEST
                            retval = @obj.__send__(*message)
                        when Session::REQUEST_BLOCK
                            retval = @obj.__send__(*message) do |*i|
                                session.send_message(Session::YIELD, i)
                            end
                        else
                            # This should not happen
                            raise ArgumentError
                    end
                    session.send_message(Session::RETVAL, retval)
                rescue Exception
                    session.send_message(Session::EXCEPTION, $!)
                end
            end
        end
    end

    # The DRuby client class.  A DRuby server must be started on the given
    # host and port before instantiating a DRuby client.
    class Client
        attr_reader :host, :port
        
        def initialize(host, port)
            @host = host
            @port = port
            @server = TCPSocket.open(@host, @port)
            @server.setsockopt(Socket::SOL_TCP, Socket::TCP_NODELAY, 1)
            @session = Session.new(@server)
        end

        # Client request handler.  The idea here is to send out a REQUEST
        # message, and get back a YIELD, EXCEPTION, or RETVAL message.
        def method_missing(method, *args)
            message = [ method, *args ]
            @session.send_message(
                (block_given?) ? Session::REQUEST_BLOCK : Session::REQUEST,
                message)
            loop do
                type, message = @session.get_message()
                case type
                    when Session::RETVAL     ; return message
                    when Session::YIELD      ; yield message
                    when Session::EXCEPTION  ; raise message
                    else                     ; raise RuntimeError
                end
            end
        end
    end

end