require 'socket'
require 'openssl'
require 'puma/server'
require 'ds9'

class Server < DS9::Server
  def initialize socket, app
    @app           = app
    @read_streams  = {}
    @write_streams = {}
    @socket        = socket
    super()
  end

  def send_event string
    @socket.write_nonblock string
  end

  def recv_event length
    case data = @socket.read_nonblock(length, nil, exception: false)
    when :wait_readable then DS9::ERR_WOULDBLOCK
    when nil            then DS9::ERR_EOF
    else
      data
    end
  end

  def on_begin_headers frame
    @read_streams[frame.stream_id] = []
  end

  def on_data_source_read stream_id, length
    @write_streams[stream_id].body.read length
  end

  def on_stream_close id, error_code
    @read_streams.delete id
    @write_streams.delete id
  end

  def submit_push_promise stream_id, headers, block
    response = Response.new(self, super(stream_id, headers), [])
    @app.call Hash[headers], response
    @write_streams[response.stream_id] = response
  end

  def on_header name, value, frame, flags
    @read_streams[frame.stream_id] << [name, value]
  end

  class Response < Struct.new :stream, :stream_id, :body
    def push headers, &block
      stream.submit_push_promise stream_id, headers, block
    end

    def submit_response headers
      stream.submit_response stream_id, headers
    end

    def finish str
      self.body = StringIO.new str
    end
  end

  def on_frame_recv frame
    return unless frame.headers?
    req_headers = @read_streams[frame.stream_id]
    response = Response.new(self, frame.stream_id, [])
    @app.call Hash[req_headers], response
    @write_streams[frame.stream_id] = response
  end

  def run
    while want_read? || want_write?
      if want_read?
        rd, _, _ = IO.select([@socket])
        return if @socket.eof?

        receive
      end

      if want_write?
        _, wr, _ = IO.select(nil, [@socket])
        send
      end
    end
  end

  def self.connect_ssl sock, ctx
    ssl_sock = OpenSSL::SSL::SSLSocket.new sock, ctx
    ssl_sock.accept
    ssl_sock
  end
end

CERT = OpenSSL::X509::Certificate.new File.read ARGV[0]
KEY = OpenSSL::PKey::RSA.new File.read ARGV[1]
PKEY = OpenSSL::PKey::EC.new "prime256v1"

class Context
  STR = "This server only supports HTTP2 requests\n"

  def initialize host, port
    @ctx               = OpenSSL::SSL::SSLContext.new
    @ctx.npn_protocols = [DS9::PROTO_VERSION_ID]

    # This needs https://bugs.ruby-lang.org/issues/11356
    @ctx.tmp_ecdh_callback = ->(ssl, export, len) { PKEY }

    @ctx.cert          = CERT
    @ctx.key           = KEY
    @authority         = ['localhost', port.to_s].join ':'
  end

  def call _, sock
    ssl_sock = Server.connect_ssl sock, @ctx

    if ssl_sock.npn_protocol == DS9::PROTO_VERSION_ID
      app = ->(headers, response) {
        puts headers[":path"]

        case headers[":path"]
        when "/favicon.ico"
          response.submit_response [[':status', '200'],
                               ["server", 'test server'],
                               ["date", 'Sat, 27 Jun 2015 17:29:21 GMT']]
          puts "PUSHING FAVICON.PNG"
          response.finish File.binread "favicon.ico"

        when "/test.png"
          response.submit_response [[':status', '200'],
                               ["server", 'test server'],
                               ["date", 'Sat, 27 Jun 2015 17:29:21 GMT']]
          puts "PUSHING TEST.PNG"
          response.finish File.binread "test.png"

        when "/"
          response.push [[":method", "GET"],
                         [":path", "/favicon.ico"],
                         [":scheme", "https"],
                         [":authority", @authority]]

          response.push [[":method", "GET"],
                         [":path", "/test.png"],
                         [":scheme", "https"],
                         [":authority", @authority]]

          response.submit_response [[':status', '200'],
                                    ["server", 'test server'],
                                    ["content-type", 'text/html'],
                                    ["date", 'Sat, 27 Jun 2015 17:29:21 GMT']]
          response.finish "<html><body><img src='/test.png' /></body></html>"
        else
          response.submit_response [[':status', '404'],
                                    ["server", 'test server'],
                                    ["content-type", 'text/plain'],
                                    ["date", 'Sat, 27 Jun 2015 17:29:21 GMT']]
          response.finish "Not Found"
        end
      }
      session = Server.new ssl_sock, app
      puts "OPENED"
      session.submit_settings [[DS9::Settings::MAX_CONCURRENT_STREAMS, 100]]
      session.run
      ssl_sock.close
      puts "CLOSED"
    else
      ssl_sock.write "HTTP/1.1 505 HTTP Version Not Supported\r\n"
      ssl_sock.write "Content-Type: text/plain\r\n"
      ssl_sock.write "Content-Length: #{STR.bytesize}\r\n"
      ssl_sock.write "Connection: close\r\n"
      ssl_sock.write "\r\n"
      ssl_sock.write STR
      ssl_sock.close
    end
  end
end

PORT = 8080
HOST = "localhost"

server = Puma::Server.new Context.new(HOST, PORT)
server.add_tcp_listener HOST, PORT
server.tcp_mode!
server.run
server.thread.join