open Import

(* Mini clone of Dune_lang.Decoder. Main advantage is that it forbids all the
   crazy stuff and is automatically bi-directional *)

(* TODO error handling is complete crap for now.

   This should be unified with [Dune_lang.Decoder] eventually. *)

type error =
  | Parse_error of
      { message : string
      ; payload : (string * Sexp.t) list
      }
  | Version_error of
      { since : int * int
      ; until : (int * int) option
      ; message : string
      ; payload : (string * Sexp.t) list
      }

let dyn_of_error =
  let open Dyn in
  function
  | Version_error { message; payload; until; since } ->
    record
      [ "message", string message
      ; "payload", list (pair string Sexp.to_dyn) payload
      ; "until", option (pair int int) until
      ; "since", (pair int int) since
      ]
  | Parse_error { message; payload } ->
    record
      [ "message", string message; "payload", list (pair string Sexp.to_dyn) payload ]
;;

exception Of_sexp of error

let raise_of_sexp ?(payload = []) message =
  raise (Of_sexp (Parse_error { message; payload }))
;;

let raise_version_error ?until ?(payload = []) ~since message =
  raise (Of_sexp (Version_error { since; until; message; payload }))
;;

let () =
  Printexc.register_printer (function
    | Of_sexp (Parse_error { message; payload }) ->
      Some (message ^ " " ^ Sexp.to_string (Sexp.record payload))
    | _ -> None)
;;

module Fields = struct
  type t = Unparsed of Sexp.t String.Map.t

  let check_empty (Unparsed s) =
    if not (String.Map.is_empty s)
    then (
      let payload =
        [ ( "unparsed"
          , Sexp.List
              (String.Map.to_list s
               |> List.map ~f:(fun (k, v) -> Sexp.List [ Sexp.Atom k; v ])) )
        ]
      in
      raise_of_sexp ~payload "unexpected fields")
  ;;

  let empty = Unparsed String.Map.empty

  let merge (Unparsed a) (Unparsed b) =
    Unparsed
      (String.Map.union a b ~f:(fun _ _ _ ->
         (* field names are guaranteed to be different at construction time in
            [Both] *)
         assert false))
  ;;

  let of_field name sexp = Unparsed (String.Map.singleton name sexp)

  let of_sexp (x : Sexp.t) =
    match x with
    | Atom _ -> raise_of_sexp "Unexpected atom"
    | List x ->
      (match
         String.Map.of_list_map x ~f:(function
           | List [ Atom s; v ] -> s, v
           | _ -> raise_of_sexp "unable to read field")
       with
       | Error (s, _, _) -> raise_of_sexp "duplicate fields" ~payload:[ "field", Atom s ]
       | Ok s -> Unparsed s)
  ;;

  let optional (Unparsed t) name =
    match String.Map.find t name with
    | None -> None, Unparsed t
    | Some v -> Some v, Unparsed (String.Map.remove t name)
  ;;

  let required t name =
    let r, t = optional t name in
    match r with
    | Some s -> s, t
    | None -> raise_of_sexp "missing required field" ~payload:[ "name", Atom name ]
  ;;

  let to_sexp (Unparsed t) : Sexp.t =
    List (String.Map.to_list t |> List.map ~f:(fun (k, v) -> Sexp.List [ Atom k; v ]))
  ;;
end

type values = Sexp.t
type fields = Fields.t

type version =
  { since : int * int
  ; until : (int * int) option
  }

