open Conv
open Core
open Extraction

exception ToSexpError of string

let tag t l = Sexp.List (Sexp.Atom t :: l)

let lit_to_sexp = function
  | LitString s -> tag "LitString" [ Sexp.Atom (str s) ]
  | LitNum (NInt n) ->
      tag "LitNum" [ Sexp.Atom "INT"; Sexp.Atom (str (string_of_Z n)) ]
  | LitNum (NFloat n) ->
      tag "LitNum"
        [
          Sexp.Atom "FLOAT";
          Sexp.Atom (Printf.sprintf "%g" (float_from_flocq n));
        ]
  | LitBool b -> tag "LitBool" [ Sexp.Atom (Bool.to_string b) ]
  | LitNull -> tag "LitNull" []

let option_to_sexp mv ~f =
  match mv with Some v -> tag "Some" [ f v ] | None -> Sexp.Atom "None"

let mode_to_sexp mode =
  Sexp.Atom (match mode with SHALLOW -> "SHALLOW" | DEEP -> "DEEP")

let rec_to_sexp r = Sexp.Atom (match r with REC -> "REC" | NONREC -> "NONREC")

let binop_to_sexp op =
  Sexp.Atom
    (match op with
    | UpdateAttrOp -> "UpdateAttrOp"
    | AddOp -> "AddOp"
    | SubOp -> "SubOp"
    | MulOp -> "MulOp"
    | DivOp -> "DivOp"
    | AndOp -> "AndOp"
    | OrOp -> "OrOp"
    | XOrOp -> "XOrOp"
    | RoundOp Ceil -> "Ceil"
    | RoundOp NearestEven -> "NearestEven"
    | RoundOp Floor -> "Floor"
    | LtOp -> "LtOp"
    | EqOp -> "EqOp"
    | HasAttrOp -> "HasAttrOp"
    | SelectAttrOp -> "SelectAttrOp"
    | DeleteAttrOp -> "DeleteAttrOp"
    | SingletonAttrOp -> "SingletonAttrOp"
    | TypeOfOp -> "TypeOfOp"
    | AppendListOp -> "AppendListOp"
    | MatchAttrOp -> "MatchAttrOp"
    | MatchListOp -> "MatchListOp"
    | MatchStringOp -> "MatchStringOp"
    | FunctionArgsOp -> "FunctionArgsOp")

let kind_to_sexp k = Sexp.Atom (match k with ABS -> "ABS" | WITH -> "WITH")

let rec expr_to_sexp = function
  | ELit l -> tag "ELit" [ lit_to_sexp l ]
  | EId (x, None) -> tag "EId" [ Sexp.Atom (str x) ]
  | EId (x, Some (k, e)) ->
      tag "EId"
        [ Sexp.Atom (str x); tag "alt" [ kind_to_sexp k; expr_to_sexp e ] ]
  | EAbs (x, e) -> tag "EAbs" [ Sexp.Atom (str x); expr_to_sexp e ]
  | EAbsMatch (ms, strict, e) ->
      tag "EAbsMatch"
        [
          Sexp.Atom (if strict then "EXACT" else "LOOSE");
          tag "formals"
            (matcher_fold
               (fun x me se ->
                 Sexp.List
                   [ Sexp.Atom (str x); option_to_sexp me ~f:expr_to_sexp ]
                 :: se)
               [] ms);
          expr_to_sexp e;
        ]
  | EApp (e1, e2) -> tag "EApp" [ expr_to_sexp e1; expr_to_sexp e2 ]
  | ELetAttr (k, e1, e2) ->
      tag "ELetAttr" [ kind_to_sexp k; expr_to_sexp e1; expr_to_sexp e2 ]
  | ESeq (mode, e1, e2) ->
      tag "ESeq" [ mode_to_sexp mode; expr_to_sexp e1; expr_to_sexp e2 ]
  | EAttr bs ->
      tag "EAttr"
        (attr_set_fold
           (fun x (Attr (r, e)) se ->
             Sexp.List [ Sexp.Atom (str x); rec_to_sexp r; expr_to_sexp e ]
             :: se)
           [] bs)
  | EList es ->
      tag "EList"
        (Internal.List.fold_right (fun e se -> expr_to_sexp e :: se) [] es)
  | EBinOp (op, e1, e2) ->
      tag "EBinOp" [ binop_to_sexp op; expr_to_sexp e1; expr_to_sexp e2 ]
  | EIf (e1, e2, e3) ->
      tag "EIf" [ expr_to_sexp e1; expr_to_sexp e2; expr_to_sexp e3 ]

let rec val_to_sexp = function
  | VLit l -> tag "VLit" [ lit_to_sexp l ]
  | VClo _ -> tag "VClo" []
  | VCloMatch _ -> tag "VCloMatch" []
  | VAttr bs ->
      tag "VAttr"
        (Extraction.thunk_map_fold
           (fun x t bs' ->
             Sexp.List [ Sexp.Atom (str x); thunk_to_sexp t ] :: bs')
           [] bs)
  | VList ts ->
      tag "VList"
        (Internal.List.fold_right (fun t st -> thunk_to_sexp t :: st) [] ts)

and env_to_sexp env =
  tag "Env"
    (Extraction.env_fold
       (fun x (k, t) envs ->
         Sexp.List
           [
             Sexp.Atom (str x);
             Sexp.Atom
               (match k with
               | Extraction.ABS -> "ABS"
               | Extraction.WITH -> "WITH");
             thunk_to_sexp t;
           ]
         :: envs)
       [] env)

and thunk_to_sexp = function
  | Thunk _ -> tag "Thunk" [ Sexp.Atom "DELAYED" ]
  | Indirect _ -> tag "Thunk" [ Sexp.Atom "INDIRECT" ]
  | Forced v -> tag "Thunk" [ Sexp.Atom "FORCED"; val_to_sexp v ]

let expr_res_to_sexp = function
  | NoFuel -> Sexp.Atom "NoFuel"
  | Res e -> tag "Res" [ option_to_sexp e ~f:expr_to_sexp ]

let val_res_to_sexp = function
  | NoFuel -> Sexp.Atom "NoFuel"
  | Res e -> tag "Res" [ option_to_sexp e ~f:val_to_sexp ]

let rec (sexp_of_import_tree : Import.tree -> Sexp.t) = function
  | { filename; deps = [] } -> Sexp.Atom filename
  | { filename; deps } ->
      Sexp.List [ Sexp.Atom filename; sexp_of_import_forest deps ]

and sexp_of_import_forest forest =
  Sexp.List (Sexp.Atom "deps" :: List.map forest ~f:sexp_of_import_tree)

exception OfSexpError of string

let rec import_tree_of_sexp : Sexp.t -> Import.tree = function
  | Sexp.Atom filename -> { filename; deps = [] }
  | Sexp.List [ Sexp.Atom filename; deps ] ->
      { filename; deps = import_forest_of_sexp deps }
  | _ -> raise (OfSexpError "Could not parse import tree")

and import_forest_of_sexp = function
  | Sexp.List (Sexp.Atom "deps" :: deps) -> List.map ~f:import_tree_of_sexp deps
  | _ -> raise (OfSexpError "Could not parse import forest")