# frozen_string_literal: true

require 'socket'
require 'openssl'
require 'base64'
require 'io/console'
require 'websocket'

module PortainerCli
  class WebsocketExec
    READ_SIZE = 4096

    def initialize(url, auth_header, ssl_verify: true, debug: false)
      @url         = URI.parse(url)
      @auth_header = auth_header
      @ssl_verify  = ssl_verify
      @debug       = debug
    end

    def run
      debug "WS URL: #{@url}"
      debug "Auth headers: #{@auth_header.map { |k, v| "#{k}: #{v[0, 8]}..." }.join(', ')}"

      socket = open_socket

      perform_handshake(socket)
      run_io_loop(socket)
    ensure
      socket&.close rescue nil
    end

    private

    def open_socket
      raw = TCPSocket.new(@url.host, @url.port)

      if @url.scheme == 'wss'
        ctx             = OpenSSL::SSL::SSLContext.new
        ctx.verify_mode = @ssl_verify ? OpenSSL::SSL::VERIFY_PEER : OpenSSL::SSL::VERIFY_NONE
        ssl             = OpenSSL::SSL::SSLSocket.new(raw, ctx)
        ssl.hostname    = @url.host
        ssl.sync_close  = true
        ssl.connect
        ssl
      else
        raw
      end
    end

    def perform_handshake(socket)
      handshake = WebSocket::Handshake::Client.new(
        url:     @url.to_s,
        headers: @auth_header
      )
      debug "Handshake request:\n#{handshake.to_s.chomp}"
      socket.write(handshake.to_s)

      header_bytes = +''
      loop do
        byte = socket.read(1)
        break if byte.nil?
        header_bytes << byte
        handshake << byte
        break if handshake.finished?
      end
      debug "Handshake response headers:\n#{header_bytes.chomp}"

      unless handshake.valid?
        # Try to read the response body for error details
        body = ''
        if (len = header_bytes.match(/Content-Length:\s*(\d+)/i)&.[](1)&.to_i) && len > 0
          body = socket.read(len).to_s
        end
        raise "WebSocket handshake failed: #{handshake.error}\n" \
              "Server response:\n#{header_bytes.chomp}\n#{body}"
      end

      handshake
    end

    def debug(msg)
      return unless @debug
      $stderr.puts "\e[2m[debug] #{msg}\e[0m"
    end

    def run_io_loop(socket)
      incoming = WebSocket::Frame::Incoming::Client.new

      # Put terminal in raw mode so all keypresses go directly to the container
      $stdin.raw do
        loop do
          readable, = IO.select([$stdin, socket], nil, nil, nil)

          readable.each do |io|
            if io == $stdin
              data = io.readpartial(READ_SIZE) rescue nil
              break if data.nil?
              frame = WebSocket::Frame::Outgoing::Client.new(
                data:    data,
                type:    :binary,
                version: 13
              )
              socket.write(frame.to_s)

            elsif io == socket
              data = io.readpartial(READ_SIZE) rescue nil
              if data.nil?
                # Server closed the connection
                return
              end
              incoming << data
              while (frame = incoming.next)
                case frame.type
                when :binary, :text
                  $stdout.write(frame.data)
                  $stdout.flush
                when :close
                  return
                when :ping
                  pong = WebSocket::Frame::Outgoing::Client.new(
                    data:    frame.data,
                    type:    :pong,
                    version: 13
                  )
                  socket.write(pong.to_s)
                end
              end
            end
          end
        end
      end
    end
  end
end