type expr = {
  r : rdesc;
  rspan : Location.span;
}

and rdesc =
  | RInt of int
  | RBool of bool
  | RString of string
  | RUnit
  | RVar of Ident.t
  | RLet of Ident.t * expr * expr
  | RLambda of Ident.t * expr
  | RApp of expr * expr
  | RIf of expr * expr * expr
  | RBinop of Syntax.binop * expr * expr
  | RNeg of expr
  | RTuple of expr list
  | RRecord of (string * expr) list
  | RField of expr * string
  | RFilter of expr
  | RMap of expr
  | RSum
  | RCount

type helper = {
  h_ident : Ident.t;
  h_params : Ident.t list;
  h_body : expr;
  h_span : Location.span;
}

type program = {
  rp_records : Syntax.record_decl list;
  rp_input : Ident.t;
  rp_input_ty : Syntax.tyexpr;
  rp_helpers : helper list;
  rp_query : Ident.t;
  rp_query_body : expr;
}

type env = {
  globals : (string * Ident.t) list;
  input : (string * Ident.t) option;
  locals : (string * Ident.t) list;
}

let empty_env = { globals = []; input = None; locals = [] }

let make r rspan = { r = r; rspan = rspan }

let lookup env name =
  let rec search = function
    | [] -> None
    | (candidate, ident) :: rest -> if candidate = name then Some ident else search rest
  in
  search env.locals

let find_opt key items =
  let rec search = function
    | [] -> None
    | (candidate, value) :: rest -> if candidate = key then Some value else search rest
  in
  search items

let hashtbl_find_opt table key = try Some (Hashtbl.find table key) with Not_found -> None

let bind env name span =
  let ident = Ident.fresh name span in
  (ident, { env with locals = (name, ident) :: env.locals })

let rec expr env syntax =
  let span = syntax.Syntax.espan in
  match syntax.Syntax.e with
  | Syntax.EInt value -> make (RInt value) span
  | Syntax.EBool value -> make (RBool value) span
  | Syntax.EString value -> make (RString value) span
  | Syntax.EUnit -> make RUnit span
  | Syntax.EVar name -> make (RVar (resolve_name env name span)) span
  | Syntax.ELambda (name, body) ->
      let ident, inner = bind env name span in
      make (RLambda (ident, expr inner body)) span
  | Syntax.ELet (name, bound, body) ->
      let bound = expr env bound in
      let ident, inner = bind env name span in
      make (RLet (ident, bound, expr inner body)) span
  | Syntax.EApp (fn, argument) -> make (RApp (expr env fn, expr env argument)) span
  | Syntax.EIf (condition, then_branch, else_branch) ->
      make (RIf (expr env condition, expr env then_branch, expr env else_branch)) span
  | Syntax.EBinop (operator, left, right) ->
      make (RBinop (operator, expr env left, expr env right)) span
  | Syntax.ENeg operand -> make (RNeg (expr env operand)) span
  | Syntax.ETuple items -> make (RTuple (List.map (expr env) items)) span
  | Syntax.ERecord fields -> make (RRecord (List.map (fun (name, value) -> (name, expr env value)) fields)) span
  | Syntax.EField (record, label) -> make (RField (expr env record, label)) span
  | Syntax.EFilter predicate -> make (RFilter (expr env predicate)) span
  | Syntax.EMap projection -> make (RMap (expr env projection)) span
  | Syntax.ESum -> make RSum span
  | Syntax.ECount -> make RCount span

and resolve_name env name span =
  match lookup env name with
  | Some ident -> ident
  | None -> (
      match env.input with
      | Some (input_name, _) when input_name = name ->
          Diagnostic.error span
            "the input collection `%s` may only be used in the query, not inside a helper function" name
      | _ -> (
          match find_opt name env.globals with
          | Some ident -> ident
          | None ->
              let known =
                List.map fst env.globals @ List.map fst env.locals
                @ (match env.input with Some (name, _) -> [ name ] | None -> [])
              in
              Diagnostic.error span "unbound variable `%s`%s" name (suggestion name known) ))

and suggestion name known =
  if List.mem name known then ""
  else
    match List.filter (fun candidate -> String.length candidate > 0 && String.lowercase_ascii candidate = String.lowercase_ascii name) known with
    | similar :: _ -> Printf.sprintf "; did you mean `%s`?" similar
    | [] -> ""

let duplicate kind span first_span name =
  Diagnostic.error span "duplicate %s `%s`, first declared at %s" kind name
    (Location.to_string first_span)

let check_unique kind items span_of name_of =
  let seen = Hashtbl.create 16 in
  List.iter
    (fun item ->
      let name = name_of item in
      if Hashtbl.mem seen name then duplicate kind (span_of item) (Hashtbl.find seen name) name
      else Hashtbl.add seen name (span_of item))
    items

