(* The Sol language server: JSON-RPC over stdio, hand rolled on
   yojson.

   Diagnostics run the same admission the runtime runs, so a rule
   with no decrease witness or a fact violating its schema is an
   error in the editor with the runtime's own message. Completion,
   hover, and definition work from a parse-level symbol scan of the
   document text, so they keep working while the buffer is broken
   mid-edit. Parse errors carry positions; admission errors anchor
   at the top of the file until the loader learns locations. *)

let send json =
  let body = Yojson.Safe.to_string json in
  Printf.printf "Content-Length: %d\r\n\r\n%s" (String.length body) body;
  flush stdout

let read_message () =
  let rec headers length =
    match In_channel.input_line stdin with
    | None -> None
    | Some line -> (
        let line = String.trim line in
        if line = "" then length
        else
          match String.index_opt line ':' with
          | Some i
            when String.lowercase_ascii (String.sub line 0 i) = "content-length"
            ->
              headers
                (int_of_string_opt
                   (String.trim
                      (String.sub line (i + 1) (String.length line - i - 1))))
          | _ -> headers length)
  in
  match headers None with
  | None -> None
  | Some length -> (
      match really_input_string stdin length with
      | body -> ( try Some (Yojson.Safe.from_string body) with _ -> None)
      | exception End_of_file -> None)

