open StdLabels

let identifier = "tree"
let description = "Build the AST"
let is_global = false
let active = ref true

type context = unit

let initialize = Fun.id
let finalize () = []

module Ast = struct
  type 'a literal = 'a T.literal = Text of string | Expression of 'a
  [@@deriving eq, show]

  type 'a variable = { pos : 'a; name : string; index : 'a expression option }
  [@@deriving eq, show]

  and 'a expression =
    | Integer of 'a * string
    | Literal of 'a * 'a expression literal list
    | Ident of 'a variable
    | BinaryOp of 'a * T.boperator * 'a expression * 'a expression
    | Op of 'a * T.uoperator * 'a expression
    | Function of 'a * T.function_ * 'a expression list
  [@@deriving eq, show]

  and 'a condition = 'a * 'a expression * 'a statement list

  and 'a statement =
    | If of {
        loc : 'a;
        then_ : 'a condition;
        elifs : 'a condition list;
        else_ : 'a statement list;
      }
    | Act of { loc : 'a; label : 'a expression; statements : 'a statement list }
    | Declaration of ('a * 'a variable * T.assignation_operator * 'a expression)
    | Expression of 'a expression
    | Comment of 'a
    | Call of 'a * T.keywords * 'a expression list
    | Location of 'a * string
  [@@deriving eq, show]
end

(** Default implementation for the expression *)
module Expression : sig
  include S.Expression with type t' = S.pos Ast.expression

  val eq : (S.pos -> S.pos -> bool) -> t' -> t' -> bool
  val hash : (S.pos -> int) -> t' -> int
  val exists : f:(t' -> bool) -> t' -> bool
end = struct
  type t = S.pos Ast.expression
  type t' = t

  let eq : (S.pos -> S.pos -> bool) -> t -> t -> bool = Ast.equal_expression

  (* Add a way to filter an expression *)
  let rec exists : f:(t -> bool) -> t -> bool =
   fun ~f -> function
    | BinaryOp (_, _, o1, o2) as op -> f op || exists ~f o1 || exists ~f o2
    | Op (_, _, expr) as op -> f op || exists ~f expr
    | Function (_, _, exprs) as fn -> f fn || List.exists exprs ~f:(exists ~f)
    | Literal (_, litts) as litt ->
        f litt
        || List.exists litts ~f:(function
             | T.Text _ -> false
             | T.Expression ex -> exists ~f ex)
    | Ident { index; _ } as ident -> (
        f ident
        || match index with None -> false | Some expr -> exists ~f expr)
    | Integer _ as int -> f int

  let rec hash : (S.pos -> int) -> t -> int =
   fun f -> function
    | Integer (pos, v) -> Hashtbl.hash (f pos, v)
    | Literal (pos, l) ->
        let litt = List.map ~f:(T.map_litteral ~f:(hash f)) l in
        Hashtbl.hash (f pos, litt)
    | Ident { pos; name; index } ->
        Hashtbl.hash (f pos, name, Option.map (hash f) index)
    | BinaryOp (pos, op, o1, o2) ->
        Hashtbl.hash (f pos, op, hash f o1, hash f o2)
    | Op (pos, op, o1) -> Hashtbl.hash (f pos, op, hash f o1)
    | Function (pos, name, args) ->
        Hashtbl.hash (f pos, name, List.map ~f:(hash f) args)

  let v : t -> t' = fun t -> t
  let integer : S.pos -> string -> t = fun pos i -> Ast.Integer (pos, i)

  let literal : S.pos -> t T.literal list -> t =
   fun pos l -> Ast.Literal (pos, l)

  let function_ : S.pos -> T.function_ -> t list -> t =
   fun pos name args -> Ast.Function (pos, name, args)

  let uoperator : S.pos -> T.uoperator -> t -> t =
   fun pos op expression -> Ast.Op (pos, op, expression)

  let boperator : S.pos -> T.boperator -> t -> t -> t =
   fun pos op op1 op2 ->
    let op1 = op1 and op2 = op2 in
    Ast.BinaryOp (pos, op, op1, op2)

  let ident : (S.pos, t) S.variable -> t =
   fun { pos; name; index } ->
    let index = Option.map (fun i -> i) index in
    Ast.Ident { pos; name; index }
end

module Instruction :
  S.Instruction
    with type t' = S.pos Ast.statement
     and type expression = Expression.t' = struct
  type t = S.pos Ast.statement
  type t' = t
  type expression = Expression.t'

  let v : t -> t' = fun t -> t

  let call : S.pos -> T.keywords -> Expression.t' list -> t =
   fun pos name args -> Ast.Call (pos, name, args)

  let location : S.pos -> string -> t =
   fun loc label -> Ast.Location (loc, label)

  let comment : S.pos -> t = fun pos -> Ast.Comment pos
  let expression : Expression.t' -> t = fun expr -> Ast.Expression expr

  let if_ :
      S.pos ->
      (Expression.t', t) S.clause ->
      elifs:(Expression.t', t) S.clause list ->
      else_:(S.pos * t list) option ->
      t =
   fun pos predicate ~elifs ~else_ ->
    let clause (pos, expr, repr) = (pos, expr, repr) in
    let elifs = List.map ~f:clause elifs
    and else_ =
      match else_ with None -> [] | Some (_, instructions) -> instructions
    in

    Ast.If { loc = pos; then_ = clause predicate; elifs; else_ }

  let act : S.pos -> label:Expression.t' -> t list -> t =
   fun pos ~label statements -> Ast.Act { loc = pos; label; statements }

  let assign :
      S.pos ->
      (S.pos, Expression.t') S.variable ->
      T.assignation_operator ->
      Expression.t' ->
      t =
   fun pos_loc { pos; name; index } op expr ->
    (*let index = Option.map (fun i -> fst @@ Expression.observe (i [])) index*)
    Ast.Declaration (pos_loc, { pos; name; index }, op, expr)
end

module Location = struct
  type t = S.pos * S.pos Ast.statement list

  let v _ = []

  let location : unit -> S.pos -> Instruction.t' list -> t =
   fun () pos block -> (pos, block)
end