open StdLabels

type pos = Lexing.position * Lexing.position

module Ast = struct
  type nonrec pos = pos

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

  and 'a expression =
    | Integer of 'a * string
    | Literal of 'a * string
    | 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 : S.Expression with type t = pos Ast.expression = struct
  type 'a obs
  type t = pos Ast.expression
  type repr = Report.t list -> t * Report.t list
  type variable = { pos : pos; name : string; index : repr option }

  let integer : pos -> string -> repr = fun pos i r -> (Ast.Integer (pos, i), r)
  let literal : pos -> string -> repr = fun pos l r -> (Ast.Literal (pos, l), r)

  let function_ : pos -> T.function_ -> repr list -> repr =
   fun pos name args r ->
    let args = List.map ~f:(fun f -> fst (f r)) args in
    (Ast.Function (pos, name, args), r)

  let uoperator : pos -> T.uoperator -> repr -> repr =
   fun pos op expression r ->
    let expression = fst (expression r) in
    (Ast.Op (pos, op, expression), r)

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

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

module Instruction :
  S.Instruction
    with type expression = Expression.repr
     and type repr = pos Ast.statement
     and type variable = Expression.variable = struct
  type repr = pos Ast.statement
  type expression = Expression.repr
  type variable = Expression.variable

  let call : pos -> T.keywords -> expression list -> repr =
   fun pos name args ->
    let args = List.map ~f:(fun f -> fst (f [])) args in
    Ast.Call (pos, name, args)

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

  let comment : pos -> repr = fun pos -> Ast.Comment pos

  let expression : expression -> repr =
   fun expr -> Ast.Expression (fst (expr []))

  type clause = pos * expression * repr list

  let if_ : pos -> clause -> elifs:clause list -> else_:repr list -> repr =
   fun pos predicate ~elifs ~else_ ->
    let clause (pos, expr, repr) = (pos, fst (expr []), repr) in
    let elifs = List.map ~f:clause elifs in

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

  let act : pos -> label:expression -> repr list -> repr =
   fun pos ~label statements ->
    let label = fst (label []) in
    Ast.Act { loc = pos; label; statements }

  let assign : pos -> variable -> T.assignation_operator -> expression -> repr =
   fun pos_loc { pos; name; index } op expr ->
    let index = Option.map (fun i -> fst (i [])) index
    and expr = fst (expr []) in
    Ast.Declaration (pos_loc, { pos; name; index }, op, expr)
end

module Location = struct
  type instruction = pos Ast.statement
  type repr = pos * instruction list

  let location : pos -> instruction list -> repr = fun pos block -> (pos, block)
end