open Ast
open Ast.IR
open Baselib

exception Error of string * Lexing.position

(* fonctions d'aide à la gestion des erreurs *)

let expr_pos expr =
  match expr with
  | Syntax.Int n -> n.pos
  | Syntax.Float n -> n.pos
  | Syntax.Char c -> c.pos
  | Syntax.Str s -> s.pos
  | Syntax.Bool b -> b.pos
  | Syntax.Var v -> v.pos
  | Syntax.Call c -> c.pos
;;

let errt expected given pos =
  raise
    (Error
       ( Printf.sprintf
           "expected %s but given %s"
           (string_of_type_t expected)
           (string_of_type_t given)
       , pos ))
;;

(* analyse sémantique *)

let rec analyze_expr expr env =
  match expr with
  | Syntax.Int n -> Int n.value, Int_t
  | Syntax.Float n -> Float n.value, Float_t (* TODO: Should be Float_t *)
  | Syntax.Char c -> Char c.value.[0], Char_t
  | Syntax.Str s -> Str s.value, Str_t
  | Syntax.Bool b -> Bool b.value, Bool_t
  | Syntax.Var va ->
    (match Env.find_opt va.name env with
    | Some v -> Var va.name, v
    | _ -> raise (Error (Printf.sprintf "unbound variable '%s'" va.name, va.pos)))
  | Syntax.Call c ->
    (match Env.find_opt c.func env with
    | Some (Func_t (rt, at)) ->
      if List.length at != List.length c.args
      then
        raise
          (Error
             ( Printf.sprintf
                 "expected %d arguments but given %d"
                 (List.length at)
                 (List.length c.args)
             , c.pos ));
      let args =
        List.map2
          (fun eat a ->
            let aa, at = analyze_expr a env in
            if at = eat then aa else errt eat at (expr_pos a))
          at
          c.args
      in
      Call (c.func, args), rt
    | Some _ -> raise (Error (Printf.sprintf "'%s' is not a function" c.func, c.pos))
    | None -> raise (Error (Printf.sprintf "undefined function '%s'" c.func, c.pos)))
;;

let rec analyze_instr instr env =
  match instr with
  | Syntax.Assign a ->
    let t =
      match Env.find_opt a.var env with
      | Some ty -> ty
      | _ -> raise (Error (Printf.sprintf "undeclared variable '%s'" a.var, a.pos))
    in
    let ae, et = analyze_expr a.expr env in
    if t = et
    then Assign (a.var, ae), Env.add a.var et env
    else
      raise
        (Error
           ( Printf.sprintf
               "changing type of variable '%s' %s -> %s"
               a.var
               (string_of_type_t t)
               (string_of_type_t et)
           , a.pos ))
  | Syntax.Decl d ->
    if Env.mem d.var env
    then raise (Error (Printf.sprintf "redefinition of '%s'" d.var, d.pos))
    else Decl d.var, Env.add d.var d.typ env
  | Syntax.Expr e ->
    let ae, t = analyze_expr e.expr env in
    Expr ae, env
  | Syntax.Return r ->
    let ae, _ = analyze_expr r.expr env in
    Return ae, env
  | Syntax.Cond c ->
    let ae, et = analyze_expr c.expr env in
    if et <> Int_t && et <> Bool_t
    then raise (Error (Printf.sprintf "not a boolean expression", c.pos))
    else (
      let abt, nenv = analyze_block c.blockt env in
      let abf, new_env = analyze_block c.blockf nenv in
      Cond (ae, abt, abf), new_env)
  | Syntax.Loop l ->
    let ae, et = analyze_expr l.expr env in
    if et <> Int_t && et <> Bool_t
    then raise (Error (Printf.sprintf "not a boolean expression", l.pos))
    else (
      let ab, nenv = analyze_block l.block env in
      Loop (ae, ab), nenv)

and analyze_block block env =
  match block with
  | [] -> [], env
  | instr :: rest ->
    let ai, new_env = analyze_instr instr env in
    let ab, nenv = analyze_block rest new_env in
    ai :: ab, nenv
;;

let rec add_arg_env args env =
  match args with
  | [] -> env
  | (ty, arg) :: rest -> add_arg_env rest (Env.add arg ty env)
;;

let rec analyze_def def env =
  match def with
  | Syntax.Func f ->
    if Env.mem f.name env
    then raise (Error (Printf.sprintf "redefinition of function '%s'" f.name, f.pos))
    else (
      let at, args = List.split f.args in
      let nenv = add_arg_env f.args (Env.add f.name (Func_t (f.typ, at)) env) in
      let b, new_env = analyze_block f.block nenv in
      Func (f.name, args, b), Env.add f.name (Func_t (f.typ, at)) env)
;;

let print_env key v = print_endline (key ^ " " ^ string_of_type_t v)

let rec analyze_prog defs env =
  match defs with
  | [] -> [], env
  | def :: rest ->
    let ad, new_env = analyze_def def env in
    let ap, nenv = analyze_prog rest new_env in
    ad :: ap, nenv
;;

let analyze parsed =
  let a, env = analyze_prog parsed Baselib._baselib_ in
  (* Env.iter print_env env; *)
  (* print_endline ("a" ^ " " ^ (string_of_type_t (Env.find "a" env))); *)
  a
;;