let member name json =
  match json with
  | `Assoc fields -> Option.value ~default:`Null (List.assoc_opt name fields)
  | _ -> `Null

let to_string_opt = function `String s -> Some s | _ -> None
let to_int_opt = function `Int i -> Some i | _ -> None

(* ------------------------------------------------------------------ *)
(* Document store and symbol scan.                                     *)
(* ------------------------------------------------------------------ *)

type doc = { lines : string array }

let documents : (string, doc) Hashtbl.t = Hashtbl.create 8

let store uri text =
  Hashtbl.replace documents uri
    { lines = Array.of_list (String.split_on_char '\n' text) }

let is_ident_char c =
  (c >= 'a' && c <= 'z')
  || (c >= 'A' && c <= 'Z')
  || (c >= '0' && c <= '9')
  || c = '_'

(* The identifier under the cursor; a leading # is part of it and
   stripped, so hovering a reference reaches its relation. *)
let token_at doc ~line ~character =
  if line < 0 || line >= Array.length doc.lines then None
  else
    let text = doc.lines.(line) in
    let n = String.length text in
    let i = min character (max 0 (n - 1)) in
    if n = 0 then None
    else
      let i = if i < n && is_ident_char text.[i] then i else i - 1 in
      if i < 0 || not (is_ident_char text.[i]) then None
      else
        let start = ref i and stop = ref i in
        while !start > 0 && is_ident_char text.[!start - 1] do decr start done;
        while !stop < n - 1 && is_ident_char text.[!stop + 1] do incr stop done;
        Some (String.sub text !start (!stop - !start + 1))

(* Definition sites, scanned per line so they survive parse errors:
   a relation is defined by its Relation directive or by the rules
   whose head applies it. *)
let find_after text needle =
  let tl = String.length text and nl = String.length needle in
  let rec go i =
    if i + nl > tl then None
    else if String.sub text i nl = needle then Some (i + nl)
    else go (i + 1)
  in
  go 0

let ident_at text i =
  let n = String.length text in
  let rec skip i = if i < n && text.[i] = ' ' then skip (i + 1) else i in
  let start = skip i in
  let stop = ref start in
  while !stop < n && is_ident_char text.[!stop] do incr stop done;
  if !stop > start then Some (String.sub text start (!stop - start), start)
  else None

let definitions doc name =
  let hits = ref [] in
  Array.iteri
    (fun lineno text ->
      let head_of needle =
        match find_after text needle with
        | Some i -> ident_at text i
        | None -> None
      in
      let relation = head_of "Relation named:" in
      let rule_head =
        match find_after text "head:" with
        | Some i -> (
            match find_after (String.sub text i (String.length text - i)) "[" with
            | Some j -> (
                match ident_at text (i + j) with
                | Some (n, col) -> Some (n, col)
                | None -> None)
            | None -> None)
        | None -> None
      in
      List.iter
        (function
          | Some (n, col) when String.equal n name ->
              hits := (lineno, col) :: !hits
          | _ -> ())
        [ relation; rule_head ])
    doc.lines;
  List.rev !hits

(* Names available for completion: declared relations and rule
   heads, every label spelled in the document, builtins, directives
   and combinators. Harvested from text, so a broken buffer still
   completes. *)
let harvest doc =
  let relations = ref [] and labels = ref [] in
  Array.iter
    (fun text ->
      let n = String.length text in
      let i = ref 0 in
      while !i < n do
        if is_ident_char text.[!i] then begin
          let start = !i in
          while !i < n && is_ident_char text.[!i] do incr i done;
          let word = String.sub text start (!i - start) in
          if !i < n && text.[!i] = ':' then begin
            if word.[0] >= 'a' && word.[0] <= 'z' then
              labels := word :: !labels
          end
          else if
            (word.[0] >= 'A' && word.[0] <= 'Z')
            && start > 0
            && (text.[start - 1] = '[' || text.[start - 1] = '#')
          then relations := word :: !relations
        end
        else if text.[!i] = '%' then i := n
        else incr i
      done)
    doc.lines;
  ( List.sort_uniq String.compare !relations
    |> List.filter (fun name -> not (Sol.Builtin.is_builtin name)),
    List.sort_uniq String.compare !labels )

(* ------------------------------------------------------------------ *)
(* What hover says.                                                    *)
(* ------------------------------------------------------------------ *)

let builtin_hover = function
  | "Integer" -> Some "[Integer value: T]\nmembership in the integers; checks once ground, residuates open"
  | "Natural" -> Some "[Natural value: T]\nmembership in the naturals (an Integer at least 0)"
  | "Rational" -> Some "[Rational value: T]\nmembership in the rationals; every Integer is one, literally"
  | "String" -> Some "[String value: T]\nmembership in the strings"
  | "Equal" -> Some "[Equal left: T right: T]\nunification"
  | "Plus" -> Some "[Plus left: T right: T result: T]\naddition over the rationals; any two ground arguments determine the third"
  | "Times" -> Some "[Times left: T right: T result: T]\nmultiplication over the rationals; solves backward except through zero"
  | "LessThan" -> Some "[LessThan left: T right: T]\ndecides once both sides are ground; contributes interval bounds while open"
  | "LessThanEqual" -> Some "[LessThanEqual left: T right: T]\ndecides once both sides are ground; contributes interval bounds while open"
  | "GreaterThan" -> Some "[GreaterThan left: T right: T]\ndecides once both sides are ground; contributes interval bounds while open"
  | "GreaterEqual" -> Some "[GreaterEqual left: T right: T]\ndecides once both sides are ground; contributes interval bounds while open"
  | "Between" -> Some "[Between value: T lower: T upper: T]\nbounded generation: enumerates the finite integer range"
  | "True" -> Some "[True]\nthe proposition that holds"
  | "World" -> Some "[World]\nthe universe; true of every tuple, scopes denial: [World] :not [...]"
  | "ForAll" -> Some "[ForAll variable: V quantifier: [...] body: [...]]\nuniversal: every quantifier solution satisfies the body; infinite domains discharge symbolically or reject"
  | "Exists" -> Some "[Exists variable: V quantifier: [...] body: [...]]\nexistential: lowers to a conjunction with V renamed apart"
  | "Insert" -> Some "[Insert relation: #R tuple: {...} before: S after: S']\ntransition relation between states (stateful worlds only)"
  | _ -> None

(* For a declared relation: its defining lines, verbatim. The
   honest hover is the source of the definition. *)
let declared_hover doc name =
  match definitions doc name with
  | [] -> None
  | sites ->
      let lines =
        List.map (fun (lineno, _) -> String.trim doc.lines.(lineno)) sites
      in
      Some (String.concat "\n" (List.sort_uniq String.compare lines))

(* ------------------------------------------------------------------ *)
(* Diagnostics.                                                        *)
(* ------------------------------------------------------------------ *)

let diagnostics text =
  match Sol.Front.parse text with
  | Error { Sol.Front.line; column; message } ->
      [ (max 0 (line - 1), max 0 column, message) ]
  | Ok statements -> (
      match Sol.Program.load statements with
      | Error error -> [ (0, 0, Format.asprintf "%a" Sol.Program.pp_error error) ]
      | Ok _ -> [])

let publish uri text =
  let items =
    List.map
      (fun (line, column, message) ->
        `Assoc
          [
            ( "range",
              `Assoc
                [
                  ( "start",
                    `Assoc [ ("line", `Int line); ("character", `Int column) ]
                  );
                  ( "end",
                    `Assoc
                      [ ("line", `Int line); ("character", `Int (column + 1)) ]
                  );
                ] );
            ("severity", `Int 1);
            ("source", `String "sol");
            ("message", `String message);
          ])
      (diagnostics text)
  in
  send
    (`Assoc
      [
        ("jsonrpc", `String "2.0");
        ("method", `String "textDocument/publishDiagnostics");
        ("params", `Assoc [ ("uri", `String uri); ("diagnostics", `List items) ]);
      ])

let respond id result =
  send (`Assoc [ ("jsonrpc", `String "2.0"); ("id", id); ("result", result) ])

