require "http"

module Kemal
  # Takes 2 parameters: *path* and a *handler* to specify
  # what action to be done if the route is matched.
  class WebSocket < HTTP::WebSocketHandler
    getter proc

    def initialize(@path : String, &@proc : HTTP::WebSocket, HTTP::Server::Context -> Void)
    end

    def error(code : Int16, message : String)
      err = WebsocketError.new(code, message)
      raise err
    end

    def call(context : HTTP::Server::Context)
      if websocket_upgrade_request?(context.request)
        response = context.response

        version = context.request.headers["Sec-WebSocket-Version"]?
        unless version == ::HTTP::WebSocket::Protocol::VERSION
          response.headers["Sec-WebSocket-Version"] = ::HTTP::WebSocket::Protocol::VERSION
          raise UpgradeRequired.new
        end

        key = context.request.headers["Sec-WebSocket-Key"]?
        raise BadRequest.new("Sec-WebSocket-Key header is missing") unless key

        accept_code = ::HTTP::WebSocket::Protocol.key_challenge(key)

        response.status_code = 101
        response.headers["Upgrade"] = "websocket"
        response.headers["Connection"] = "Upgrade"
        response.headers["Sec-WebSocket-Accept"] = accept_code

        response.upgrade do |io|
          socket = ::HTTP::WebSocket.new(io)
          @proc.call(socket, context)
          socket.run
        rescue error : Exception
          # TODO: check if 500 is what we're supposed to give
          context.response.status_code = 500

          if error.is_a?(WebsocketError)
            code = error.code.to_i32
            message = error.status_message
          else
            code = 1011_i32
            message = "Exception"
          end

          raw = uninitialized UInt8[2]
          IO::ByteFormat::BigEndian.encode(code, raw.to_slice)
          socket.not_nil!.close(
            HTTP::WebSocket::CloseCode.new(code),
            message: String.new(raw.to_slice) + message
          )

          raise error unless error.is_a?(WebsocketError)
        end
      else
        raise UpgradeRequired.new
      end
    end
  end
end