module EventMachine
  module ProxyServer
    class Connection < EventMachine::Connection
      attr_accessor :debug
      
      ##### Proxy Methods
      def on_data(&blk); @on_data = blk; end
      def on_response(&blk); @on_response = blk; end
      def on_finish(&blk); @on_finish = blk; end
      def on_connect(&blk); blk.call; end

      ##### EventMachine
      def initialize(options)
        @debug = options[:debug] || false
        @servers = {}
      end

      def receive_data(data)
        debug [:connection, data]
        processed = @on_data.call(data)

        return if processed == :async or processed.nil?
        relay_to_servers(processed)
      end

      def relay_to_servers(processed)
        if processed.is_a? Array
          data, servers = *processed

          # guard for "unbound" servers
          servers = servers.collect {|s| @servers[s]}.compact
        else
          data = processed
          servers ||= @servers.values.compact
        end

        servers.each do |s|
          s.send_data data unless data.nil?
        end
      end
      
      #
      # initialize connections to backend servers
      #
      def server(name, opts)
        srv = EventMachine::connect(opts[:host], opts[:port], EventMachine::ProxyServer::Backend, @debug) do |c|
          c.name = name
          c.plexer = self
          c.proxy_incoming_to(self, 10240) if opts[:relay_server]
        end
        self.proxy_incoming_to(srv, 10240) if opts[:relay_client]

        @servers[name] = srv
      end

      #
      # [ip, port] of the connected client
      #
      def peer
        @peer ||= Socket.unpack_sockaddr_in(get_peername).reverse
      end

      #
      # relay data from backend server to client
      #
      def relay_from_backend(name, data)
        debug [:relay_from_backend, name, data]

        data = @on_response.call(name, data)
        send_data data unless data.nil?
      end

      def unbind
        debug [:unbind, :connection]
      
        # terminate any unfinished connections
        @servers.values.compact.each do |s|
          s.close_connection
        end
      end

      def unbind_backend(name)
        debug [:unbind_backend, name]
        @servers[name] = nil

        # if all connections are terminated downstream, then notify client
        close_connection_after_writing if @servers.values.compact.size.zero?
        
        if @on_finish
          @on_finish.call(name)

          # not sure if this is required
          # @on_finish.call(:done) if @servers.values.compact.size.zero?
        end
      end

      private
      
      def debug(*data)
        if @debug
          require 'pp'
          pp data
          puts
        end
      end
    end
  end
end