require 'uri'
require 'socket'
require 'openssl'
require 'timeout'
require 'websocket'

module Hypertube
  module Core
    module WebSocketClient
      class WebSocketClient
        MAX_ATTEMPTS = 2
        DEFAULT_TIMEOUT_SECONDS = 60
        TRANSIENT_ERRORS = [
          Errno::EPIPE,
          Errno::ECONNRESET,
          Errno::ETIMEDOUT,
          IOError,
          SystemCallError,
          Timeout::Error,
          OpenSSL::SSL::SSLError,
          WebSocket::Error
        ].freeze

        # Cache: URI(string) -> Hypertube::Core::WebSocketClient::WebSocketClient
        @clients = {}
        @clients_lock = Mutex.new

        class << self
          # Equivalent to C# SendMessage(Uri, byte[])
          # message can be String (bytes) or Array<Integer 0..255>
          def send_message(connection_data, message)
            raise ArgumentError, "message cannot be nil" if message.nil?
            key = connection_key(connection_data)

            client = create_or_get_client(connection_data, expected_current: nil)

            last_ex = nil

            MAX_ATTEMPTS.times do |attempt|
              begin
                # Serialize send+receive per connection instance (like _sendReceiveLock)
                client.with_send_receive_lock do
                  client.send_bytes(message)
                  return client.receive_bytes
                end
              rescue *TRANSIENT_ERRORS => ex
                last_ex = ex
                client.mark_closed
                break if attempt == MAX_ATTEMPTS - 1

                # Recreate semantics: only replace if the cached instance equals the one that failed
                client = create_or_get_client(connection_data, expected_current: client)
              end
            end

            if last_ex
              e = RuntimeError.new("WebSocket send to #{key} failed after #{MAX_ATTEMPTS} attempts. Last error: #{last_ex.message}")
              e.set_backtrace(last_ex.backtrace)
              raise e
            end
            raise RuntimeError, "WebSocket send to #{key} failed after #{MAX_ATTEMPTS} attempts."
          end

          # Equivalent to C# Close(uri)
          # suppress_close_frame can emulate "do not send CLOSE" during reconnect handling
          def close(connection_data, suppress_close_frame: false)
            key = connection_key(connection_data)
            @clients_lock.synchronize do
              if (client = @clients[key])
                begin
                  client.close(suppress_close_frame: suppress_close_frame)
                ensure
                  @clients.delete(key)
                end
              end
            end
          end

          # Equivalent to C# GetState(uri, out state)
          # Returns :OPEN, :CLOSED, or nil if no cached client
          def get_state(connection_data)
            key = connection_key(connection_data)
            @clients_lock.synchronize do
              c = @clients[key]
              return nil unless c
              c.open? ? :OPEN : :CLOSED
            end
          end

          private

          def connection_key(connection_data)
            if connection_data.respond_to?(:to_canonical_key)
              return connection_data.to_canonical_key
            end
            url = connection_data.to_s
            "websocket|#{url}"
          end

          # Mirrors the C# consolidated CreateOrGetClient(uri, expectedCurrent)
          def create_or_get_client(connection_data, expected_current:)
            @clients_lock.synchronize do
              key = connection_key(connection_data)
              if expected_current.nil?
                # add-or-get
                if (current = @clients[key])
                  if current.open?
                    return current
                  else
                    # Replace stale/broken
                    begin current.close(suppress_close_frame: true) rescue nil end
                    new_client = new(connection_data)
                    @clients[key] = new_client
                    return new_client
                  end
                end

                created = new(connection_data)
                @clients[key] = created
                created
              else
                # Recreate semantics: reuse a fresh client if another thread already replaced the failed one.
                if (current = @clients[key])
                  return current if !current.equal?(expected_current) && current.open?

                  begin current.close(suppress_close_frame: true) rescue nil end
                  @clients.delete(key)
                end

                new_client = new(connection_data)
                @clients[key] = new_client
                new_client
              end
            end
          end
        end

        def initialize(connection_data)
          @connection_data = connection_data
          @closed = false
          url = connection_data.respond_to?(:hostname) ? connection_data.hostname : connection_data.to_s
          @connect_timeout = connection_data.respond_to?(:connect_timeout) ? connection_data.connect_timeout : nil
          @uri = URI.parse(url)
          raise "WebSocket URL must use ws:// or wss://. Got: #{url}" unless %w[ws wss].include?(@uri.scheme)

          @host = @uri.host
          @port = @uri.port || default_port
          @path = (@uri.path.nil? || @uri.path.empty?) ? '/' : @uri.path
          @path += "?#{@uri.query}" if @uri.query

          @send_receive_lock = Mutex.new

          begin
            @socket = open_socket
            @handshake = WebSocket::Handshake::Client.new(url: url, headers: request_headers)
            perform_handshake
          rescue Errno::ECONNREFUSED, Errno::ETIMEDOUT, Errno::EHOSTUNREACH, SocketError, IOError, SystemCallError,
                 OpenSSL::SSL::SSLError => e
            mark_closed
            raise "WebSocket connection to #{url} failed: #{e.message}. Check that the server is running, the URL is correct (ws:// or wss://), and that no firewall or proxy is blocking the connection."
          rescue WebSocket::Error => e
            mark_closed
            raise "WebSocket connection to #{url} failed: #{e.message}. Check that the server is running and supports WebSocket."
          end

          @incoming = WebSocket::Frame::Incoming::Client.new(version: @handshake.version)
        end

        # Serialize send+receive for this connection instance
        def with_send_receive_lock
          @send_receive_lock.synchronize { yield }
        end

        def mark_closed
          @closed = true
        end

        # Send bytes (String bytes or Array of bytes)
        def send_bytes(message)
          data =
            case message
            when String then message.b
            when Array  then message.pack('C*')
            else raise ArgumentError, "Unsupported message type: #{message.class}"
            end

          frame = WebSocket::Frame::Outgoing::Client.new(
            version: @handshake.version,
            data: data,
            type: :binary
          )
          write_all(frame.to_s)
          true
        end

        # Receive a full message payload as Array<Integer> (bytes)
        # Raises on CLOSE like the C# code (instead of returning nil).
        def receive_bytes(timeout: DEFAULT_TIMEOUT_SECONDS)
          deadline = Time.now + timeout
          payload_chunks = []

          loop do
            if (frame = @incoming.next)
              case frame.type
              when :binary, :text, :continuation
                payload_chunks << frame.data
                # Note: if you want to support fragmented continuation properly, you’d track FIN;
                # the websocket gem’s Incoming client typically handles reassembly at this layer.
                return payload_chunks.join.bytes
              when :ping
                send_pong(frame.data)
              when :close
                # Mirror C#: close and throw
                send_close_reply
                raise WebSocket::Error, "WebSocket connection to #{@uri} was closed by the remote endpoint. The server may have shut down or closed the connection."
              end
              next
            end

            remaining = deadline - Time.now
            raise Timeout::Error, "WebSocket receive from #{@uri} timed out. The server did not respond in time." if remaining <= 0

            if IO.select([@socket], nil, nil, [remaining, 0.05].max)
              chunk = @socket.read_nonblock(64 * 1024, exception: false)
              case chunk
              when nil
                raise EOFError, "peer closed"
              when :wait_readable
                next
              else
                @incoming << chunk
              end
            end
          end
        end

        def open?
          return false if @closed
          return false if @socket.closed?
          true
        rescue IOError, SystemCallError, OpenSSL::SSL::SSLError
          false
        end

        def close(suppress_close_frame: false)
          begin
            send_close_reply unless suppress_close_frame || @socket.closed?
          rescue Errno::EPIPE, Errno::ECONNRESET, IOError, SystemCallError, OpenSSL::SSL::SSLError, WebSocket::Error
            # ignore
          ensure
            mark_closed
            @socket.close unless @socket.closed?
          end
        end

        private

        def request_headers
          return {} unless @connection_data.respond_to?(:headers)

          @connection_data.get_headers.each_with_object({}) do |(key, value), headers|
            next if key.nil? || value.nil?
            if restricted_websocket_header?(key)
              warn "WebSocket header '#{key}' is reserved and was ignored."
              next
            end

            headers[key.to_s] = value.to_s
          end
        end

        def restricted_websocket_header?(key)
          normalized = key.to_s.downcase
          %w[host upgrade connection].include?(normalized) || normalized.start_with?('sec-websocket-')
        end

        def default_port
          @uri.scheme == 'wss' ? 443 : 80
        end

        def open_socket
          tcp = TCPSocket.new(@host, @port)
          tcp.sync = true
          return tcp unless @uri.scheme == 'wss'

          ctx = OpenSSL::SSL::SSLContext.new
          # Production: VERIFY_PEER and trust store/ca_file.
          ctx.set_params(verify_mode: OpenSSL::SSL::VERIFY_NONE)

          ssl = OpenSSL::SSL::SSLSocket.new(tcp, ctx)
          ssl.sync_close = true
          ssl.hostname = @host if ssl.respond_to?(:hostname=)
          ssl.connect
          ssl
        end

        def perform_handshake
          write_all(@handshake.to_s)
          loop do
            break if @handshake.finished?
            wait_timeout = @connect_timeout && @connect_timeout < 0 ? nil : (@connect_timeout || DEFAULT_TIMEOUT_SECONDS)
            raise "WebSocket connection to #{@uri}: handshake timeout. The server did not complete the WebSocket upgrade in time." unless IO.select([@socket], nil, nil, wait_timeout)

            data = @socket.read_nonblock(4096, exception: false)
            case data
            when nil
              break
            when :wait_readable
              next
            else
              @handshake << data
            end
          end

          raise "WebSocket connection to #{@uri}: handshake failed. The server did not return a valid WebSocket upgrade response." unless @handshake.finished? && @handshake.valid?
        end

        def write_all(data)
          total = 0
          while total < data.bytesize
            written = @socket.write_nonblock(data.byteslice(total..-1), exception: false)
            case written
            when :wait_writable
              IO.select(nil, [@socket])
            when 0, nil
              raise IOError, "socket closed while writing"
            else
              total += written
            end
          end
        rescue IO::WaitWritable
          IO.select(nil, [@socket])
          retry
        end

        def send_pong(data)
          pong = WebSocket::Frame::Outgoing::Client.new(
            version: @handshake.version, type: :pong, data: data
          )
          write_all(pong.to_s)
        rescue Errno::EPIPE, Errno::ECONNRESET, IOError, SystemCallError, OpenSSL::SSL::SSLError
          # ignore
        end

        def send_close_reply
          return if @socket.closed?
          frame = WebSocket::Frame::Outgoing::Client.new(
            version: @handshake.version, type: :close
          )
          write_all(frame.to_s)
        rescue Errno::EPIPE, Errno::ECONNRESET, IOError, SystemCallError, OpenSSL::SSL::SSLError
          # ignore
        end
      end
    end
  end
end