type ('a, 'kind) t =
  | String : (string, values) t
  | Int : (int, values) t
  | Float : (float, values) t
  | Unit : (unit, values) t
  | Char : (char, values) t
  | Iso : ('a, 'kind) t * ('a -> 'b) * ('b -> 'a) -> ('b, 'kind) t
  | Iso_result : ('a, 'kind) t * ('a -> ('b, exn) result) * ('b -> 'a) -> ('b, 'kind) t
  | Version : ('a, 'kind) t * version -> ('a, 'kind) t
  | Both :
      (* Invariant: field names must be different *)
      ('a, fields) t
      * ('b, fields) t
      -> ('a * 'b, fields) t
  | Sexp : (Sexp.t, values) t
  | List : ('a, values) t -> ('a list, values) t
  | Field : string * 'a field -> ('a, fields) t
  | Enum : (string * 'a) list -> ('a, values) t
  | Sum : 'a econstr list * ('a -> case) -> ('a, values) t
  | Pair : ('a, values) t * ('b, values) t -> ('a * 'b, values) t
  | Triple : ('a, values) t * ('b, values) t * ('c, values) t -> ('a * 'b * 'c, values) t
  | Fdecl : int * ('a, 'k) t Fdecl.t -> ('a, 'k) t
  | Either :
      (* Invariant: field names must be different *)
      ('a, fields) t
      * ('b, fields) t
      -> (('a, 'b) Either.t, fields) t
  | Record : ('a, fields) t -> ('a, values) t

and ('a, 'arg) constr =
  { (* TODO allow constructors without an argument *)
    name : string
  ; arg : ('arg, values) t
  ; inj : 'arg -> 'a
  }

and 'a econstr = Constr : ('a, 'arg) constr -> 'a econstr
and case = Case : 'arg * ('a, 'arg) constr -> case

and 'a field =
  | Required : ('a, values) t -> 'a field
  | Optional : ('a, values) t -> 'a option field

and 'k ret =
  | Values : values ret
  | Fields : Fields.t -> fields ret

type 'a value = ('a, values) t

let case a c = Case (a, c)
let constr name arg inj = { name; arg; inj }
let econstr c = Constr c
let both x y = Both (x, y)
let list x = List x
let sum x y = Sum (x, y)
let pair x y = Pair (x, y)
let triple x y z = Triple (x, y, z)

let discard_values ((a, x) : _ * values ret) =
  match (x : values ret) with
  | Values -> a
;;

let string = String
let int = Int
let float = Float
let unit = Unit

let option x =
  let none = constr "None" unit (fun () -> None) in
  let some = constr "Some" x (fun x -> Some x) in
  sum
    [ econstr none; econstr some ]
    (function
      | None -> case () none
      | Some s -> case s some)
;;

let char = Char

let sexp_for_digest t =
  let rec iter : type a b. int list -> (a, b) t -> Sexp.t =
    fun ids -> function
      | String -> Atom "String"
      | Int -> Atom "Int"
      | Float -> Atom "Float"
      | Unit -> Atom "Unit"
      | Char -> Atom "Char"
      | Iso (t, _, _) -> List [ Atom "Iso"; iter ids t ]
      | Iso_result (t, _, _) -> List [ Atom "Iso_result"; iter ids t ]
      | Version (t, { since = a, b; until }) ->
        let items : Sexp.t list =
          [ Atom "Version"
          ; iter ids t
          ; List [ Atom "since"; Atom (Int.to_string a); Atom (Int.to_string b) ]
          ]
        in
        let items =
          match until with
          | None -> items
          | Some (a, b) ->
            items
            @ [ List [ Atom "until"; Atom (Int.to_string a); Atom (Int.to_string b) ] ]
        in
        List items
      | Both (a, b) -> List [ Atom "Both"; iter ids a; iter ids b ]
      | Sexp -> Atom "Sexp"
      | List t -> List [ Atom "List"; iter ids t ]
      | Field (name, field) ->
        let field : Sexp.t =
          match field with
          | Required t -> List [ Atom "Required"; iter ids t ]
          | Optional t -> List [ Atom "Optional"; iter ids t ]
        in
        List [ Atom "Field"; Atom name; field ]
      | Enum cases ->
        List (Atom "Enum" :: List.map cases ~f:(fun (name, _) : Sexp.t -> Atom name))
      | Sum (constrs, _) ->
        List
          (Atom "Sum"
           :: List.map constrs ~f:(fun (Constr { name; arg; inj = _ }) : Sexp.t ->
             List [ Atom name; iter ids arg ]))
      | Pair (a, b) -> List [ Atom "Pair"; iter ids a; iter ids b ]
      | Triple (a, b, c) -> List [ Atom "Triple"; iter ids a; iter ids b; iter ids c ]
      | Fdecl (id, fdecl) ->
        (* Although the id is represented as an auto-incrementing integer, we
         find De Bruijn indices to put in the digest so that equivalent
         structures produce the same digest. *)
        (match List.findi ids ~f:(Int.equal id) with
         | Some (_, index) -> List [ Atom "Recurse"; Atom (Int.to_string index) ]
         | None -> List [ Atom "Fixpoint"; iter (id :: ids) (Fdecl.get fdecl) ])
      | Either (a, b) -> List [ Atom "Either"; iter ids a; iter ids b ]
      | Record t -> List [ Atom "Record"; iter ids t ]
  in
  iter [] t
;;

let to_sexp : 'a. ('a, values) t -> 'a -> Sexp.t =
  fun t a ->
  let rec loop : type a k. (a, k) t -> a -> k =
    fun t a ->
    match t with
    | String -> Atom a
    | Int -> Atom (Int.to_string a)
    | Float -> Atom (Float.to_string a)
    | Unit -> List []
    | Char -> Atom (String.make 1 a)
    | Sexp -> a
    | Version (t, _) -> loop t a
    | Fdecl (_, t) -> loop (Fdecl.get t) a
    | List t -> List (List.map a ~f:(loop t))
    | Pair (x, y) ->
      let a, b = a in
      List [ loop x a; loop y b ]
    | Triple (x, y, z) ->
      let a, b, c = a in
      List [ loop x a; loop y b; loop z c ]
    | Record r ->
      let fields = loop r a in
      Fields.to_sexp fields
    | Field (name, spec) ->
      (match spec with
       | Required t -> Fields.of_field name (loop t a)
       | Optional t ->
         (match a with
          | None -> Fields.empty
          | Some a -> Fields.of_field name (loop t a)))
    | Iso_result (t, _, from) -> loop t (from a)
    | Iso (t, _, from) -> loop t (from a)
    | Both (x, y) ->
      let x = loop x (fst a) in
      let y = loop y (snd a) in
      Fields.merge x y
    | Either (x, y) ->
      (match a with
       | Left a -> loop x a
       | Right a -> loop y a)
    | Sum (_, constr) ->
      let (Case (a, constr)) = constr a in
      let arg = loop constr.arg a in
      Sexp.List [ Atom constr.name; arg ]
    | Enum choices ->
      (match
         List.find_map choices ~f:(fun (s, a') ->
           if Poly.equal a a' then Some s else None)
       with
       | Some v -> Atom v
       | None ->
         let open Dyn in
         Code_error.raise
           "enum does not include this value"
           [ "valid values", list (fun (x, _) -> string x) choices ])
  in
  loop t a
;;

let check_version ~version ~since ~until _ctx =
  if
    version < since
    ||
    match until with
    | None -> false
    | Some until -> version > until
  then raise_version_error ?until ~since "invalid version"
;;

let of_sexp : 'a. ('a, values) t -> version:int * int -> Sexp.t -> 'a =
  fun t ~version sexp ->
  let rec loop : type a k. (a, k) t -> k -> a * k ret =
    fun (type a k) (t : (a, k) t) (ctx : k) : (a * k ret) ->
    match t with
    | String ->
      (match ctx with
       | Atom s -> s, Values
       | List _ as list ->
         raise_of_sexp ~payload:[ "list", list ] "string: expected atom. received list")
    | Int ->
      (match ctx with
       | List _ as list ->
         raise_of_sexp ~payload:[ "list", list ] "int: expected atom. received list"
       | Atom s ->
         (match Int.of_string s with
          | None -> raise_of_sexp "unable to read int"
          | Some i -> i, Values))
    | Float ->
      (match ctx with
       | List _ as list ->
         raise_of_sexp ~payload:[ "list", list ] "float: expected atom. received list"
       | Atom s ->
         (match Float.of_string_opt s with
          | None -> raise_of_sexp "unable to read float"
          | Some i -> i, Values))
    | Unit ->
      (match ctx with
       | List [] -> (), Values
       | _ -> raise_of_sexp "expected empty list")
    | Char ->
      (match ctx with
       | Atom s ->
         if String.length s = 1
         then s.[0], Values
         else raise_of_sexp "expected only a single character"
       | List _ -> raise_of_sexp "expected a string of length 1")
    | Sexp -> ctx, Values
    | Version (t, { since; until }) ->
      check_version ~version ~since ~until ctx;
      loop t ctx
    | Fdecl (_, t) -> loop (Fdecl.get t) ctx
    | List t ->
      (match ctx with
       | List xs -> List.map xs ~f:(fun x -> discard_values (loop t x)), Values
       | Atom _ -> raise_of_sexp "expected list")
    | Pair (x, y) ->
      (match ctx with
       | List [ a; b ] ->
         let a, Values = loop x a in
         let b, Values = loop y b in
         (a, b), Values
       | _ -> raise_of_sexp "expected field entry")
    | Triple (x, y, z) ->
      (match ctx with
       | List [ a; b; c ] ->
         let a, Values = loop x a in
         let b, Values = loop y b in
         let c, Values = loop z c in
         (a, b, c), Values
       | _ -> raise_of_sexp "expected field entry")
    | Record (r : (a, fields) t) ->
      let (fields : Fields.t) = Fields.of_sexp ctx in
      let a, Fields f = loop r fields in
      Fields.check_empty f;
      a, Values
    | Field (name, spec) ->
      (match spec with
       | Required v ->
         let field, rest = Fields.required ctx name in
         let t, Values = loop v field in
         t, Fields rest
       | Optional v ->
         let field, rest = Fields.optional ctx name in
         let t =
           match field with
           | None -> None
           | Some f ->
             let a, Values = loop v f in
             Some a
         in
         t, Fields rest)
    | Either (x, y) ->
      (try
         (* TODO share computation somehow *)
         let a, x = loop x ctx in
         Left a, x
       with
       | Of_sexp _ ->
         let a, y = loop y ctx in
         Right a, y)
    | Iso (t, f, _) ->
      let a, k = loop t ctx in
      f a, k
    | Iso_result (t, f, _) ->
      let a, k = loop t ctx in
      (match f a with
       | Error exn -> raise exn
       | Ok a -> a, k)
    | Both (x, y) ->
      let a, Fields k = loop x ctx in
      let b, k = loop y k in
      (a, b), k
    | Sum (constrs, _) ->
      (match ctx with
       | List [ Atom head; args ] ->
         (match
            List.find_map constrs ~f:(fun (Constr c) ->
              if head = c.name
              then
                Some
                  (let a, k = loop c.arg args in
                   c.inj a, k)
              else None)
          with
          | None -> raise_of_sexp "invalid constructor name"
          | Some p -> p)
       | _ -> raise_of_sexp "expected constructor")
    | Enum choices ->
      (match ctx with
       | List _ -> raise_of_sexp "expected list"
       | Atom a ->
         (match List.assoc choices a with
          | None -> raise_of_sexp "unable to read enum"
          | Some s -> s, Values))
  in
  discard_values (loop t sexp)
;;

let of_sexp conv ~version sexp =
  match of_sexp conv ~version sexp with
  | s -> Ok s
  | exception Of_sexp e -> Error e
;;

let record r = Record r
let either x y = Either (x, y)
let iso a t f = Iso (a, t, f)
let iso_result a t f = Iso_result (a, t, f)
let version ?until t ~since = Version (t, { until; since })
let field name spec = Field (name, spec)
let enum choices = Enum choices

let three a b c =
  iso (Both (a, Both (b, c))) (fun (x, (y, z)) -> x, y, z) (fun (x, y, z) -> x, (y, z))
;;

let four a b c d =
  iso
    (both (both a b) (both c d))
    (fun ((w, x), (y, z)) -> w, x, y, z)
    (fun (w, x, y, z) -> (w, x), (y, z))
;;

let five a b c d e =
  iso
    (both (both a b) (three c d e))
    (fun ((a, b), (c, d, e)) -> a, b, c, d, e)
    (fun (a, b, c, d, e) -> (a, b), (c, d, e))
;;

let six a b c d e f =
  iso
    (both (three a b c) (three d e f))
    (fun ((a, b, c), (d, e, f)) -> a, b, c, d, e, f)
    (fun (a, b, c, d, e, f) -> (a, b, c), (d, e, f))
;;

let seven a b c d e f g =
  iso
    (both (three a b c) (four d e f g))
    (fun ((a, b, c), (d, e, f, g)) -> a, b, c, d, e, f, g)
    (fun (a, b, c, d, e, f, g) -> (a, b, c), (d, e, f, g))
;;

let eight a b c d e f g h =
  iso
    (both (four a b c d) (four e f g h))
    (fun ((a, b, c, d), (e, f, g, h)) -> a, b, c, d, e, f, g, h)
    (fun (a, b, c, d, e, f, g, h) -> (a, b, c, d), (e, f, g, h))
;;

let sexp = Sexp
let required x = Required x
let optional x = Optional x
let fdecl_id = ref 0

let fixpoint f =
  let fdecl = Fdecl.create Dyn.opaque in
  let id = !fdecl_id in
  incr fdecl_id;
  let result = Fdecl (id, fdecl) in
  Fdecl.set fdecl (f result);
  result
;;

let error e = raise (Of_sexp e)