let label_spans records =
  let seen = Hashtbl.create 32 in
  List.iter
    (fun record ->
      List.iter
        (fun (label, tyexpr) ->
          match hashtbl_find_opt seen label with
          | Some first ->
              Diagnostic.error tyexpr.Syntax.tyspan
                "duplicate field label `%s`, already declared in record `%s`; field labels must be globally unique"
                label first
          | None -> Hashtbl.add seen label record.Syntax.rd_name)
        record.Syntax.rd_fields)
    records

let rec refs_of_expr expr =
  match expr.r with
  | RVar ident -> [ ident ]
  | RLet (_, bound, body) -> refs_of_expr bound @ refs_of_expr body
  | RLambda (_, body) -> refs_of_expr body
  | RApp (fn, argument) -> refs_of_expr fn @ refs_of_expr argument
  | RIf (a, b, c) -> refs_of_expr a @ refs_of_expr b @ refs_of_expr c
  | RBinop (_, a, b) -> refs_of_expr a @ refs_of_expr b
  | RNeg a -> refs_of_expr a
  | RTuple items -> List.concat (List.map refs_of_expr items)
  | RRecord fields -> List.concat (List.map (fun (_, value) -> refs_of_expr value) fields)
  | RField (record, _) -> refs_of_expr record
  | RFilter p -> refs_of_expr p
  | RMap m -> refs_of_expr m
  | RInt _ | RBool _ | RString _ | RUnit | RSum | RCount -> []

let check_recursion helpers =
  let table = Hashtbl.create 16 in
  List.iter (fun helper -> Hashtbl.add table (Ident.stamp helper.h_ident) helper) helpers;
  let state = Hashtbl.create 16 in
  let rec visit helper =
    match hashtbl_find_opt state (Ident.stamp helper.h_ident) with
    | Some `In_progress ->
        Diagnostic.error helper.h_span "helper function `%s` is recursive; recursive definitions are not supported"
          (Ident.name helper.h_ident)
    | Some `Done -> ()
    | None ->
        Hashtbl.add state (Ident.stamp helper.h_ident) `In_progress;
        List.iter
          (fun ident ->
            match hashtbl_find_opt table (Ident.stamp ident) with
            | Some callee -> visit callee
            | None -> ())
          (refs_of_expr helper.h_body);
        Hashtbl.replace state (Ident.stamp helper.h_ident) `Done
  in
  List.iter visit helpers

let program syntax =
  check_unique "record type" syntax.Syntax.prog_records
    (fun record -> record.Syntax.rd_span)
    (fun record -> record.Syntax.rd_name);
  check_unique "helper function" syntax.Syntax.prog_helpers
    (fun helper -> helper.Syntax.h_span)
    (fun helper -> helper.Syntax.h_name);
  label_spans syntax.Syntax.prog_records;
  let input_name = syntax.Syntax.prog_input.Syntax.in_name in
  List.iter
    (fun helper ->
      if helper.Syntax.h_name = input_name then
        Diagnostic.error helper.Syntax.h_name_span
          "the helper function `%s` collides with the input collection" helper.Syntax.h_name;
      let seen = Hashtbl.create 8 in
      List.iter
        (fun parameter ->
          if Hashtbl.mem seen parameter then
            Diagnostic.error helper.Syntax.h_span "duplicate parameter `%s`" parameter
          else Hashtbl.add seen parameter ())
        helper.Syntax.h_params)
    syntax.Syntax.prog_helpers;
  let input_ident = Ident.fresh input_name syntax.Syntax.prog_input.Syntax.in_span in
  let helper_idents =
    List.map
      (fun helper -> (helper.Syntax.h_name, Ident.fresh helper.Syntax.h_name helper.Syntax.h_name_span))
      syntax.Syntax.prog_helpers
  in
  let globals = helper_idents in
  let helpers =
    List.map2
      (fun syntax_helper (_, ident) ->
        let env = { empty_env with globals = globals; input = Some (input_name, input_ident) } in
        let params, body_env =
          List.fold_left
            (fun (params, env) parameter ->
              let name_span = syntax_helper.Syntax.h_span in
              let ident, env = bind env parameter name_span in
              (params @ [ ident ], env))
            ([], env) syntax_helper.Syntax.h_params
        in
        {
          h_ident = ident;
          h_params = params;
          h_body = expr body_env syntax_helper.Syntax.h_body;
          h_span = syntax_helper.Syntax.h_span;
        })
      syntax.Syntax.prog_helpers helper_idents
  in
  check_recursion helpers;
  let query_ident = Ident.fresh syntax.Syntax.prog_query.Syntax.q_name syntax.Syntax.prog_query.Syntax.q_name_span in
  let query_env =
    { empty_env with globals = globals; input = Some (input_name, input_ident); locals = [] }
  in
  let query_env = { query_env with input = None; locals = [ (input_name, input_ident) ] } in
  {
    rp_records = syntax.Syntax.prog_records;
    rp_input = input_ident;
    rp_input_ty = syntax.Syntax.prog_input.Syntax.in_ty;
    rp_helpers = helpers;
    rp_query = query_ident;
    rp_query_body = expr query_env syntax.Syntax.prog_query.Syntax.q_body;
  }