let check_notification params =
  let uri = to_string_opt (member "uri" (member "textDocument" params)) in
  let text =
    match to_string_opt (member "text" (member "textDocument" params)) with
    | Some text -> Some text
    | None -> (
        match member "contentChanges" params with
        | `List changes -> (
            match List.rev changes with
            | last :: _ -> to_string_opt (member "text" last)
            | [] -> None)
        | _ -> None)
  in
  match (uri, text) with
  | Some uri, Some text ->
      store uri text;
      publish uri text
  | _ -> ()

(* ------------------------------------------------------------------ *)
(* Requests.                                                           *)
(* ------------------------------------------------------------------ *)

let position params =
  let uri = to_string_opt (member "uri" (member "textDocument" params))
  and line = to_int_opt (member "line" (member "position" params))
  and character = to_int_opt (member "character" (member "position" params)) in
  match (uri, line, character) with
  | Some uri, Some line, Some character -> (
      match Hashtbl.find_opt documents uri with
      | Some doc -> Some (uri, doc, line, character)
      | None -> None)
  | _ -> None

let handle_hover params =
  match position params with
  | None -> `Null
  | Some (_, doc, line, character) -> (
      match token_at doc ~line ~character with
      | None -> `Null
      | Some token -> (
          let text =
            match builtin_hover token with
            | Some text -> Some text
            | None -> declared_hover doc token
          in
          match text with
          | None -> `Null
          | Some value ->
              `Assoc
                [
                  ( "contents",
                    `Assoc
                      [ ("kind", `String "plaintext"); ("value", `String value) ]
                  );
                ]))

let handle_definition uri params =
  match position params with
  | None -> `Null
  | Some (_, doc, line, character) -> (
      match token_at doc ~line ~character with
      | None -> `Null
      | Some token ->
          let locations =
            List.map
              (fun (lineno, col) ->
                `Assoc
                  [
                    ("uri", `String uri);
                    ( "range",
                      `Assoc
                        [
                          ( "start",
                            `Assoc
                              [ ("line", `Int lineno); ("character", `Int col) ]
                          );
                          ( "end",
                            `Assoc
                              [
                                ("line", `Int lineno);
                                ( "character",
                                  `Int (col + String.length token) );
                              ] );
                        ] );
                  ])
              (definitions doc token)
          in
          `List locations)

let directive_items = [ "Relation"; "Rule"; "Query" ]
let combinator_items = [ ":and"; ":or"; ":not" ]
let form_items = [ "ForAll"; "Exists" ]

let handle_completion params =
  match position params with
  | None -> `Null
  | Some (_, doc, _, _) ->
      let relations, labels = harvest doc in
      let item ?detail kind label =
        `Assoc
          (List.concat
             [
               [ ("label", `String label); ("kind", `Int kind) ];
               (match detail with
               | Some d -> [ ("detail", `String d) ]
               | None -> []);
             ])
      in
      (* LSP completion kinds: 7 class, 5 field, 14 keyword, 3
         function. The client filters by prefix. *)
      let items =
        List.concat
          [
            List.map (item 7 ~detail:"relation") relations;
            List.map
              (fun name -> item 3 ~detail:"builtin" name)
              (Sol.Builtin.names @ form_items);
            List.map (fun l -> item 5 ~detail:"label" (l ^ ":")) labels;
            List.map (item 14 ~detail:"directive") directive_items;
            List.map (item 14 ~detail:"combinator") combinator_items;
          ]
      in
      `Assoc [ ("isIncomplete", `Bool false); ("items", `List items) ]

(* ------------------------------------------------------------------ *)
(* The loop.                                                           *)
(* ------------------------------------------------------------------ *)

let run () =
  set_binary_mode_in stdin true;
  set_binary_mode_out stdout true;
  let running = ref true in
  while !running do
    match read_message () with
    | None -> running := false
    | Some message -> (
        let id = member "id" message in
        let params = member "params" message in
        let uri = to_string_opt (member "uri" (member "textDocument" params)) in
        match to_string_opt (member "method" message) with
        | Some "initialize" ->
            respond id
              (`Assoc
                [
                  ( "capabilities",
                    `Assoc
                      [
                        ("textDocumentSync", `Int 1);
                        ("hoverProvider", `Bool true);
                        ("definitionProvider", `Bool true);
                        ( "completionProvider",
                          `Assoc
                            [
                              ( "triggerCharacters",
                                `List [ `String "["; `String "#"; `String ":" ]
                              );
                            ] );
                      ] );
                  ( "serverInfo",
                    `Assoc
                      [
                        ("name", `String "sol-lsp"); ("version", `String "0.2.0");
                      ] );
                ])
        | Some "shutdown" -> respond id `Null
        | Some "exit" -> running := false
        | Some ("textDocument/didOpen" | "textDocument/didChange") ->
            check_notification params
        | Some "textDocument/didClose" -> (
            match uri with
            | Some uri ->
                Hashtbl.remove documents uri;
                send
                  (`Assoc
                    [
                      ("jsonrpc", `String "2.0");
                      ("method", `String "textDocument/publishDiagnostics");
                      ( "params",
                        `Assoc
                          [ ("uri", `String uri); ("diagnostics", `List []) ] );
                    ])
            | None -> ())
        | Some "textDocument/hover" -> respond id (handle_hover params)
        | Some "textDocument/definition" ->
            respond id
              (handle_definition (Option.value ~default:"" uri) params)
        | Some "textDocument/completion" -> respond id (handle_completion params)
        | Some _ when id <> `Null ->
            send
              (`Assoc
                [
                  ("jsonrpc", `String "2.0");
                  ("id", id);
                  ( "error",
                    `Assoc
                      [
                        ("code", `Int (-32601));
                        ("message", `String "method not found");
                      ] );
                ])
        | _ -> ())
  done