require "json"
require "open3"
require "uri"

# All server/protocol work stays in this optional plugin. Host callbacks only enqueue.
module VampLsp
  MAX_BODY = 4 * 1024 * 1024
  MAX_HEADER = 16 * 1024
  QUEUE_SIZE = 128
  MAX_VERSION = (1 << 31) - 1

  def self.clean(text)
    raise ArgumentError, "LSP message must be UTF-8 text" unless text.is_a?(String) && text.valid_encoding?
    text.gsub(/[[:cntrl:]]/, " ")
  end

  def self.uri(path)
    "file://#{URI::DEFAULT_PARSER.escape(File.expand_path(path), /[^A-Za-z0-9\-._~\/]/)}"
  end

  def self.path(uri)
    raise ArgumentError, "definition URI must be a String" unless uri.is_a?(String)
    parsed = URI.parse(uri)
    unless parsed.scheme == "file" && [nil, "", "localhost"].include?(parsed.host) &&
           parsed.userinfo.nil? && parsed.port.nil? && parsed.query.nil? && parsed.fragment.nil? &&
           parsed.path.start_with?("/") && !parsed.path.match?(/%(?![0-9a-fA-F]{2})/)
      raise ArgumentError, "unsupported local file URI: #{clean(uri)}"
    end
    path = URI::DEFAULT_PARSER.unescape(parsed.path).force_encoding(Encoding::UTF_8)
    raise ArgumentError, "invalid UTF-8 file URI" unless path.valid_encoding? && !path.match?(/[[:cntrl:]]/)
    File.expand_path(path)
  end

  def self.identity(path)
    File.realpath(path)
  rescue SystemCallError
    File.expand_path(path)
  end

  def self.lines(text)
    text.gsub(/\r\n?/, "\n").split("\n", -1)
  end

  def self.integer(value, label)
    raise ArgumentError, "invalid LSP #{label}" unless value.is_a?(Integer) && value.between?(0, MAX_VERSION)
    value
  end

  # Expand positions inside a surrogate pair or combining cluster to whole graphemes.
  def self.position(lines, value, round_up: false)
    raise ArgumentError, "invalid LSP position" unless value.is_a?(Hash)
    row = integer(value["line"], "line")
    units = integer(value["character"], "character")
    line = lines.fetch(row) { raise ArgumentError, "LSP line is outside the document" }
    offset = 0
    column = 0
    line.each_grapheme_cluster do |cluster|
      return { "row" => row, "column" => column } if units == offset
      width = cluster.encode(Encoding::UTF_16LE).bytesize / 2
      if units < offset + width
        return { "row" => row, "column" => column + (round_up ? 1 : 0) }
      end
      offset += width
      column += 1
    end
    # LSP specifies oversized character offsets as end-of-line.
    { "row" => row, "column" => column }
  end

  def self.range(lines, value)
    raise ArgumentError, "invalid LSP range" unless value.is_a?(Hash)
    first = value["start"]
    last = value["end"]
    start = position(lines, first)
    finish = position(lines, last, round_up: true)
    if ([integer(first["line"], "line"), integer(first["character"], "character")] <=>
        [integer(last["line"], "line"), integer(last["character"], "character")]) == 1
      raise ArgumentError, "reversed LSP range"
    end
    [start, finish]
  end

  def self.outbound_position(text, row, column)
    line = lines(text).fetch(row) { raise ArgumentError, "definition source row is outside the document" }
    clusters = line.each_grapheme_cluster.to_a
    raise ArgumentError, "definition source column is outside the line" unless column.between?(0, clusters.length)
    { "line" => row, "character" => clusters.take(column).join.encode(Encoding::UTF_16LE).bytesize / 2 }
  end

  class Server
    attr_reader :root, :config, :documents, :pending, :requests, :ready
    attr_accessor :sync, :definition

    def initialize(config, root, events)
      @config, @root, @events = config, root, events
      @documents, @pending, @requests = {}, [], {}
      @ready, @dead, @closing = false, false, false
      @next_id = 1
      @writes = SizedQueue.new(QUEUE_SIZE)
      @stdin, @stdout, @stderr, @wait = Open3.popen3(*config.fetch("command"), chdir: root)
      [@stdin, @stdout, @stderr].each(&:binmode)
      @reader = Thread.new do
        loop { @events.push([:incoming, self, read_frame], true) }
      rescue StandardError => error
        failed(error)
      end
      @writer = Thread.new do
        loop do
          bytes = @writes.pop
          break if bytes.nil?
          @stdin.write(bytes)
          @stdin.flush
        end
      rescue StandardError => error
        failed(error)
      end
      @drain = Thread.new do
        # Server stderr is intentionally not editor status or host protocol output.
        loop { @stderr.readpartial(8192) }
      rescue EOFError, IOError
        nil
      end
      request("initialize", {
        "processId" => Process.pid,
        "clientInfo" => { "name" => "vamp" },
        "rootUri" => VampLsp.uri(root),
        "workspaceFolders" => [{ "uri" => VampLsp.uri(root), "name" => File.basename(root) }],
        "capabilities" => {
          "general" => { "positionEncodings" => ["utf-16"] },
          "workspace" => { "configuration" => true, "workspaceFolders" => true },
          "textDocument" => {
            "synchronization" => { "dynamicRegistration" => false, "willSave" => false,
                                   "willSaveWaitUntil" => false, "didSave" => true },
            "definition" => { "dynamicRegistration" => false, "linkSupport" => true },
            "publishDiagnostics" => { "relatedInformation" => false, "versionSupport" => true }
          }
        }
      }, :initialize)
    rescue StandardError
      stop
      raise
    end

    def dead?
      @dead
    end

    def initialized!
      @ready = true
    end

    def request(method, params, context)
      raise ArgumentError, "LSP request IDs exhausted" if @next_id > MAX_VERSION
      raise ThreadError, "LSP pending request queue is full" if @requests.length >= QUEUE_SIZE
      id = @next_id
      @next_id += 1
      @requests[id] = context
      frame = { "jsonrpc" => "2.0", "id" => id, "method" => method }
      frame["params"] = params unless params.nil?
      send_frame(frame)
    end

    def notify(method, params)
      frame = { "jsonrpc" => "2.0", "method" => method }
      frame["params"] = params unless params.nil?
      send_frame(frame)
    end

    def reply(id, result: nil, error: nil)
      frame = { "jsonrpc" => "2.0", "id" => id }
      frame[error ? "error" : "result"] = error || result
      send_frame(frame)
    end

    def shutdown
      return if @dead || @closing
      @closing = true
      if @ready
        request("shutdown", nil, :shutdown)
      else
        stop
      end
    end

    def stop
      return if @dead
      @dead = true
      [@stdin, @stdout, @stderr].compact.each { |io| io.close rescue nil }
      [@reader, @writer, @drain].compact.each { |thread| thread.kill unless thread == Thread.current }
      if @wait
        (Process.kill("TERM", @wait.pid) rescue nil) if @wait.alive?
        # Reaping and forced termination never block the host callback/dispatcher.
        Thread.new do
          unless @wait.join(0.2)
            Process.kill("KILL", @wait.pid) rescue nil
            @wait.join
          end
        end
      end
    end

    private

    def send_frame(frame)
      raise IOError, "LSP server is unavailable" if @dead
      body = JSON.generate(frame).encode(Encoding::UTF_8)
      raise ArgumentError, "LSP outgoing frame exceeds 4 MiB" if body.bytesize > MAX_BODY
      @writes.push("Content-Length: #{body.bytesize}\r\n\r\n".b + body.b, true)
    rescue ThreadError
      raise ThreadError, "LSP outgoing queue is full"
    end

    def read_frame
      size = 0
      length = nil
      loop do
        line = @stdout.gets("\r\n", MAX_HEADER + 1)
        raise EOFError, "LSP server exited" unless line
        size += line.bytesize
        raise ArgumentError, "LSP headers exceed 16 KiB" if size > MAX_HEADER
        raise ArgumentError, "LSP headers must be ASCII" unless line.ascii_only?
        raise ArgumentError, "invalid LSP header newline" unless line.end_with?("\r\n")
        break if line == "\r\n"
        name, value = line.chomp("\r\n").split(":", 2)
        raise ArgumentError, "invalid LSP header" unless value
        if name.downcase == "content-length"
          raise ArgumentError, "invalid LSP Content-Length" if length || !value.strip.match?(/\A[0-9]+\z/)
          length = value.to_i
        elsif name.downcase == "content-type" && value.match?(/charset=/i) &&
              !value.match?(/charset\s*=\s*utf-?8\s*\z/i)
          raise ArgumentError, "unsupported LSP content encoding"
        end
      end
      raise ArgumentError, "invalid LSP body length" unless length && length.between?(1, MAX_BODY)
      body = @stdout.read(length)
      raise EOFError, "truncated LSP body" unless body && body.bytesize == length
      body.force_encoding(Encoding::UTF_8)
      raise ArgumentError, "invalid UTF-8 LSP body" unless body.valid_encoding?
      JSON.parse(body)
    end

    def failed(error)
      return if @dead || @closing
      @events.push([:failure, self, error.message], true)
    rescue ThreadError
      begin
        Vamp.publish([{ "type" => "message", "text" => "LSP incoming queue is full" }])
      ensure
        Kernel.exit!(1) # Host EOF retires the feature and clears all owned diagnostics.
      end
    end
  end

  def self.clear(doc)
    { "type" => "diagnostics", "buffer_id" => doc.fetch("id"),
      "revision" => doc.fetch("revision"), "items" => [] }
  end

  class Client
    def initialize(options)
      servers = options["servers"]
      unless servers.is_a?(Array) && !servers.empty? && servers.length <= 16
        raise ArgumentError, "LSP options.servers must contain 1..16 server configurations"
      end
      @configs = servers.map do |config|
        unless config.is_a?(Hash) && config["language_id"].is_a?(String) && !config["language_id"].empty? &&
               config["extensions"].is_a?(Array) && !config["extensions"].empty? &&
               config["extensions"].all? { |ext| ext.is_a?(String) && ext.match?(/\A[^.\/\\\s\x00]+\z/) } &&
               config["command"].is_a?(Array) && !config["command"].empty? &&
               config["command"].all? { |arg| arg.is_a?(String) && !arg.include?("\0") } &&
               !config["command"].first.empty?
          raise ArgumentError, "LSP server requires language_id, extensions, and command argv"
        end
        config
      end
      @events = SizedQueue.new(QUEUE_SIZE)
      @documents, @servers = {}, {}
      @failure = nil
      @worker = Thread.new { run }
    end

    def enqueue(*event)
      raise IOError, @failure if @failure
      @events.push(event, true)
      nil
    rescue ThreadError
      raise IOError, "LSP document/command queue is full"
    end

    private

    def run
      loop do
        event, *args = @events.pop
        case event
        when :document then document(*args)
        when :definition then definition(*args)
        when :diagnostics then diagnostics_status(*args)
        when :incoming then incoming(*args)
        when :failure then fail_server(*args)
        when :shutdown
          @servers.values.compact.each(&:shutdown)
          # The host process-group cleanup is the final safety net, including descendants.
        end
      rescue StandardError => error
        server = [:incoming, :failure].include?(event) ? args.first : nil
        if server.is_a?(Server) && !error.is_a?(ThreadError)
          fail_server(server, error.message)
        else
          @failure = "LSP client: #{VampLsp.clean(error.message)}"
          @servers.values.compact.each(&:stop)
          begin
            Vamp.publish([{ "type" => "message", "text" => @failure }])
          ensure
            Kernel.exit!(1)
          end
        end
      end
    end

    def message(text, request_id: nil)
      Vamp.publish([{ "type" => "message", "text" => "LSP: #{VampLsp.clean(text)}" }], request_id: request_id)
    end

    def server_for(doc)
      return unless doc["path"]
      extension = File.extname(doc["path"]).delete_prefix(".")
      index = @configs.index { |config| config["extensions"].include?(extension) }
      return unless index
      key = [index, File.expand_path(doc.fetch("root"))]
      return @servers[key] if @servers.key?(key)
      begin
        @servers[key] = Server.new(@configs[index], key.last, @events)
      rescue StandardError => error
        # Remember failed starts: no automatic retries or restart loop.
        @servers[key] = nil
        message("#{@configs[index]["language_id"]} server failed to start: #{error.message}")
      end
    end

    def document(event)
      kind = event.fetch("event")
      snapshot = event.fetch("document")
      id = snapshot.fetch("id")
      old = @documents[id]
      if kind == "close"
        if old && (server = old["server"])
          close_document(server, old)
        end
        @documents.delete(id)
        return
      end
      doc = snapshot.merge("uri" => snapshot["path"] && VampLsp.uri(snapshot["path"]),
                           "server" => old ? old["server"] : server_for(snapshot),
                           "version" => old ? old["version"] : 1,
                           "opened" => old ? old["opened"] : false,
                           "saved" => kind == "save")
      server = doc["server"]
      @documents[id] = doc
      return unless server && !server.dead?
      if old && old["revision"] != doc["revision"]
        raise ArgumentError, "LSP document version exhausted" if doc["version"] == MAX_VERSION
        doc["version"] += 1
        Vamp.publish([VampLsp.clear(doc)])
      end
      server.documents[doc["uri"]] = doc
      return unless server.ready
      if !doc["opened"]
        open_document(server, doc)
      elsif old && old["revision"] != doc["revision"]
        change_document(server, old, doc)
      end
      save_document(server, doc) if kind == "save"
    end

    def open_document(server, doc)
      doc["opened"] = true
      return unless server.sync["openClose"]
      server.notify("textDocument/didOpen", {
        "textDocument" => { "uri" => doc["uri"], "languageId" => server.config["language_id"],
                            "version" => doc["version"], "text" => doc["text"] }
      })
    end

    def change_document(server, old, doc)
      kind = server.sync["change"]
      return if kind == 0
      change = { "text" => doc["text"] }
      if kind == 2
        old_lines = VampLsp.lines(old["text"])
        change["range"] = { "start" => { "line" => 0, "character" => 0 },
                            "end" => { "line" => old_lines.length - 1,
                                       "character" => old_lines.last.encode(Encoding::UTF_16LE).bytesize / 2 } }
      end
      server.notify("textDocument/didChange", {
        "textDocument" => { "uri" => doc["uri"], "version" => doc["version"] },
        "contentChanges" => [change]
      })
    end

    def save_document(server, doc)
      save = server.sync["save"]
      return unless save
      params = { "textDocument" => { "uri" => doc["uri"] } }
      params["text"] = doc["text"] if save.is_a?(Hash) && save["includeText"]
      server.notify("textDocument/didSave", params)
      doc["saved"] = false
    end

    def close_document(server, doc)
      server.documents.delete(doc["uri"])
      return unless server.ready && !server.dead? && doc["opened"] && server.sync["openClose"]
      server.notify("textDocument/didClose", { "textDocument" => { "uri" => doc["uri"] } })
    end

    def definition(source)
      doc = @documents[source.fetch("id")]
      return unless doc && doc["revision"] == source["revision"]
      server = doc["server"]
      unless server && !server.dead?
        message("no running server for this document", request_id: source["request_id"])
        return
      end
      unless server.ready
        raise ThreadError, "LSP pending command queue is full" if server.pending.length == QUEUE_SIZE
        server.pending << source
        return
      end
      unless server.definition
        message("server does not support go-to-definition", request_id: source["request_id"])
        return
      end
      server.request("textDocument/definition", {
        "textDocument" => { "uri" => doc["uri"] },
        "position" => VampLsp.outbound_position(source["text"], source["row"], source["column"])
      }, source)
    end

    def diagnostics_status(id)
      doc = @documents[id]
      server = doc && doc["server"]
      message(if !server || server.dead?
                "no running server for this document"
              elsif !server.ready
                "server initializing"
              else
                "diagnostics update automatically from #{server.config["language_id"]}"
              end)
    end

    def incoming(server, frame)
      return if server.dead?
      unless frame.is_a?(Hash) && frame["jsonrpc"] == "2.0"
        raise ArgumentError, "invalid LSP JSON-RPC frame"
      end
      if frame.key?("method")
        method = frame["method"]
        raise ArgumentError, "invalid LSP method" unless method.is_a?(String)
        if frame.key?("id")
          server_request(server, frame)
        elsif method == "textDocument/publishDiagnostics"
          publish_diagnostics(server, frame["params"])
        elsif method == "window/showMessage"
          params = frame["params"]
          raise ArgumentError, "invalid LSP showMessage" unless params.is_a?(Hash)
          message(params["message"])
        end
        return
      end
      id = frame["id"]
      unless id.is_a?(Integer) && server.requests.key?(id) &&
             (frame.key?("result") ^ frame.key?("error"))
        raise ArgumentError, "unexpected LSP response"
      end
      context = server.requests.delete(id)
      if context == :shutdown
        server.notify("exit", nil)
        return
      end
      if frame.key?("error")
        error = frame["error"]
        unless error.is_a?(Hash) && error["code"].is_a?(Integer) && error["message"].is_a?(String)
          raise ArgumentError, "invalid LSP error response"
        end
        raise IOError, "initialization failed: #{error["message"]}" if context == :initialize
        message("definition failed: #{error["message"]}", request_id: context["request_id"])
      elsif context == :initialize
        initialize_result(server, frame["result"])
      else
        definition_result(context, frame["result"])
      end
    end

    def initialize_result(server, result)
      capabilities = result.is_a?(Hash) && result["capabilities"]
      raise ArgumentError, "invalid LSP initialize result" unless capabilities.is_a?(Hash)
      encoding = capabilities.fetch("positionEncoding", "utf-16")
      raise ArgumentError, "server selected unsupported position encoding #{encoding}" unless encoding == "utf-16"
      sync = capabilities.fetch("textDocumentSync", 0)
      if sync.is_a?(Integer) && sync.between?(0, 2)
        sync = { "openClose" => sync != 0, "change" => sync, "save" => false }
      elsif sync.is_a?(Hash)
        sync = { "openClose" => sync.fetch("openClose", false), "change" => sync.fetch("change", 0),
                 "save" => sync.fetch("save", false) }
      else
        raise ArgumentError, "invalid LSP textDocumentSync capability"
      end
      unless [true, false].include?(sync["openClose"]) && [0, 1, 2].include?(sync["change"]) &&
             ([true, false].include?(sync["save"]) ||
              (sync["save"].is_a?(Hash) && [true, false].include?(sync["save"].fetch("includeText", false))))
        raise ArgumentError, "invalid LSP synchronization options"
      end
      server.sync = sync
      server.definition = capabilities["definitionProvider"] == true || capabilities["definitionProvider"].is_a?(Hash)
      server.notify("initialized", {})
      server.initialized!
      server.documents.each_value do |doc|
        open_document(server, doc)
        save_document(server, doc) if doc["saved"]
      end
      pending = server.pending.dup
      server.pending.clear
      pending.each { |source| definition(source) }
    end

    def server_request(server, frame)
      id = frame["id"]
      raise ArgumentError, "invalid LSP server request ID" unless id.is_a?(Integer) || id.is_a?(String)
      case frame["method"]
      when "workspace/configuration"
        params = frame["params"]
        items = params.is_a?(Hash) && params["items"]
        raise ArgumentError, "invalid LSP configuration request" unless items.is_a?(Array)
        server.reply(id, result: items.map { nil })
      when "workspace/workspaceFolders"
        server.reply(id, result: [{ "uri" => VampLsp.uri(server.root), "name" => File.basename(server.root) }])
      when "window/showMessageRequest"
        server.reply(id, result: nil)
      else
        server.reply(id, error: { "code" => -32601, "message" => "Method not supported by vamp" })
      end
    end

    def publish_diagnostics(server, params)
      unless params.is_a?(Hash) && params["uri"].is_a?(String) && params["diagnostics"].is_a?(Array)
        raise ArgumentError, "invalid LSP publishDiagnostics"
      end
      doc = server.documents[params["uri"]]
      return unless doc
      if params.key?("version") && !params["version"].nil?
        version = params["version"]
        unless version.is_a?(Integer) && version.between?(-(1 << 31), MAX_VERSION)
          raise ArgumentError, "invalid LSP diagnostic version"
        end
        return if version < doc["version"]
        raise ArgumentError, "LSP diagnostics refer to an unknown future version" if version > doc["version"]
      end
      # LSP permits no version: only arrival-time revision is knowable in that case.
      text_lines = VampLsp.lines(doc["text"])
      items = params["diagnostics"].map do |item|
        raise ArgumentError, "invalid LSP diagnostic" unless item.is_a?(Hash)
        start, finish = VampLsp.range(text_lines, item["range"])
        severity = { 1 => "error", 2 => "warning", 3 => "information", 4 => "hint" }[item.fetch("severity", 1)]
        raise ArgumentError, "invalid LSP diagnostic severity" unless severity
        { "start" => start, "end" => finish, "severity" => severity,
          "message" => VampLsp.clean(item["message"]) }
      end
      Vamp.publish([{ "type" => "diagnostics", "buffer_id" => doc["id"],
                      "revision" => doc["revision"], "items" => items }])
    end

    def definition_result(source, result)
      doc = @documents[source["id"]]
      return unless doc && doc["revision"] == source["revision"]
      locations = result.nil? ? [] : (result.is_a?(Array) ? result : [result])
      if locations.empty?
        message("no definition found", request_id: source["request_id"])
        return
      end
      location = locations.first
      raise ArgumentError, "invalid LSP definition location" unless location.is_a?(Hash)
      link = location.key?("targetUri")
      path = VampLsp.path(location[link ? "targetUri" : "uri"])
      identity = VampLsp.identity(path)
      cached = @documents.values.find { |candidate| candidate["path"] && VampLsp.identity(candidate["path"]) == identity }
      text = cached ? cached["text"] : File.read(path, encoding: Encoding::UTF_8)
      raise ArgumentError, "definition target is not UTF-8" unless text.valid_encoding?
      text_lines = VampLsp.lines(text)
      VampLsp.range(text_lines, location["targetRange"]) if link
      start, = VampLsp.range(text_lines, location[link ? "targetSelectionRange" : "range"])
      effects = [{ "type" => "navigate", "buffer_id" => source["id"], "revision" => source["revision"],
                   "path" => path, "row" => start["row"], "column" => start["column"] }]
      if locations.length > 1
        effects << { "type" => "message", "text" => "LSP: selected first of #{locations.length} definitions" }
      end
      Vamp.publish(effects, request_id: source["request_id"])
    rescue SystemCallError, ArgumentError => error
      message("definition failed: #{error.message}", request_id: source["request_id"])
    end

    def fail_server(server, reason)
      return if server.dead?
      server.stop
      Vamp.publish(server.documents.values.map { |doc| VampLsp.clear(doc) } +
                   [{ "type" => "message", "text" => "LSP server stopped: #{VampLsp.clean(reason)}" }])
    end
  end
end

Vamp.plugin "lsp", requires: ["text", "diagnostics"], async: true
client = VampLsp::Client.new(Vamp.options)
Vamp.on_document { |event| client.enqueue(:document, event) }
Vamp.command "lsp_definition" do |editor, _args|
  client.enqueue(:definition, { "id" => editor.buffer_id, "revision" => editor.revision,
                                "text" => editor.text, "row" => editor.row - 1,
                                "column" => editor.column - 1, "request_id" => editor.request_id })
end
Vamp.command("lsp_diagnostics") { |editor, _args| client.enqueue(:diagnostics, editor.buffer_id) }
Vamp.on_shutdown { client.enqueue(:shutdown) }